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

kiln_ai/adapters/test_ml_embedding_model_list.py

@@ -0,0 +1,239 @@
+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,
+    KilnEmbeddingModelFamily,
+    KilnEmbeddingModelProvider,
+    built_in_embedding_models,
+    built_in_embedding_models_from_provider,
+    get_model_by_name,
+)
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+from kiln_ai.datamodel.embedding import EmbeddingConfig
+
+
+@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:
+    """Test cases for KilnEmbeddingModelProvider model"""
+
+    def test_basic_provider_creation(self):
+        """Test creating a basic provider with required fields"""
+        provider = KilnEmbeddingModelProvider(
+            name=ModelProviderName.openai,
+            model_id="text-embedding-3-small",
+            max_input_tokens=8192,
+            n_dimensions=1536,
+            supports_custom_dimensions=True,
+        )
+
+        assert provider.name == ModelProviderName.openai
+        assert provider.model_id == "text-embedding-3-small"
+        assert provider.max_input_tokens == 8192
+        assert provider.n_dimensions == 1536
+        assert provider.supports_custom_dimensions is True
+
+    def test_provider_with_optional_fields_unspecified(self):
+        """Test creating a provider with optional fields not specified"""
+        provider = KilnEmbeddingModelProvider(
+            name=ModelProviderName.gemini_api,
+            model_id="text-embedding-004",
+            n_dimensions=768,
+        )
+
+        assert provider.name == ModelProviderName.gemini_api
+        assert provider.model_id == "text-embedding-004"
+        assert provider.max_input_tokens is None
+        assert provider.n_dimensions == 768
+        assert provider.supports_custom_dimensions is False
+
+
+class TestKilnEmbeddingModel:
+    """Test cases for KilnEmbeddingModel model"""
+
+    def test_basic_model_creation(self):
+        """Test creating a basic model with required fields"""
+        providers = [
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.openai,
+                model_id="text-embedding-3-small",
+                n_dimensions=1536,
+                max_input_tokens=8192,
+            )
+        ]
+
+        model = KilnEmbeddingModel(
+            family=KilnEmbeddingModelFamily.openai,
+            name=EmbeddingModelName.openai_text_embedding_3_small,
+            friendly_name="Text Embedding 3 Small",
+            providers=providers,
+        )
+
+        assert model.family == KilnEmbeddingModelFamily.openai
+        assert model.name == EmbeddingModelName.openai_text_embedding_3_small
+        assert model.friendly_name == "Text Embedding 3 Small"
+        assert len(model.providers) == 1
+        assert model.providers[0].name == ModelProviderName.openai
+
+    def test_model_with_multiple_providers(self):
+        """Test creating a model with multiple providers"""
+        providers = [
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.openai,
+                model_id="model-1",
+                n_dimensions=1536,
+                max_input_tokens=8192,
+            ),
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.anthropic,
+                model_id="model-1",
+                n_dimensions=1536,
+                max_input_tokens=8192,
+            ),
+        ]
+
+        model = KilnEmbeddingModel(
+            family=KilnEmbeddingModelFamily.openai,
+            name=EmbeddingModelName.openai_text_embedding_3_small,
+            friendly_name="text-embedding-3-small",
+            providers=providers,
+        )
+
+        assert len(model.providers) == 2
+        assert model.providers[0].name == ModelProviderName.openai
+        assert model.providers[1].name == ModelProviderName.anthropic
+
+
+class TestGetModelByName:
+    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")  # type: ignore
+
+    @pytest.mark.parametrize(
+        "model_name",
+        [model.name for model in built_in_embedding_models],
+    )
+    def test_model_retrieval(self, model_name):
+        """Test retrieving models with parametrized test cases"""
+        model = get_model_by_name(model_name)
+        assert model.family == model.family
+        assert model.friendly_name == model.friendly_name
+
+
+class TestBuiltInEmbeddingModelsFromProvider:
+    @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 == 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
+        )
+
+    def test_get_nonexistent_model_returns_none(self):
+        """Test that getting a nonexistent model returns None"""
+        provider = built_in_embedding_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="nonexistent_model",
+        )
+        assert provider is None
+
+    def test_get_wrong_provider_for_model_returns_none(self):
+        """Test that getting wrong provider for a model returns None"""
+        provider = built_in_embedding_models_from_provider(
+            provider_name=ModelProviderName.gemini_api,
+            model_name=EmbeddingModelName.openai_text_embedding_3_small,
+        )
+        assert provider is None
+
+
+class TestGenerateEmbedding:
+    """Test cases for generate_embedding function"""
+
+    @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 model_provider is not None
+
+        embedding = embedding_adapter_from_type(
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name=provider_name,
+                model_name=model_name,
+                properties={},
+            )
+        )
+        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(
+        "model_name,provider_name", get_all_embedding_models_and_providers()
+    )
+    @pytest.mark.paid
+    async def test_generate_embedding_with_user_supplied_dimensions(
+        self, model_name, provider_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
+
+        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

@@ -5,6 +5,7 @@ import pytest
 from kiln_ai.adapters.ml_model_list import (
     ModelName,
     built_in_models,
+    built_in_models_from_provider,
     default_structured_output_mode_for_model_provider,
     get_model_by_name,
 )
@@ -161,6 +162,186 @@ class TestDefaultStructuredOutputModeForModelProvider:
         assert result == first_provider.structured_output_mode
 
 
+class TestBuiltInModelsFromProvider:
+    """Test cases for built_in_models_from_provider function"""
+
+    def test_valid_model_and_provider_returns_provider(self):
+        """Test that valid model and provider returns the correct provider configuration"""
+        # GPT 4.1 has OpenAI provider
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="gpt_4_1",
+        )
+        assert result is not None
+        assert result.name == ModelProviderName.openai
+        assert result.model_id == "gpt-4.1"
+        assert result.supports_logprobs is True
+        assert result.suggested_for_data_gen is True
+
+    def test_valid_model_different_provider_returns_correct_provider(self):
+        """Test that different providers for the same model return different configurations"""
+        # GPT 4.1 has multiple providers with different configurations
+        openai_provider = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="gpt_4_1",
+        )
+        openrouter_provider = built_in_models_from_provider(
+            provider_name=ModelProviderName.openrouter,
+            model_name="gpt_4_1",
+        )
+
+        assert openai_provider is not None
+        assert openrouter_provider is not None
+        assert openai_provider.name == ModelProviderName.openai
+        assert openrouter_provider.name == ModelProviderName.openrouter
+        assert openai_provider.model_id == "gpt-4.1"
+        assert openrouter_provider.model_id == "openai/gpt-4.1"
+
+    def test_invalid_model_name_returns_none(self):
+        """Test that invalid model name returns None"""
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="invalid_model_name",
+        )
+        assert result is None
+
+    def test_valid_model_invalid_provider_returns_none(self):
+        """Test that valid model but invalid provider returns None"""
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.gemini_api,  # GPT 4.1 doesn't have gemini_api provider
+            model_name="gpt_4_1",
+        )
+        assert result is None
+
+    def test_model_with_single_provider(self):
+        """Test model that only has one provider"""
+        # Find a model with only one provider for this test
+        model = get_model_by_name(ModelName.gpt_4_1_nano)
+        assert len(model.providers) >= 1  # Verify it has providers
+
+        first_provider = model.providers[0]
+        result = built_in_models_from_provider(
+            provider_name=first_provider.name,
+            model_name="gpt_4_1_nano",
+        )
+        assert result is not None
+        assert result.name == first_provider.name
+        assert result.model_id == first_provider.model_id
+
+    def test_model_with_multiple_providers(self):
+        """Test model that has multiple providers"""
+        # GPT 4.1 has multiple providers
+        openai_result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="gpt_4_1",
+        )
+        openrouter_result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openrouter,
+            model_name="gpt_4_1",
+        )
+        azure_result = built_in_models_from_provider(
+            provider_name=ModelProviderName.azure_openai,
+            model_name="gpt_4_1",
+        )
+
+        assert openai_result is not None
+        assert openrouter_result is not None
+        assert azure_result is not None
+        assert openai_result.name == ModelProviderName.openai
+        assert openrouter_result.name == ModelProviderName.openrouter
+        assert azure_result.name == ModelProviderName.azure_openai
+
+    def test_case_sensitive_model_name(self):
+        """Test that model name matching is case sensitive"""
+        # Should return None for case-mismatched model name
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="GPT_4_1",  # Wrong case
+        )
+        assert result is None
+
+    def test_provider_specific_attributes(self):
+        """Test that provider-specific attributes are correctly returned"""
+        # Test a provider with specific attributes like multimodal_capable
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="gpt_4_1",
+        )
+        assert result is not None
+        assert result.multimodal_capable is True
+        assert result.supports_doc_extraction is True
+        assert result.multimodal_mime_types is not None
+        assert len(result.multimodal_mime_types) > 0
+
+    def test_provider_without_special_attributes(self):
+        """Test provider that doesn't have special attributes"""
+        # Test a simpler provider configuration
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="gpt_4_1_nano",
+        )
+        assert result is not None
+        assert result.multimodal_capable is False  # Should be False for nano
+        assert result.supports_doc_extraction is False
+
+    @pytest.mark.parametrize(
+        "model_name,provider,expected_model_id",
+        [
+            ("gpt_4o", ModelProviderName.openai, "gpt-4o"),
+            (
+                "claude_3_5_haiku",
+                ModelProviderName.anthropic,
+                "claude-3-5-haiku-20241022",
+            ),
+            ("gemini_2_5_pro", ModelProviderName.gemini_api, "gemini-2.5-pro"),
+            ("llama_3_1_8b", ModelProviderName.groq, "llama-3.1-8b-instant"),
+        ],
+    )
+    def test_parametrized_valid_combinations(
+        self, model_name, provider, expected_model_id
+    ):
+        """Test multiple valid model/provider combinations"""
+        result = built_in_models_from_provider(
+            provider_name=provider,
+            model_name=model_name,
+        )
+        assert result is not None
+        assert result.name == provider
+        assert result.model_id == expected_model_id
+
+    def test_empty_string_model_name(self):
+        """Test that empty string model name returns None"""
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name="",
+        )
+        assert result is None
+
+    def test_none_model_name(self):
+        """Test that None model name returns None"""
+        result = built_in_models_from_provider(
+            provider_name=ModelProviderName.openai,
+            model_name=None,  # type: ignore
+        )
+        assert result is None
+
+    def test_all_built_in_models_have_valid_providers(self):
+        """Test that all built-in models have at least one valid provider"""
+        for model in built_in_models:
+            assert len(model.providers) > 0, f"Model {model.name} has no providers"
+            for provider in model.providers:
+                # Test that we can retrieve each provider
+                result = built_in_models_from_provider(
+                    provider_name=provider.name,
+                    model_name=model.name,
+                )
+                assert result is not None, (
+                    f"Could not retrieve provider {provider.name} for model {model.name}"
+                )
+                assert result.name == provider.name
+                assert result.model_id == provider.model_id
+
+
 def test_uncensored():
     """Test that uncensored is set correctly"""
     model = get_model_by_name(ModelName.grok_3_mini)
@@ -179,6 +360,27 @@ def test_uncensored():
         assert provider.suggested_for_uncensored_data_gen
 
 
+def test_no_empty_multimodal_mime_types():
+    """Ensure that multimodal fields are self-consistent as they are interdependent"""
+    for model in built_in_models:
+        for provider in model.providers:
+            # a multimodal model should always have mime types it supports
+            if provider.multimodal_capable:
+                assert provider.multimodal_mime_types is not None
+                assert len(provider.multimodal_mime_types) > 0
+
+            # a model that specifies mime types is necessarily multimodal
+            if (
+                provider.multimodal_mime_types is not None
+                and len(provider.multimodal_mime_types) > 0
+            ):
+                assert provider.multimodal_capable
+
+            # a model that supports doc extraction is necessarily multimodal
+            if provider.supports_doc_extraction:
+                assert provider.multimodal_capable
+
+
 def test_no_reasoning_for_structured_output():
     """Test that no reasoning is returned for structured output"""
     # get all models