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

kiln_ai/adapters/test_ml_embedding_model_list.py

@@ -0,0 +1,429 @@
+import pytest
+
+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
+
+
+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"
+        )
+
+
+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 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")
+
+    @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",
+            ),
+        ],
+    )
+    def test_parametrized_model_retrieval(
+        self, model_name, expected_family, expected_friendly_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
+
+
+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,
+        )
+
+        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 is None
+
+    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
+
+    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
+
+    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,
+        )
+
+        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 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
+
+    @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,
+            ),
+        ],
+    )
+    def test_parametrized_provider_retrieval(
+        self, provider_name, model_name, expected_model_id, expected_dimensions
+    ):
+        """Test retrieving providers with parametrized test cases"""
+        provider = built_in_embedding_models_from_provider(provider_name, model_name)
+
+        assert provider is not None
+        assert provider.model_id == expected_model_id
+        assert provider.n_dimensions == expected_dimensions

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,37 @@ 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:
+        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