Merge pull request #1759 from dimagi/sk/custom-models

replace custom models

Author: snopoke

apps/service_providers/llm_service/default_models.py

@@ -137,6 +137,7 @@ def _update_llm_provider_models(LlmProviderModel):
     for m in existing_custom_by_team.values():
         existing_custom_global[(m.type, m.name)].append(m)
 
+    created_models = dict()
     for provider_type, provider_models in DEFAULT_LLM_PROVIDER_MODELS.items():
         for model in provider_models:
             key = (provider_type, model.name)
@@ -147,7 +148,7 @@ def _update_llm_provider_models(LlmProviderModel):
                     existing_global_model.max_token_limit = model.token_limit
                     existing_global_model.save()
             else:
-                LlmProviderModel.objects.create(
+                created_models[(provider_type, model.name)] = LlmProviderModel.objects.create(
                     team=None,
                     type=provider_type,
                     name=model.name,
@@ -170,6 +171,22 @@ def _update_llm_provider_models(LlmProviderModel):
 
         provider_model.delete()
 
+    # replace existing custom models with the new global model and delete the custom models
+    for key, model in created_models.items():
+        if key in existing_custom_global:
+            for custom_model in existing_custom_global[key]:
+                related_objects = get_related_objects(custom_model)
+                for obj in related_objects:
+                    field = [f for f in obj._meta.fields if f.related_model == LlmProviderModel][0]
+                    setattr(obj, field.attname, model.id)
+                    obj.save(update_fields=[field.name])
+
+                related_pipeline_nodes = get_related_pipelines_queryset(custom_model, "llm_provider_model_id")
+                for node in related_pipeline_nodes.select_related("pipeline").all():
+                    _update_pipeline_node_param(node.pipeline, node, "llm_provider_model_id", model.id)
+
+                custom_model.delete()
+
 
 def _get_or_create_custom_model(team_object, key, global_model, existing_custom_by_team):
     """Check the `existing_custom_by_team` mapping for a custom model for the given team and key.

apps/service_providers/tests/test_default_models.py

@@ -1,7 +1,10 @@
+from unittest.mock import patch
+
 import pytest
 
 from apps.service_providers.llm_service.default_models import (
     DEFAULT_LLM_PROVIDER_MODELS,
+    Model,
     get_default_model,
     update_llm_provider_models,
 )
@@ -109,6 +112,52 @@ def test_converts_old_global_models_to_custom_models_pipelines():
     assert node_data[0]["data"]["params"]["llm_provider_model_id"] == custom_model.id
 
 
+@pytest.mark.django_db()
+def test_replaces_custom_models_with_global_models():
+    experiment = ExperimentFactory()
+    old_custom_model = experiment.llm_provider_model
+    assert old_custom_model.team is not None  # Ensure it's a custom model
+
+    # no global model should exist
+    assert not LlmProviderModel.objects.filter(
+        team=None, type=old_custom_model.type, name=old_custom_model.name
+    ).exists()
+
+    defaults = {old_custom_model.type: [Model(old_custom_model.name, old_custom_model.max_token_limit)]}
+    with patch("apps.service_providers.llm_service.default_models.DEFAULT_LLM_PROVIDER_MODELS", defaults):
+        update_llm_provider_models()
+
+    new_global_model = LlmProviderModel.objects.get(team=None, type=old_custom_model.type, name=old_custom_model.name)
+    # custom model should now point to the global model
+    assert not LlmProviderModel.objects.filter(id=old_custom_model.id).exists()
+    experiment.refresh_from_db()
+    assert experiment.llm_provider_model_id == new_global_model.id
+
+
+@pytest.mark.django_db()
+def test_converts_custom_models_to_global_models_pipelines():
+    custom_model = LlmProviderModelFactory()
+    pipeline = get_pipeline(custom_model)
+
+    # no custom model should exist
+    assert not LlmProviderModel.objects.filter(team=None, type=custom_model.type, name=custom_model.name).exists()
+
+    defaults = {custom_model.type: [Model(custom_model.name, custom_model.max_token_limit)]}
+    with patch("apps.service_providers.llm_service.default_models.DEFAULT_LLM_PROVIDER_MODELS", defaults):
+        update_llm_provider_models()
+
+    # custom model is removed
+    assert not LlmProviderModel.objects.filter(id=custom_model.id).exists()
+
+    # global model is created
+    global_model = LlmProviderModel.objects.get(team=None, type=custom_model.type, name=custom_model.name)
+    # pipeline is updated to use the custom model
+    pipeline.refresh_from_db()
+    assert pipeline.node_set.get(type="LLMResponseWithPrompt").params["llm_provider_model_id"] == global_model.id
+    node_data = [node for node in pipeline.data["nodes"] if node["data"]["type"] == "LLMResponseWithPrompt"]
+    assert node_data[0]["data"]["params"]["llm_provider_model_id"] == global_model.id
+
+
 def get_pipeline(llm_provider_model):
     pipeline = PipelineFactory()
     pipeline.data["nodes"].append(