Separate admin interface for LLM and Embedding to allow setting specific fields

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)
 
 

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")
 
 

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"

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",
             ),

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",
+            ),
+        ),
+    ]