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

kiln_ai/adapters/test_ml_embedding_model_list.py

@@ -1,5 +1,8 @@
+from typing import List
+
 import pytest
 
+from kiln_ai.adapters.embedding.embedding_registry import embedding_adapter_from_type
 from kiln_ai.adapters.ml_embedding_model_list import (
     EmbeddingModelName,
     KilnEmbeddingModel,
@@ -10,24 +13,28 @@ from kiln_ai.adapters.ml_embedding_model_list import (
     get_model_by_name,
 )
 from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+from kiln_ai.datamodel.embedding import EmbeddingConfig
 
 
-class TestEmbeddingModelName:
-    """Test cases for EmbeddingModelName enum"""
-
-    def test_enum_values(self):
-        """Test that enum values are correctly defined"""
-        assert (
-            EmbeddingModelName.openai_text_embedding_3_small
-            == "openai_text_embedding_3_small"
-        )
-        assert (
-            EmbeddingModelName.openai_text_embedding_3_large
-            == "openai_text_embedding_3_large"
-        )
-        assert (
-            EmbeddingModelName.gemini_text_embedding_004 == "gemini_text_embedding_004"
+@pytest.fixture
+def litellm_adapter():
+    adapter = embedding_adapter_from_type(
+        EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name=ModelProviderName.openai,
+            model_name=EmbeddingModelName.openai_text_embedding_3_small,
+            properties={},
         )
+    )
+    return adapter
+
+
+def get_all_embedding_models_and_providers() -> List[tuple[str, str]]:
+    return [
+        (model.name, provider.name)
+        for model in built_in_embedding_models
+        for provider in model.providers
+    ]
 
 
 class TestKilnEmbeddingModelProvider:
@@ -120,222 +127,40 @@ class TestKilnEmbeddingModel:
         assert model.providers[1].name == ModelProviderName.anthropic
 
 
-class TestEmbeddingModelsList:
-    """Test cases for the embedding_models list"""
-
-    def test_embedding_models_not_empty(self):
-        """Test that the embedding_models list is not empty"""
-        assert len(built_in_embedding_models) > 0
-
-    def test_all_models_have_required_fields(self):
-        """Test that all models in the list have required fields"""
-        for model in built_in_embedding_models:
-            assert hasattr(model, "family")
-            assert hasattr(model, "name")
-            assert hasattr(model, "friendly_name")
-            assert hasattr(model, "providers")
-            assert isinstance(model.name, str)
-            assert isinstance(model.friendly_name, str)
-            assert isinstance(model.providers, list)
-            assert len(model.providers) > 0
-
-    def test_all_providers_have_required_fields(self):
-        """Test that all providers in all models have required fields"""
-        for model in built_in_embedding_models:
-            for provider in model.providers:
-                assert hasattr(provider, "name")
-                assert isinstance(provider.name, ModelProviderName)
-
-    def test_model_names_are_unique(self):
-        """Test that all model names in the list are unique"""
-        model_names = [model.name for model in built_in_embedding_models]
-        assert len(model_names) == len(set(model_names))
-
-    def test_specific_models_exist(self):
-        """Test that specific expected models exist in the list"""
-        model_names = [model.name for model in built_in_embedding_models]
-
-        assert EmbeddingModelName.openai_text_embedding_3_small in model_names
-        assert EmbeddingModelName.openai_text_embedding_3_large in model_names
-        assert EmbeddingModelName.gemini_text_embedding_004 in model_names
-
-    def test_openai_embedding_models(self):
-        """Test specific OpenAI embedding models"""
-        openai_models = [
-            model
-            for model in built_in_embedding_models
-            if model.family == KilnEmbeddingModelFamily.openai
-        ]
-
-        assert len(openai_models) >= 2  # Should have at least 2 OpenAI models
-
-        # Check for specific OpenAI models
-        openai_model_names = [model.name for model in openai_models]
-        assert EmbeddingModelName.openai_text_embedding_3_small in openai_model_names
-        assert EmbeddingModelName.openai_text_embedding_3_large in openai_model_names
-
-    def test_gemini_embedding_models(self):
-        """Test specific Gemini embedding models"""
-        gemini_models = [
-            model
-            for model in built_in_embedding_models
-            if model.family == KilnEmbeddingModelFamily.gemini
-        ]
-
-        assert len(gemini_models) >= 1  # Should have at least 1 Gemini model
-
-        # Check for specific Gemini model
-        gemini_model_names = [model.name for model in gemini_models]
-        assert EmbeddingModelName.gemini_text_embedding_004 in gemini_model_names
-
-    def test_openai_text_embedding_3_small_details(self):
-        """Test specific details of OpenAI text-embedding-3-small model"""
-        model = get_model_by_name(EmbeddingModelName.openai_text_embedding_3_small)
-
-        assert model.family == KilnEmbeddingModelFamily.openai
-        assert model.friendly_name == "Text Embedding 3 Small"
-        assert len(model.providers) == 1
-
-        provider = model.providers[0]
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-small"
-        assert provider.n_dimensions == 1536
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
-
-    def test_openai_text_embedding_3_large_details(self):
-        """Test specific details of OpenAI text-embedding-3-large model"""
-        model = get_model_by_name(EmbeddingModelName.openai_text_embedding_3_large)
-
-        assert model.family == KilnEmbeddingModelFamily.openai
-        assert model.friendly_name == "Text Embedding 3 Large"
-        assert len(model.providers) == 1
-
-        provider = model.providers[0]
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-large"
-        assert provider.n_dimensions == 3072
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
-
-    def test_gemini_text_embedding_004_details(self):
-        """Test specific details of Gemini text-embedding-004 model"""
-        model = get_model_by_name(EmbeddingModelName.gemini_text_embedding_004)
-
-        assert model.family == KilnEmbeddingModelFamily.gemini
-        assert model.friendly_name == "Text Embedding 004"
-        assert len(model.providers) == 1
-
-        provider = model.providers[0]
-        assert provider.name == ModelProviderName.gemini_api
-        assert provider.model_id == "text-embedding-004"
-        assert provider.n_dimensions == 768
-        assert provider.max_input_tokens == 2048
-        assert provider.supports_custom_dimensions is False
-
-
 class TestGetModelByName:
-    """Test cases for get_model_by_name function"""
-
-    def test_get_existing_model(self):
-        """Test getting an existing model by name"""
-        model = get_model_by_name(EmbeddingModelName.openai_text_embedding_3_small)
-        assert model.name == EmbeddingModelName.openai_text_embedding_3_small
-        assert model.family == KilnEmbeddingModelFamily.openai
-
-    def test_get_all_existing_models(self):
-        """Test getting all existing models by name"""
-        for model_name in EmbeddingModelName:
-            model = get_model_by_name(model_name)
-            assert model.name == model_name
-
     def test_get_nonexistent_model_raises_error(self):
         """Test that getting a nonexistent model raises ValueError"""
         with pytest.raises(
             ValueError, match="Embedding model nonexistent_model not found"
         ):
-            get_model_by_name("nonexistent_model")
-
-    def test_get_model_with_invalid_enum_value(self):
-        """Test that getting a model with invalid enum value raises ValueError"""
-        with pytest.raises(ValueError, match="Embedding model invalid_enum not found"):
-            get_model_by_name("invalid_enum")
+            get_model_by_name("nonexistent_model")  # type: ignore
 
     @pytest.mark.parametrize(
-        "model_name,expected_family,expected_friendly_name",
-        [
-            (
-                EmbeddingModelName.openai_text_embedding_3_small,
-                KilnEmbeddingModelFamily.openai,
-                "Text Embedding 3 Small",
-            ),
-            (
-                EmbeddingModelName.openai_text_embedding_3_large,
-                KilnEmbeddingModelFamily.openai,
-                "Text Embedding 3 Large",
-            ),
-            (
-                EmbeddingModelName.gemini_text_embedding_004,
-                KilnEmbeddingModelFamily.gemini,
-                "Text Embedding 004",
-            ),
-        ],
+        "model_name",
+        [model.name for model in built_in_embedding_models],
     )
-    def test_parametrized_model_retrieval(
-        self, model_name, expected_family, expected_friendly_name
-    ):
+    def test_model_retrieval(self, model_name):
         """Test retrieving models with parametrized test cases"""
         model = get_model_by_name(model_name)
-        assert model.family == expected_family
-        assert model.friendly_name == expected_friendly_name
+        assert model.family == model.family
+        assert model.friendly_name == model.friendly_name
 
 
 class TestBuiltInEmbeddingModelsFromProvider:
-    """Test cases for built_in_embedding_models_from_provider function"""
-
-    def test_get_existing_provider_for_model(self):
-        """Test getting an existing provider for a model"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.openai,
-            model_name=EmbeddingModelName.openai_text_embedding_3_small,
-        )
+    @pytest.mark.parametrize(
+        "model_name,provider_name", get_all_embedding_models_and_providers()
+    )
+    def test_get_all_existing_models_and_providers(self, model_name, provider_name):
+        provider = built_in_embedding_models_from_provider(provider_name, model_name)
 
         assert provider is not None
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-small"
-        assert provider.n_dimensions == 1536
-
-    def test_get_all_existing_provider_model_combinations(self):
-        """Test getting all existing provider-model combinations"""
-        combinations = [
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_small,
-            ),
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_large,
-            ),
-            (
-                ModelProviderName.gemini_api,
-                EmbeddingModelName.gemini_text_embedding_004,
-            ),
-        ]
-
-        for provider_name, model_name in combinations:
-            provider = built_in_embedding_models_from_provider(
-                provider_name, model_name
-            )
-            assert provider is not None
-            assert provider.name == provider_name
-
-    def test_get_nonexistent_provider_returns_none(self):
-        """Test that getting a nonexistent provider returns None"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.anthropic,  # Not used for embeddings
-            model_name=EmbeddingModelName.openai_text_embedding_3_small,
+        assert provider.name == provider_name
+        assert provider.model_id == provider.model_id
+        assert provider.n_dimensions == provider.n_dimensions
+        assert provider.max_input_tokens == provider.max_input_tokens
+        assert (
+            provider.supports_custom_dimensions == provider.supports_custom_dimensions
         )
-        assert provider is None
 
     def test_get_nonexistent_model_returns_none(self):
         """Test that getting a nonexistent model returns None"""
@@ -353,77 +178,62 @@ class TestBuiltInEmbeddingModelsFromProvider:
         )
         assert provider is None
 
-    def test_get_openai_text_embedding_3_small_provider_details(self):
-        """Test specific details of OpenAI text-embedding-3-small provider"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.openai,
-            model_name=EmbeddingModelName.openai_text_embedding_3_small,
-        )
 
-        assert provider is not None
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-small"
-        assert provider.n_dimensions == 1536
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
+class TestGenerateEmbedding:
+    """Test cases for generate_embedding function"""
 
-    def test_get_openai_text_embedding_3_large_provider_details(self):
-        """Test specific details of OpenAI text-embedding-3-large provider"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.openai,
-            model_name=EmbeddingModelName.openai_text_embedding_3_large,
+    @pytest.mark.parametrize(
+        "model_name,provider_name", get_all_embedding_models_and_providers()
+    )
+    @pytest.mark.paid
+    async def test_generate_embedding(self, model_name, provider_name):
+        """Test generating an embedding"""
+        model_provider = built_in_embedding_models_from_provider(
+            provider_name, model_name
         )
-
-        assert provider is not None
-        assert provider.name == ModelProviderName.openai
-        assert provider.model_id == "text-embedding-3-large"
-        assert provider.n_dimensions == 3072
-        assert provider.max_input_tokens == 8192
-        assert provider.supports_custom_dimensions is True
-
-    def test_get_gemini_text_embedding_004_provider_details(self):
-        """Test specific details of Gemini text-embedding-004 provider"""
-        provider = built_in_embedding_models_from_provider(
-            provider_name=ModelProviderName.gemini_api,
-            model_name=EmbeddingModelName.gemini_text_embedding_004,
+        assert model_provider is not None
+
+        embedding = embedding_adapter_from_type(
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name=provider_name,
+                model_name=model_name,
+                properties={},
+            )
         )
-
-        assert provider is not None
-        assert provider.name == ModelProviderName.gemini_api
-        assert provider.model_id == "text-embedding-004"
-        assert provider.n_dimensions == 768
-        assert provider.max_input_tokens == 2048
-        assert provider.supports_custom_dimensions is False
+        embedding = await embedding.generate_embeddings(["Hello, world!"])
+        assert len(embedding.embeddings) == 1
+        assert len(embedding.embeddings[0].vector) == model_provider.n_dimensions
 
     @pytest.mark.parametrize(
-        "provider_name,model_name,expected_model_id,expected_dimensions",
-        [
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_small,
-                "text-embedding-3-small",
-                1536,
-            ),
-            (
-                ModelProviderName.openai,
-                EmbeddingModelName.openai_text_embedding_3_large,
-                "text-embedding-3-large",
-                3072,
-            ),
-            (
-                ModelProviderName.gemini_api,
-                EmbeddingModelName.gemini_text_embedding_004,
-                "text-embedding-004",
-                768,
-            ),
-        ],
+        "model_name,provider_name", get_all_embedding_models_and_providers()
     )
-    def test_parametrized_provider_retrieval(
-        self, provider_name, model_name, expected_model_id, expected_dimensions
+    @pytest.mark.paid
+    async def test_generate_embedding_with_user_supplied_dimensions(
+        self, model_name, provider_name
     ):
-        """Test retrieving providers with parametrized test cases"""
-        provider = built_in_embedding_models_from_provider(provider_name, model_name)
+        """Test generating an embedding with user supplied dimensions"""
+        model_provider = built_in_embedding_models_from_provider(
+            provider_name=provider_name,
+            model_name=model_name,
+        )
+        assert model_provider is not None
 
-        assert provider is not None
-        assert provider.model_id == expected_model_id
-        assert provider.n_dimensions == expected_dimensions
+        if not model_provider.supports_custom_dimensions:
+            pytest.skip("Model does not support custom dimensions")
+
+        # max dim
+        max_dimensions = model_provider.n_dimensions
+        dimensions_target = max_dimensions // 2
+
+        embedding = embedding_adapter_from_type(
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name=provider_name,
+                model_name=model_name,
+                properties={"dimensions": dimensions_target},
+            )
+        )
+        embedding = await embedding.generate_embeddings(["Hello, world!"])
+        assert len(embedding.embeddings) == 1
+        assert len(embedding.embeddings[0].vector) == dimensions_target

kiln_ai/adapters/test_ml_model_list.py

@@ -360,16 +360,6 @@ def test_uncensored():
         assert provider.suggested_for_uncensored_data_gen
 
 
-def test_multimodal_capable():
-    """Test that multimodal_capable is set correctly"""
-    model = get_model_by_name(ModelName.gpt_4_1)
-    for provider in model.providers:
-        assert provider.multimodal_capable
-        assert provider.supports_doc_extraction
-        assert provider.multimodal_mime_types is not None
-        assert len(provider.multimodal_mime_types) > 0
-
-
 def test_no_empty_multimodal_mime_types():
     """Ensure that multimodal fields are self-consistent as they are interdependent"""
     for model in built_in_models:

kiln_ai/datamodel/basemodel.py

@@ -9,7 +9,7 @@ from abc import ABCMeta
 from builtins import classmethod
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Type, TypeVar
+from typing import Any, Callable, Dict, List, Optional, Set, Type, TypeVar
 
 from pydantic import (
     BaseModel,
@@ -78,9 +78,9 @@ def string_to_valid_name(name: str) -> str:
     # https://docs.python.org/3/library/unicodedata.html#unicodedata.normalize
     valid_name = unicodedata.normalize("NFKD", name)
     # Replace any forbidden chars with an underscore
-    valid_name = re.sub(FORBIDDEN_CHARS_REGEX, "_", valid_name)
+    valid_name = re.sub(FORBIDDEN_CHARS_REGEX, " ", valid_name)
     # Replace control characters with an underscore
-    valid_name = re.sub(r"[\x00-\x1F]", "_", valid_name)
+    valid_name = re.sub(r"[\x00-\x1F]", " ", valid_name)
     # Replace consecutive whitespace with a single space
     valid_name = re.sub(r"\s+", " ", valid_name)
     # Replace consecutive underscores with a single underscore
@@ -594,6 +594,34 @@ class KilnParentedModel(KilnBaseModel, metaclass=ABCMeta):
                     return child
         return None
 
+    @classmethod
+    def from_ids_and_parent_path(
+        cls: Type[PT], ids: Set[str], parent_path: Path | None
+    ) -> Dict[str, PT]:
+        """
+        Bulk equivalent of from_id_and_parent_path, much faster for large collections.
+
+        It picks out the matching models from the directory only once. This avoids
+        doing individual costly lookups that scan the whole directory in scenarios
+        where we need to iterate over a large collection of models (e.g. bulk tagging).
+        """
+        if parent_path is None:
+            return {}
+
+        children = {}
+
+        # Note: we're using the in-file ID. We could make this faster using the path-ID if this becomes perf bottleneck, but it's better to have 1 source of truth.
+        for child_path in cls.iterate_children_paths_of_parent_path(parent_path):
+            child_id = ModelCache.shared().get_model_id(child_path, cls)
+            if child_id in ids:
+                children[child_id] = cls.load_from_file(child_path)
+            if child_id is None:
+                child = cls.load_from_file(child_path)
+                if child.id in ids:
+                    children[child.id] = child
+
+        return children
+
 
 # Parent create methods for all child relationships
 # You must pass in parent_of in the subclass definition, defining the child relationships