diff --git a/labs/core/admin.py b/labs/core/admin.py index 82169fe..6fbed17 100644 --- a/labs/core/admin.py +++ b/labs/core/admin.py @@ -1,17 +1,39 @@ from django.contrib import admin from django.urls import reverse -from .forms import ProjectForm, ModelFormSet +from .forms import ProjectForm, EmbeddingModelFormSet, LLMModelFormSet from .mixins import JSONFormatterMixin -from .models import Model, Project, Prompt, Variable, VectorizerModel, WorkflowResult +from .models import ( + EmbeddingModel, LLMModel, Project, Prompt, + Variable, VectorizerModel, WorkflowResult +) + +@admin.register(EmbeddingModel) +class EmbeddingModelAdmin(admin.ModelAdmin): + list_display = ("id", "provider", "name", "active") + list_display_links = ("id",) + list_editable = ("name", "active") + list_filter = ("provider", "name") + search_fields = ("provider", "name") + + def has_add_permission(self, request): + return False + + def has_delete_permission(self, request, obj=None): + return False -@admin.register(Model) -class ModelAdmin(admin.ModelAdmin): - list_display = ("id", "model_type", "provider", "model_name", "active") + def get_changelist_formset(self, request, **kwargs): + kwargs["formset"] = EmbeddingModelFormSet + return super().get_changelist_formset(request, **kwargs) + + +@admin.register(LLMModel) +class LLMModelAdmin(admin.ModelAdmin): + list_display = ("id", "provider", "name", "max_output_tokens", "active") list_display_links = ("id",) - list_editable = ("model_type", "provider", "model_name", "active") - list_filter = ("provider", "model_name") - search_fields = ("provider", "model_name") + list_editable = ("name", "max_output_tokens", "active") + list_filter = ("provider", "name") + search_fields = ("provider", "name", "max_output_tokens") def has_add_permission(self, request): return False @@ -20,7 +42,7 @@ def has_delete_permission(self, request, obj=None): return False def get_changelist_formset(self, request, **kwargs): - kwargs["formset"] = ModelFormSet + kwargs["formset"] = LLMModelFormSet return super().get_changelist_formset(request, **kwargs) diff --git a/labs/core/factories.py b/labs/core/factories.py index cea9dd8..ebc0445 100644 --- a/labs/core/factories.py +++ b/labs/core/factories.py @@ -2,8 +2,8 @@ from factory.django import DjangoModelFactory from .models import ( - Model, - ModelTypeEnum, + LLMModel, + EmbeddingModel, Project, Prompt, ProviderEnum, @@ -30,13 +30,22 @@ def default_vectorizer_value_validation(self, create, extracted, **kwargs): raise ValueError("Invalid vectorizer value") -class ModelFactory(DjangoModelFactory): +class LLMModelFactory(DjangoModelFactory): class Meta: - model = Model + model = LLMModel - model_type = factory.Iterator([mt.name for mt in ModelTypeEnum]) provider = factory.Iterator([provider.name for provider in ProviderEnum if provider != ProviderEnum.NO_PROVIDER]) - model_name = factory.Faker("word") + name = factory.Faker("word") + active = factory.Faker("boolean") + max_output_tokens = factory.Faker("integer") + + +class EmbeddingModelFactory(DjangoModelFactory): + class Meta: + model = EmbeddingModel + + provider = factory.Iterator([provider.name for provider in ProviderEnum if provider != ProviderEnum.NO_PROVIDER]) + name = factory.Faker("word") active = factory.Faker("boolean") diff --git a/labs/core/forms.py b/labs/core/forms.py index 82ba5a4..3c61219 100644 --- a/labs/core/forms.py +++ b/labs/core/forms.py @@ -1,6 +1,6 @@ import os -from core.models import Project, ModelTypeEnum +from core.models import Project from django.forms import ( ModelForm, BaseModelFormSet, ValidationError, ChoiceField ) @@ -30,30 +30,22 @@ class Meta: fields = ["name", "description", "path", "url"] -class ModelFormSet(BaseModelFormSet): +class SingleActiveFormSet(BaseModelFormSet): + model_name = "Model" + def clean(self): super().clean() + active_count = sum( + 1 for form in self.forms + if form.cleaned_data and form.cleaned_data.get("active", False) + ) + if active_count != 1: + raise ValidationError( + f"You must have exactly 1 active {self.model_name}, found {active_count}." + ) + +class EmbeddingModelFormSet(SingleActiveFormSet): + model_name = "Embedding" - embedding_name = ModelTypeEnum.EMBEDDING.name - llm_name = ModelTypeEnum.LLM.name - counts = { - embedding_name: 0, - llm_name: 0, - } - errors = [] - - for form in self.forms: - model_type = form.cleaned_data["model_type"] - if model_type in counts and form.cleaned_data["active"]: - counts[model_type] += 1 - - for type_name, count in counts.items(): - if count != 1: - errors.append( - ValidationError( - f"You must have exactly 1 active '{type_name}', found {count}." - ) - ) - - if errors: - raise ValidationError(errors) +class LLMModelFormSet(SingleActiveFormSet): + model_name = "LLM" \ No newline at end of file diff --git a/labs/core/migrations/0006_model_unique_active_embedding_and_more.py b/labs/core/migrations/0006_model_unique_active_embedding_and_more.py index 46b8063..71c9125 100644 --- a/labs/core/migrations/0006_model_unique_active_embedding_and_more.py +++ b/labs/core/migrations/0006_model_unique_active_embedding_and_more.py @@ -1,6 +1,5 @@ # Generated by Django 5.1.5 on 2025-02-21 10:50 -import core.models from django.db import migrations, models @@ -15,7 +14,7 @@ class Migration(migrations.Migration): constraint=models.UniqueConstraint( condition=models.Q( ("active", True), - ("model_type", core.models.ModelTypeEnum["EMBEDDING"]), + ("model_type", "Embedding"), ), fields=("model_type",), name="unique_active_embedding", @@ -24,7 +23,7 @@ class Migration(migrations.Migration): migrations.AddConstraint( model_name="model", constraint=models.UniqueConstraint( - condition=models.Q(("active", True), ("model_type", core.models.ModelTypeEnum["LLM"])), + condition=models.Q(("active", True), ("model_type", "LLM")), fields=("model_type",), name="unique_active_llm", ), diff --git a/labs/core/migrations/0008_alter_variable_value_alter_variable_unique_together.py b/labs/core/migrations/0008_embeddingmodel_llmmodel_delete_model_and_more.py similarity index 100% rename from labs/core/migrations/0008_alter_variable_value_alter_variable_unique_together.py rename to labs/core/migrations/0008_embeddingmodel_llmmodel_delete_model_and_more.py diff --git a/labs/core/migrations/0009_embeddingmodel_llmmodel_delete_model_and_more.py b/labs/core/migrations/0009_embeddingmodel_llmmodel_delete_model_and_more.py new file mode 100644 index 0000000..8d5ce16 --- /dev/null +++ b/labs/core/migrations/0009_embeddingmodel_llmmodel_delete_model_and_more.py @@ -0,0 +1,139 @@ +# Generated by Django 5.1.5 on 2025-02-28 11:11 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("core", "0008_embeddingmodel_llmmodel_delete_model_and_more"), + ] + + operations = [ + migrations.CreateModel( + name="EmbeddingModel", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, + primary_key=True, + serialize=False, + verbose_name="ID", + ), + ), + ( + "provider", + models.CharField( + choices=[ + ("NO_PROVIDER", "No provider"), + ("OPENAI", "OpenAI"), + ("OLLAMA", "Ollama"), + ("GEMINI", "Gemini"), + ("ANTHROPIC", "Anthropic"), + ] + ), + ), + ( + "name", + models.CharField( + help_text="Ensure this Embedding exists and is downloaded.", + max_length=255, + ), + ), + ( + "active", + models.BooleanField( + default=True, help_text="Only one Embedding can be active." + ), + ), + ], + options={ + "verbose_name": "Embedding", + "verbose_name_plural": "Embeddings", + }, + ), + migrations.CreateModel( + name="LLMModel", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, + primary_key=True, + serialize=False, + verbose_name="ID", + ), + ), + ( + "provider", + models.CharField( + choices=[ + ("NO_PROVIDER", "No provider"), + ("OPENAI", "OpenAI"), + ("OLLAMA", "Ollama"), + ("GEMINI", "Gemini"), + ("ANTHROPIC", "Anthropic"), + ] + ), + ), + ( + "name", + models.CharField( + help_text="Ensure this LLM exists and is downloaded.", + max_length=255, + ), + ), + ( + "active", + models.BooleanField( + default=True, help_text="Only one LLM can be active." + ), + ), + ( + "max_output_tokens", + models.IntegerField( + blank=True, + default=None, + help_text="Leave blank for auto-detection, set only if required.", + null=True, + ), + ), + ], + options={ + "verbose_name": "LLM", + "verbose_name_plural": "LLMs", + }, + ), + migrations.DeleteModel( + name="Model", + ), + migrations.AddIndex( + model_name="embeddingmodel", + index=models.Index( + fields=["provider", "name"], name="core_embedd_provide_04143f_idx" + ), + ), + migrations.AddConstraint( + model_name="embeddingmodel", + constraint=models.UniqueConstraint( + condition=models.Q(("active", True)), + fields=("provider",), + name="unique_active_embedding", + ), + ), + migrations.AddIndex( + model_name="llmmodel", + index=models.Index( + fields=["provider", "name"], name="core_llmmod_provide_8330c3_idx" + ), + ), + migrations.AddConstraint( + model_name="llmmodel", + constraint=models.UniqueConstraint( + condition=models.Q(("active", True)), + fields=("provider",), + name="unique_active_llm", + ), + ), + ] diff --git a/labs/core/models.py b/labs/core/models.py index 0887ae0..c77eb71 100644 --- a/labs/core/models.py +++ b/labs/core/models.py @@ -1,6 +1,6 @@ import os from enum import Enum -from typing import Literal, Tuple +from typing import Tuple from django.core.exceptions import ValidationError from django.db import models @@ -25,17 +25,10 @@ "ANTHROPIC": {"llm": AnthropicRequester}, } -vectorizer_model_class = {"CHUNK_VECTORIZER": ChunkVectorizer, "PYTHON_VECTORIZER": PythonVectorizer} - - -class ModelTypeEnum(Enum): - EMBEDDING = "Embedding" - LLM = "LLM" - - @classmethod - def choices(cls): - return [(prop.name, prop.value) for prop in cls] - +vectorizer_model_class = { + "CHUNK_VECTORIZER": ChunkVectorizer, + "PYTHON_VECTORIZER": PythonVectorizer, +} class ProviderEnum(Enum): NO_PROVIDER = "No provider" @@ -91,54 +84,91 @@ def _default_vectorizer_value_validation(self): f"The only possible values for DEFAULT_VECTORIZER are: {', '.join(allowed_vectorizer_values)}" ) - def __str__(self): + def __str__(self) -> str: return self.name class Meta: unique_together = ("provider", "name") -class Model(models.Model): - model_type = models.CharField(choices=ModelTypeEnum.choices()) +class EmbeddingModel(models.Model): provider = models.CharField(choices=ProviderEnum.choices()) - model_name = models.CharField(max_length=255) - active = models.BooleanField(default=True) + name = models.CharField(max_length=255,help_text="Ensure this Embedding exists and is downloaded.") + active = models.BooleanField(default=True, help_text="Only one Embedding can be active.") - @staticmethod - def get_active_embedding_model() -> Tuple[Embedder, str]: - return Model._get_active_provider_model("embedding") - - @staticmethod - def get_active_llm_model() -> Tuple[Requester, str]: - return Model._get_active_provider_model("llm") + @classmethod + def get_active_model(cls) -> Tuple[Embedder, "EmbeddingModel"]: + try: + model = cls.objects.get(active=True) + except cls.DoesNotExist: + raise ValueError("No active embedding model configured") - @staticmethod - def _get_active_provider_model(model_type: Literal["embedding", "llm", "vectorizer"]): - queryset = Model.objects.filter(model_type=model_type.upper(), active=True) - if not queryset.exists(): - raise ValueError(f"No {model_type} model configured") + Variable.load_provider_keys(model.provider) - model = queryset.first() + try: + embedder_class = provider_model_class[model.provider]["embedding"] + except KeyError: + raise ValueError(f"Provider '{model.provider}' is missing or has no embedder class defined.") - # Load associated provider variables - Variable.load_provider_keys(model.provider) - return provider_model_class[model.provider][model_type], model.model_name + return embedder_class, model - def __str__(self): - return f"{self.model_type} {self.provider} {self.model_name}" + def __str__(self) -> str: + return f"EmbeddingModel [{self.provider}] {self.name}" class Meta: + verbose_name = "Embedding" + verbose_name_plural = "Embeddings" constraints = [ models.UniqueConstraint( - fields=["model_type"], - condition=Q(model_type=ModelTypeEnum.EMBEDDING, active=True), + fields=["provider"], + condition=Q(active=True), name="unique_active_embedding", - ), + ) + ] + indexes = [models.Index(fields=["provider", "name"])] + + +class LLMModel(models.Model): + provider = models.CharField(choices=ProviderEnum.choices()) + name = models.CharField(max_length=255, help_text="Ensure this LLM exists and is downloaded.") + active = models.BooleanField(default=True, help_text="Only one LLM can be active.") + max_output_tokens = models.IntegerField( + null=True, + blank=True, + default=None, + help_text="Leave blank for auto-detection, set only if required." + ) + + @classmethod + def get_active_model(cls) -> Tuple[Requester, "LLMModel"]: + try: + model = cls.objects.get(active=True) + except cls.DoesNotExist: + raise ValueError("No active llm model configured") + + Variable.load_provider_keys(model.provider) + + try: + llm_class = provider_model_class[model.provider]["llm"] + except KeyError: + raise ValueError(f"Provider '{model.provider}' is missing or has no LLM class defined.") + + return llm_class, model + + def __str__(self) -> str: + return f"LLMModel [{self.provider}] {self.name}" + + class Meta: + verbose_name = "LLM" + verbose_name_plural = "LLMs" + constraints = [ models.UniqueConstraint( - fields=["model_type"], condition=Q(model_type=ModelTypeEnum.LLM, active=True), name="unique_active_llm" - ), + fields=["provider"], + condition=Q(active=True), + name="unique_active_llm", + ) ] - indexes = [models.Index(fields=["provider", "model_name"])] + indexes = [models.Index(fields=["provider", "name"])] class Project(models.Model): @@ -161,7 +191,7 @@ def save(self, *args, **kwargs): VectorizerModel.objects.create(project=self) Prompt.objects.create(project=self) - def __str__(self): + def __str__(self) -> str: return self.name class Meta: @@ -174,16 +204,20 @@ class VectorizerModel(models.Model): vectorizer_type = models.CharField(choices=VectorizerEnum.choices(), default=Variable.get_default_vectorizer_value) @staticmethod - def get_active_vectorizer(project_id) -> Vectorizer: - queryset = VectorizerModel.objects.filter(project__id=project_id) - if not queryset.exists(): - raise ValueError("No vectorizer configured") + def get_active_vectorizer(project_id: int) -> Vectorizer: + vector_model = VectorizerModel.objects.filter(project__id=project_id).first() + if not vector_model: + raise ValueError("No vectorizer configured for this project.") + + try: + vec_class = vectorizer_model_class[vector_model.vectorizer_type] + except KeyError: + raise ValueError(f"Unrecognized vectorizer type '{vector_model.vectorizer_type}'") - vectorizer_model = queryset.first() - return vectorizer_model_class[vectorizer_model.vectorizer_type] + return vec_class - def __str__(self): - return self.vectorizer_type + def __str__(self) -> str: + return f"{self.vectorizer_type} for {self.project.name}" class Meta: verbose_name = "Vectorizer" @@ -202,7 +236,7 @@ class WorkflowResult(models.Model): pre_commit_error = models.TextField(null=True) created_at = models.DateTimeField(auto_now_add=True) - def __str__(self): + def __str__(self) -> str: return f"{self.task_id}" class Meta: @@ -243,7 +277,7 @@ def get_instruction(project_id: int) -> str: return queryset.first().instruction - def __str__(self): + def __str__(self) -> str: return f"{self.persona[:50]}..., {self.instruction[:50]}..." class Meta: diff --git a/labs/embeddings/gemini.py b/labs/embeddings/gemini.py index cff5364..953e925 100644 --- a/labs/embeddings/gemini.py +++ b/labs/embeddings/gemini.py @@ -3,8 +3,8 @@ from embeddings.embedder import Embeddings class GeminiEmbedder: - def __init__(self, model: str): - self._model_name = model + def __init__(self, model): + self._model_name = model.name api_key = os.environ.get("GEMINI_API_KEY") genai.configure(api_key=api_key) diff --git a/labs/embeddings/ollama.py b/labs/embeddings/ollama.py index e65f2de..1c693f6 100644 --- a/labs/embeddings/ollama.py +++ b/labs/embeddings/ollama.py @@ -5,7 +5,7 @@ class OllamaEmbedder: def __init__(self, model): - self._model_name = model + self._model_name = model.name self._client = Client(settings.LOCAL_LLM_HOST) def embed(self, prompt, *args, **kwargs) -> Embeddings: diff --git a/labs/embeddings/openai.py b/labs/embeddings/openai.py index 28c9e7a..981eac4 100644 --- a/labs/embeddings/openai.py +++ b/labs/embeddings/openai.py @@ -6,7 +6,7 @@ class OpenAIEmbedder: def __init__(self, model): - self._model_name = model + self._model_name = model.name openai.api_key = os.environ.get("OPENAI_API_KEY") def embed(self, prompt, *args, **kwargs) -> Embeddings: diff --git a/labs/fixtures/embedding.json b/labs/fixtures/embedding.json new file mode 100644 index 0000000..413a21f --- /dev/null +++ b/labs/fixtures/embedding.json @@ -0,0 +1,29 @@ +[ + { + "model": "core.embeddingmodel", + "pk": 1, + "fields": { + "provider": "OLLAMA", + "name": "nomic-embed-text:latest", + "active": false + } + }, + { + "model": "core.embeddingmodel", + "pk": 2, + "fields": { + "provider": "OPENAI", + "name": "text-embedding-3-small", + "active": true + } + }, + { + "model": "core.embeddingmodel", + "pk": 3, + "fields": { + "provider": "GEMINI", + "name": "models/text-embedding-004", + "active": false + } + } +] \ No newline at end of file diff --git a/labs/fixtures/llm.json b/labs/fixtures/llm.json new file mode 100644 index 0000000..71e15ec --- /dev/null +++ b/labs/fixtures/llm.json @@ -0,0 +1,40 @@ +[ + { + "model": "core.llmmodel", + "pk": 1, + "fields": { + "provider": "OPENAI", + "name": "gpt-4o", + "active": true + } + }, + { + "model": "core.llmmodel", + "pk": 2, + "fields": { + "provider": "OLLAMA", + "name": "llama3.2:latest", + "max_output_tokens": 2048, + "active": false + } + }, + { + "model": "core.llmmodel", + "pk": 3, + "fields": { + "provider": "GEMINI", + "name": "gemini-2.0-flash", + "active": false + } + }, + { + "model": "core.llmmodel", + "pk": 4, + "fields": { + "provider": "ANTHROPIC", + "name": "claude-3-5-sonnet-20241022", + "max_output_tokens": 8192, + "active": false + } + } +] \ No newline at end of file diff --git a/labs/fixtures/model.json b/labs/fixtures/model.json deleted file mode 100644 index 7edb11e..0000000 --- a/labs/fixtures/model.json +++ /dev/null @@ -1,72 +0,0 @@ -[ - { - "model": "core.model", - "pk": 1, - "fields": { - "model_type": "EMBEDDING", - "provider": "OLLAMA", - "model_name": "nomic-embed-text:latest", - "active": false - } - }, - { - "model": "core.model", - "pk": 2, - "fields": { - "model_type": "LLM", - "provider": "OPENAI", - "model_name": "gpt-4o", - "active": true - } - }, - { - "model": "core.model", - "pk": 3, - "fields": { - "model_type": "EMBEDDING", - "provider": "OPENAI", - "model_name": "text-embedding-3-small", - "active": true - } - }, - { - "model": "core.model", - "pk": 4, - "fields": { - "model_type": "LLM", - "provider": "OLLAMA", - "model_name": "llama3.2:latest", - "active": false - } - }, - { - "model": "core.model", - "pk": 5, - "fields": { - "model_type": "LLM", - "provider": "GEMINI", - "model_name": "gemini-2.0-flash", - "active": false - } - }, - { - "model": "core.model", - "pk": 6, - "fields": { - "model_type": "EMBEDDING", - "provider": "GEMINI", - "model_name": "models/text-embedding-004", - "active": false - } - }, - { - "model": "core.model", - "pk": 7, - "fields": { - "model_type": "LLM", - "provider": "ANTHROPIC", - "model_name": "claude-3-5-sonnet-20241022", - "active": false - } - } -] diff --git a/labs/llm/anthropic.py b/labs/llm/anthropic.py index 2741690..57e2821 100644 --- a/labs/llm/anthropic.py +++ b/labs/llm/anthropic.py @@ -9,8 +9,9 @@ class AnthropicRequester: - def __init__(self, model: str): - self._model_name = model + def __init__(self, model): + self._model_name = model.name + self._model_max_output_tokens = model.max_output_tokens api_key = os.environ.get("ANTHROPIC_API_KEY") self.client = Anthropic(api_key=api_key) @@ -25,7 +26,10 @@ def completion_without_proxy( try: response = self.client.messages.create( - model=self._model_name, system=system_prompt, messages=user_messages, max_tokens=8192 + model=self._model_name, + system=system_prompt, + messages=user_messages, + max_tokens=self._model_max_output_tokens, ) response_steps = self.response_to_steps(response) diff --git a/labs/llm/gemini.py b/labs/llm/gemini.py index c00fc34..2bc1899 100644 --- a/labs/llm/gemini.py +++ b/labs/llm/gemini.py @@ -6,8 +6,8 @@ class GeminiRequester: - def __init__(self, model: str): - self._model_name = model + def __init__(self, model): + self._model_name = model.name api_key = os.environ.get("GEMINI_API_KEY") genai.configure(api_key=api_key) self.generative_model = genai.GenerativeModel(self._model_name) diff --git a/labs/llm/ollama.py b/labs/llm/ollama.py index 4a2fb57..5781afa 100644 --- a/labs/llm/ollama.py +++ b/labs/llm/ollama.py @@ -1,18 +1,20 @@ from django.conf import settings +from typing import Tuple, Dict, Any from ollama import Client class OllamaRequester: def __init__(self, model): - self._model_name = model + self._model_name = model.name + self._model_max_output_tokens = model.max_output_tokens self._client = Client(settings.LOCAL_LLM_HOST) - def completion_without_proxy(self, messages, *args, **kwargs): + def completion_without_proxy(self, messages, *args, **kwargs) -> Tuple[str, Dict[str, Any]]: response = self._client.chat( model=self._model_name, messages=messages, format="json", - options={"num_ctx": 8192}, + options={"num_ctx": self._model_max_output_tokens}, *args, **kwargs, ) diff --git a/labs/llm/openai.py b/labs/llm/openai.py index 6543110..420f101 100644 --- a/labs/llm/openai.py +++ b/labs/llm/openai.py @@ -1,4 +1,5 @@ import os +from typing import Tuple, Dict, Any import openai from litellm import completion @@ -6,10 +7,10 @@ class OpenAIRequester: def __init__(self, model): - self._model_name = model + self._model_name = model.name openai.api_key = os.environ.get("OPENAI_API_KEY") - def completion_without_proxy(self, messages, *args, **kwargs): + def completion_without_proxy(self, messages, *args, **kwargs) -> Tuple[str, Dict[str, Any]]: return self._model_name, completion( model=self._model_name, messages=messages, diff --git a/labs/tasks/llm.py b/labs/tasks/llm.py index 71005f8..ff1efab 100644 --- a/labs/tasks/llm.py +++ b/labs/tasks/llm.py @@ -4,7 +4,7 @@ from config.celery import app from config.redis_client import RedisVariable, redis_client -from core.models import Model, Project, VectorizerModel +from core.models import LLMModel, EmbeddingModel, Project, VectorizerModel from django.conf import settings from embeddings.embedder import Embedder from embeddings.vectorizers.vectorizer import Vectorizer @@ -17,7 +17,7 @@ def get_llm_response(prompt): - llm_requester, *llm_requester_args = Model.get_active_llm_model() + llm_requester, *llm_requester_args = LLMModel.get_active_model() requester = Requester(llm_requester, *llm_requester_args) retries, max_retries = 0, 5 @@ -49,7 +49,7 @@ def vectorize_repository_task(prefix="", project_id=0): if not (project_path := redis_client.get(RedisVariable.PROJECT_PATH, prefix=prefix)): project_path = Project.objects.get(id=project_id).path - embedder_class, *embeder_args = Model.get_active_embedding_model() + embedder_class, *embeder_args = EmbeddingModel.get_active_model() embedder = Embedder(embedder_class, *embeder_args) vectorizer_class = VectorizerModel.get_active_vectorizer(project_id) @@ -70,7 +70,7 @@ def find_embeddings_task( ): project_id = redis_client.get(RedisVariable.PROJECT, prefix=prefix, default=project_id) - embedder_class, *embeder_args = Model.get_active_embedding_model() + embedder_class, *embeder_args = EmbeddingModel.get_active_model() files_path = Embedder(embedder_class, *embeder_args).retrieve_files_path( redis_client.get(RedisVariable.ISSUE_BODY, prefix=prefix, default=issue_body), project_id, diff --git a/labs/tasks/logging.py b/labs/tasks/logging.py index 8548356..ec0f5c7 100644 --- a/labs/tasks/logging.py +++ b/labs/tasks/logging.py @@ -1,12 +1,12 @@ from config.celery import app from config.redis_client import RedisVariable, redis_client -from core.models import Model, WorkflowResult +from core.models import EmbeddingModel, LLMModel, WorkflowResult @app.task def save_workflow_result_task(prefix): - _, embedding_model_name = Model.get_active_embedding_model() - _, llm_model_name = Model.get_active_llm_model() + _, embedding_model_name = EmbeddingModel.get_active_model() + _, llm_model_name = LLMModel.get_active_model() project_id = redis_client.get(RedisVariable.PROJECT, prefix) embeddings = redis_client.get(RedisVariable.EMBEDDINGS, prefix) context = redis_client.get(RedisVariable.CONTEXT, prefix) diff --git a/labs/tests/conftest.py b/labs/tests/conftest.py index 1ea4ebe..8709707 100644 --- a/labs/tests/conftest.py +++ b/labs/tests/conftest.py @@ -1,7 +1,7 @@ from typing import List import pytest -from core.models import Model, ModelTypeEnum, Project, ProviderEnum, Variable +from core.models import LLMModel, EmbeddingModel, Project, ProviderEnum, Variable from embeddings.models import Embedding from tests.constants import ( ANTHROPIC_LLM_MODEL_NAME, @@ -50,10 +50,9 @@ def create_multiple_embeddings(create_test_project): @pytest.fixture @pytest.mark.django_db def create_test_ollama_embedding_config(): - return Model.objects.create( - model_type=ModelTypeEnum.EMBEDDING.name, + return EmbeddingModel.objects.create( provider=ProviderEnum.OLLAMA.name, - model_name=OLLAMA_EMBEDDING_MODEL_NAME, + name=OLLAMA_EMBEDDING_MODEL_NAME, active=True, ) @@ -61,10 +60,9 @@ def create_test_ollama_embedding_config(): @pytest.fixture @pytest.mark.django_db def create_test_ollama_llm_config(): - return Model.objects.create( - model_type=ModelTypeEnum.LLM.name, + return LLMModel.objects.create( provider=ProviderEnum.OLLAMA.name, - model_name=OLLAMA_LLM_MODEL_NAME, + name=OLLAMA_LLM_MODEL_NAME, active=True, ) @@ -72,10 +70,9 @@ def create_test_ollama_llm_config(): @pytest.fixture @pytest.mark.django_db def create_test_openai_embedding_config(): - return Model.objects.create( - model_type=ModelTypeEnum.EMBEDDING.name, + return EmbeddingModel.objects.create( provider=ProviderEnum.OPENAI.name, - model_name=OPENAI_EMBEDDING_MODEL_NAME, + name=OPENAI_EMBEDDING_MODEL_NAME, active=True, ) @@ -83,10 +80,9 @@ def create_test_openai_embedding_config(): @pytest.fixture @pytest.mark.django_db def create_test_openai_llm_config(): - return Model.objects.create( - model_type=ModelTypeEnum.LLM.name, + return LLMModel.objects.create( provider=ProviderEnum.OPENAI.name, - model_name=OPENAI_LLM_MODEL_NAME, + name=OPENAI_LLM_MODEL_NAME, active=True, ) @@ -94,10 +90,9 @@ def create_test_openai_llm_config(): @pytest.fixture @pytest.mark.django_db def create_test_gemini_embedding_config(): - return Model.objects.create( - model_type=ModelTypeEnum.EMBEDDING.name, + return EmbeddingModel.objects.create( provider=ProviderEnum.GEMINI.name, - model_name=GEMINI_EMBEDDING_MODEL_NAME, + name=GEMINI_EMBEDDING_MODEL_NAME, active=True, ) @@ -105,10 +100,9 @@ def create_test_gemini_embedding_config(): @pytest.fixture @pytest.mark.django_db def create_test_gemini_llm_config(): - return Model.objects.create( - model_type=ModelTypeEnum.LLM.name, + return LLMModel.objects.create( provider=ProviderEnum.GEMINI.name, - model_name=GEMINI_LLM_MODEL_NAME, + name=GEMINI_LLM_MODEL_NAME, active=True, ) @@ -116,9 +110,8 @@ def create_test_gemini_llm_config(): @pytest.fixture @pytest.mark.django_db def create_test_anthropic_llm_config(): - return Model.objects.create( - model_type=ModelTypeEnum.LLM.name, + return LLMModel.objects.create( provider=ProviderEnum.ANTHROPIC.name, - model_name=ANTHROPIC_LLM_MODEL_NAME, + name=ANTHROPIC_LLM_MODEL_NAME, active=True, ) diff --git a/labs/tests/database/test_embeddings.py b/labs/tests/database/test_embeddings.py index 8c6ebea..cc71eb9 100644 --- a/labs/tests/database/test_embeddings.py +++ b/labs/tests/database/test_embeddings.py @@ -6,6 +6,7 @@ from embeddings.embedder import Embedder, Embeddings from embeddings.models import Embedding from embeddings.openai import OpenAIEmbedder +from core.factories import EmbeddingModelFactory from tests.constants import MULTIPLE_EMBEDDINGS, OPENAI_EMBEDDING_MODEL_NAME, SINGLE_EMBEDDING @@ -56,7 +57,13 @@ def test_reembed_code(create_test_project): ], ) - Embedder(OpenAIEmbedder, model=OPENAI_EMBEDDING_MODEL_NAME).reembed_code( + embedding_model = EmbeddingModelFactory( + provider="OPENAI", + name=OPENAI_EMBEDDING_MODEL_NAME, + active=True, + ) + + Embedder(OpenAIEmbedder, model=embedding_model).reembed_code( project_id=project.id, files_texts=files_texts, embeddings=embeddings, diff --git a/labs/tests/test_llm.py b/labs/tests/test_llm.py index e18e1ba..c351b05 100644 --- a/labs/tests/test_llm.py +++ b/labs/tests/test_llm.py @@ -2,7 +2,7 @@ from unittest.mock import patch import pytest -from core.models import Model, VectorizerModel +from core.models import LLMModel, EmbeddingModel, VectorizerModel from embeddings.embedder import Embedder from embeddings.gemini import GeminiEmbedder from embeddings.ollama import OllamaEmbedder @@ -31,7 +31,7 @@ def call_llm_with_context(issue_summary, project): if not issue_summary: raise ValueError("issue_summary cannot be empty.") - embedder_class, *embeder_args = Model.get_active_embedding_model() + embedder_class, *embeder_args = EmbeddingModel.get_active_model() embedder = Embedder(embedder_class, *embeder_args) vectorizer_class = VectorizerModel.get_active_vectorizer(project.id) @@ -146,49 +146,49 @@ def test_local_llm_redirect( class TestLLMRequester: @pytest.mark.django_db def test_openai_llm_requester(self, create_test_openai_llm_config): - requester, model_name = Model.get_active_llm_model() + requester, model = LLMModel.get_active_model() assert issubclass(requester, OpenAIRequester) - assert model_name == OPENAI_LLM_MODEL_NAME + assert model.name == OPENAI_LLM_MODEL_NAME @pytest.mark.django_db def test_openai_embedder(self, create_test_openai_embedding_config): - embedder, model_name = Model.get_active_embedding_model() + embedder, model = EmbeddingModel.get_active_model() assert issubclass(embedder, OpenAIEmbedder) - assert model_name == OPENAI_EMBEDDING_MODEL_NAME + assert model.name == OPENAI_EMBEDDING_MODEL_NAME @pytest.mark.django_db def test_ollama_llm_requester(self, create_test_ollama_llm_config): - requester, model_name = Model.get_active_llm_model() + requester, model = LLMModel.get_active_model() assert issubclass(requester, OllamaRequester) - assert model_name == OLLAMA_LLM_MODEL_NAME + assert model.name == OLLAMA_LLM_MODEL_NAME @pytest.mark.django_db def test_ollama_embedder(self, create_test_ollama_embedding_config): - embedder, model_name = Model.get_active_embedding_model() + embedder, model = EmbeddingModel.get_active_model() assert issubclass(embedder, OllamaEmbedder) - assert model_name == OLLAMA_EMBEDDING_MODEL_NAME + assert model.name == OLLAMA_EMBEDDING_MODEL_NAME @pytest.mark.django_db def test_gemini_llm_requester(self, create_test_gemini_llm_config): - requester, model_name = Model.get_active_llm_model() + requester, model = LLMModel.get_active_model() assert issubclass(requester, GeminiRequester) - assert model_name == GEMINI_LLM_MODEL_NAME + assert model.name == GEMINI_LLM_MODEL_NAME @pytest.mark.django_db def test_gemini_embedder(self, create_test_gemini_embedding_config): - embedder, model_name = Model.get_active_embedding_model() + embedder, model = EmbeddingModel.get_active_model() assert issubclass(embedder, GeminiEmbedder) - assert model_name == GEMINI_EMBEDDING_MODEL_NAME + assert model.name == GEMINI_EMBEDDING_MODEL_NAME @pytest.mark.django_db def test_anthropic_llm_requester(self, create_test_anthropic_llm_config): - requester, model_name = Model.get_active_llm_model() + requester, model = LLMModel.get_active_model() assert issubclass(requester, AnthropicRequester) - assert model_name == ANTHROPIC_LLM_MODEL_NAME + assert model.name == ANTHROPIC_LLM_MODEL_NAME