kiln-ai 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/adapters/remote_config.py

@@ -3,28 +3,53 @@ import json
 import logging
 import os
 import threading
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, List
 
 import requests
 from pydantic import ValidationError
 
+from kiln_ai.adapters.ml_embedding_model_list import (
+    KilnEmbeddingModel,
+    KilnEmbeddingModelProvider,
+    built_in_embedding_models,
+)
+from kiln_ai.datamodel.datamodel_enums import KilnMimeType
+
 from .ml_model_list import KilnModel, KilnModelProvider, built_in_models
 
 logger = logging.getLogger(__name__)
 
 
-def serialize_config(models: List[KilnModel], path: str | Path) -> None:
-    data = {"model_list": [m.model_dump(mode="json") for m in models]}
+@dataclass
+class KilnRemoteConfig:
+    model_list: List[KilnModel]
+    embedding_model_list: List[KilnEmbeddingModel]
+
+
+def serialize_config(
+    models: List[KilnModel],
+    embedding_models: List[KilnEmbeddingModel],
+    path: str | Path,
+) -> None:
+    data = {
+        "model_list": [m.model_dump(mode="json") for m in models],
+        "embedding_model_list": [m.model_dump(mode="json") for m in embedding_models],
+    }
     Path(path).write_text(json.dumps(data, indent=2, sort_keys=True))
 
 
-def deserialize_config_at_path(path: str | Path) -> List[KilnModel]:
+def deserialize_config_at_path(
+    path: str | Path,
+) -> KilnRemoteConfig:
     raw = json.loads(Path(path).read_text())
     return deserialize_config_data(raw)
 
 
-def deserialize_config_data(config_data: Any) -> List[KilnModel]:
+def deserialize_config_data(
+    config_data: Any,
+) -> KilnRemoteConfig:
     if not isinstance(config_data, dict):
         raise ValueError(f"Remote config expected dict, got {type(config_data)}")
 
@@ -34,6 +59,12 @@ def deserialize_config_data(config_data: Any) -> List[KilnModel]:
             f"Remote config expected list of models, got {type(model_list)}"
         )
 
+    embedding_model_data = config_data.get("embedding_model_list", [])
+    if not isinstance(embedding_model_data, list):
+        raise ValueError(
+            f"Remote config expected list of embedding models, got {type(embedding_model_data)}"
+        )
+
     # We must be careful here, because some of the JSON data may be generated from a forward
     # version of the code that has newer fields / versions of the fields, that may cause
     # the current client this code is running on to fail to validate the item into a KilnModel.
@@ -48,6 +79,14 @@ def deserialize_config_data(config_data: Any) -> List[KilnModel]:
             providers = []
             for provider_data in providers_list:
                 try:
+                    # we filter out the mime types that we don't support
+                    mime_types = provider_data.get("multimodal_mime_types")
+                    if mime_types is not None:
+                        provider_data["multimodal_mime_types"] = [
+                            mime_type
+                            for mime_type in mime_types
+                            if mime_type in list(KilnMimeType)
+                        ]
                     provider = KilnModelProvider.model_validate(provider_data)
                     providers.append(provider)
                 except ValidationError as e:
@@ -72,10 +111,38 @@ def deserialize_config_data(config_data: Any) -> List[KilnModel]:
                 model_data,
                 e,
             )
-    return models
+
+    embedding_models = []
+    for embedding_model_data in embedding_model_data:
+        try:
+            provider_list = embedding_model_data.get("providers", [])
+            providers = []
+            for provider_data in provider_list:
+                try:
+                    provider = KilnEmbeddingModelProvider.model_validate(provider_data)
+                    providers.append(provider)
+                except ValidationError as e:
+                    logger.warning(
+                        "Failed to validate an embedding model provider from remote config. Upgrade Kiln to use this model. Details %s: %s",
+                        provider_data,
+                        e,
+                    )
+
+            embedding_model_data["providers"] = []
+            embedding_model = KilnEmbeddingModel.model_validate(embedding_model_data)
+            embedding_model.providers = providers
+            embedding_models.append(embedding_model)
+        except ValidationError as e:
+            logger.warning(
+                "Failed to validate an embedding model from remote config. Upgrade Kiln to use this model. Details %s: %s",
+                embedding_model_data,
+                e,
+            )
+
+    return KilnRemoteConfig(model_list=models, embedding_model_list=embedding_models)
 
 
-def load_from_url(url: str) -> List[KilnModel]:
+def load_from_url(url: str) -> KilnRemoteConfig:
     response = requests.get(url, timeout=10)
     response.raise_for_status()
     data = response.json()
@@ -83,7 +150,11 @@ def load_from_url(url: str) -> List[KilnModel]:
 
 
 def dump_builtin_config(path: str | Path) -> None:
-    serialize_config(built_in_models, path)
+    serialize_config(
+        models=built_in_models,
+        embedding_models=built_in_embedding_models,
+        path=path,
+    )
 
 
 def load_remote_models(url: str) -> None:
@@ -93,7 +164,8 @@ def load_remote_models(url: str) -> None:
     def fetch_and_replace() -> None:
         try:
             models = load_from_url(url)
-            built_in_models[:] = models
+            built_in_models[:] = models.model_list
+            built_in_embedding_models[:] = models.embedding_model_list
         except Exception as exc:
             # Do not crash startup, but surface the issue
             logger.warning("Failed to fetch remote model list from %s: %s", url, exc)

kiln_ai/adapters/repair/test_repair_task.py

@@ -229,21 +229,20 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         "rating": 8,
     }
 
+    run_config = RunConfigProperties(
+        model_name="llama_3_1_8b",
+        model_provider_name="ollama",
+        prompt_id="simple_prompt_builder",
+        structured_output_mode="json_schema",
+    )
+
     with patch.object(LiteLlmAdapter, "_run", new_callable=AsyncMock) as mock_run:
         mock_run.return_value = (
             RunOutput(output=mocked_output, intermediate_outputs=None),
             None,
         )
 
-        adapter = adapter_for_task(
-            repair_task,
-            RunConfigProperties(
-                model_name="llama_3_1_8b",
-                model_provider_name="ollama",
-                prompt_id="simple_prompt_builder",
-                structured_output_mode="json_schema",
-            ),
-        )
+        adapter = adapter_for_task(repair_task, run_config)
 
         run = await adapter.invoke(repair_task_input.model_dump())
 
@@ -264,6 +263,10 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     }
     assert run.input_source.type == DataSourceType.human
     assert "created_by" in run.input_source.properties
+    assert run.output.source is not None
+    assert run.output.source.run_config is not None
+    saved_run_config = run.output.source.run_config.model_dump()
+    assert saved_run_config == run_config.model_dump()
 
     # Verify that the mock was called
     mock_run.assert_called_once()

kiln_ai/adapters/run_output.py

@@ -3,9 +3,12 @@ from typing import Dict
 
 from litellm.types.utils import ChoiceLogprobs
 
+from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam
+
 
 @dataclass
 class RunOutput:
     output: Dict | str
     intermediate_outputs: Dict[str, str] | None
     output_logprobs: ChoiceLogprobs | None = None
+    trace: list[ChatCompletionMessageParam] | None = None