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

kiln_ai/adapters/test_ml_model_list.py

@@ -1,8 +1,11 @@
+from collections import Counter
+
 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,
 )
@@ -132,7 +135,7 @@ class TestDefaultStructuredOutputModeForModelProvider:
             ("llama_3_1_8b", ModelProviderName.groq, StructuredOutputMode.default),
             (
                 "qwq_32b",
-                ModelProviderName.fireworks_ai,
+                ModelProviderName.together_ai,
                 StructuredOutputMode.json_instructions,
             ),
         ],
@@ -159,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)
@@ -177,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
@@ -186,3 +400,39 @@ def test_no_reasoning_for_structured_output():
                 assert provider.reasoning_capable, (
                     f"{model.name} {provider.name} has reasoning_optional_for_structured_output but is not reasoning capable. This field should only be defined for models that are reasoning capable."
                 )
+
+
+def test_unique_providers_per_model():
+    """Test that each model can only have one entry per provider"""
+    for model in built_in_models:
+        provider_names = [provider.name for provider in model.providers]
+        unique_provider_names = set(provider_names)
+
+        if len(provider_names) != len(unique_provider_names):
+            # Find which providers have duplicates
+            provider_counts = Counter(provider_names)
+            duplicates = {
+                name: count for name, count in provider_counts.items() if count > 1
+            }
+
+            # Show details about duplicates
+            duplicate_details = []
+            for provider_name, count in duplicates.items():
+                duplicate_providers = [
+                    p for p in model.providers if p.name == provider_name
+                ]
+                model_ids = [p.model_id for p in duplicate_providers]
+                duplicate_details.append(
+                    f"{provider_name} (appears {count} times with model_ids: {model_ids})"
+                )
+
+            assert False, (
+                f"Model {model.name} has duplicate providers:\n"
+                f"Expected: 1 entry per provider\n"
+                f"Found: {len(provider_names)} total entries, {len(unique_provider_names)} unique providers\n"
+                f"Duplicates: {', '.join(duplicate_details)}\n"
+                f"This suggests either:\n"
+                f"1. A bug where the same provider is accidentally duplicated, or\n"
+                f"2. Intentional design where the same provider offers different model variants\n"
+                f"If this is intentional, the test should be updated to allow multiple entries per provider."
+            )

kiln_ai/adapters/test_ollama_tools.py

@@ -1,13 +1,85 @@
 import json
+from unittest.mock import patch
 
+import pytest
+
+from kiln_ai.adapters.ml_embedding_model_list import (
+    KilnEmbeddingModel,
+    KilnEmbeddingModelProvider,
+)
+
+# Mock data for testing - using proper Pydantic model instances
+from kiln_ai.adapters.ml_model_list import KilnModel, KilnModelProvider
 from kiln_ai.adapters.ollama_tools import (
     OllamaConnection,
+    ollama_embedding_model_installed,
     ollama_model_installed,
     parse_ollama_tags,
 )
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
 
+MOCK_BUILT_IN_MODELS = [
+    KilnModel(
+        family="phi",
+        name="phi3.5",
+        friendly_name="phi3.5",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                model_id="phi3.5",
+                ollama_model_aliases=None,
+            )
+        ],
+    ),
+    KilnModel(
+        family="gemma",
+        name="gemma2",
+        friendly_name="gemma2",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                model_id="gemma2:2b",
+                ollama_model_aliases=None,
+            )
+        ],
+    ),
+    KilnModel(
+        family="llama",
+        name="llama3.1",
+        friendly_name="llama3.1",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                model_id="llama3.1",
+                ollama_model_aliases=None,
+            )
+        ],
+    ),
+]
 
-def test_parse_ollama_tags_no_models():
+MOCK_BUILT_IN_EMBEDDING_MODELS = [
+    KilnEmbeddingModel(
+        family="gemma",
+        name="embeddinggemma",
+        friendly_name="embeddinggemma",
+        providers=[
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.ollama,
+                model_id="embeddinggemma:300m",
+                n_dimensions=768,
+                ollama_model_aliases=["embeddinggemma"],
+            )
+        ],
+    ),
+]
+
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
+def test_parse_ollama_tags_models():
     json_response = '{"models":[{"name":"scosman_net","model":"scosman_net:latest"},{"name":"phi3.5:latest","model":"phi3.5:latest","modified_at":"2024-10-02T12:04:35.191519822-04:00","size":2176178843,"digest":"61819fb370a3c1a9be6694869331e5f85f867a079e9271d66cb223acb81d04ba","details":{"parent_model":"","format":"gguf","family":"phi3","families":["phi3"],"parameter_size":"3.8B","quantization_level":"Q4_0"}},{"name":"gemma2:2b","model":"gemma2:2b","modified_at":"2024-09-09T16:46:38.64348929-04:00","size":1629518495,"digest":"8ccf136fdd5298f3ffe2d69862750ea7fb56555fa4d5b18c04e3fa4d82ee09d7","details":{"parent_model":"","format":"gguf","family":"gemma2","families":["gemma2"],"parameter_size":"2.6B","quantization_level":"Q4_0"}},{"name":"llama3.1:latest","model":"llama3.1:latest","modified_at":"2024-09-01T17:19:43.481523695-04:00","size":4661230720,"digest":"f66fc8dc39ea206e03ff6764fcc696b1b4dfb693f0b6ef751731dd4e6269046e","details":{"parent_model":"","format":"gguf","family":"llama","families":["llama"],"parameter_size":"8.0B","quantization_level":"Q4_0"}}]}'
     tags = json.loads(json_response)
     conn = parse_ollama_tags(tags)
@@ -16,7 +88,52 @@ def test_parse_ollama_tags_no_models():
     assert "llama3.1:latest" in conn.supported_models
     assert "scosman_net:latest" in conn.untested_models
 
+    # there should be no embedding models because the tags response does not include any embedding models
+    # that are in the built-in embedding models list
+    assert len(conn.supported_embedding_models) == 0
 
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
+@pytest.mark.parametrize("json_response", ["{}", '{"models": []}'])
+def test_parse_ollama_tags_no_models(json_response):
+    tags = json.loads(json_response)
+    conn = parse_ollama_tags(tags)
+    assert (
+        conn.message
+        == "Ollama is running, but no supported models are installed. Install one or more supported model, like 'ollama pull phi3.5'."
+    )
+    assert len(conn.supported_models) == 0
+    assert len(conn.untested_models) == 0
+    assert len(conn.supported_embedding_models) == 0
+
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
+def test_parse_ollama_tags_empty_models():
+    """Test parsing Ollama tags response with empty models list"""
+    json_response = '{"models": []}'
+    tags = json.loads(json_response)
+    conn = parse_ollama_tags(tags)
+
+    # Check that connection indicates no supported models
+    assert conn.supported_models == []
+    assert conn.untested_models == []
+    assert conn.supported_embedding_models == []
+    assert "no supported models are installed" in conn.message
+
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
 def test_parse_ollama_tags_only_untested_models():
     json_response = '{"models":[{"name":"scosman_net","model":"scosman_net:latest"}]}'
     tags = json.loads(json_response)
@@ -24,12 +141,17 @@ def test_parse_ollama_tags_only_untested_models():
     assert conn.supported_models == []
     assert conn.untested_models == ["scosman_net:latest"]
 
+    # there should be no embedding models because the tags response does not include any embedding models
+    # that are in the built-in embedding models list
+    assert len(conn.supported_embedding_models) == 0
+
 
 def test_ollama_model_installed():
     conn = OllamaConnection(
         supported_models=["phi3.5:latest", "gemma2:2b", "llama3.1:latest"],
         message="Connected",
         untested_models=["scosman_net:latest"],
+        supported_embedding_models=["embeddinggemma:300m"],
     )
     assert ollama_model_installed(conn, "phi3.5:latest")
     assert ollama_model_installed(conn, "phi3.5")
@@ -39,3 +161,220 @@ def test_ollama_model_installed():
     assert ollama_model_installed(conn, "scosman_net:latest")
     assert ollama_model_installed(conn, "scosman_net")
     assert not ollama_model_installed(conn, "unknown_model")
+
+    # use the ollama_embedding_model_installed for testing embedding models installed, not ollama_model_installed
+    assert not ollama_model_installed(conn, "embeddinggemma:300m")
+    assert not ollama_model_installed(conn, "embeddinggemma")
+
+
+def test_ollama_model_installed_embedding_models():
+    conn = OllamaConnection(
+        message="Connected",
+        supported_models=["phi3.5:latest", "gemma2:2b", "llama3.1:latest"],
+        untested_models=["scosman_net:latest"],
+        supported_embedding_models=["embeddinggemma:300m", "embeddinggemma:latest"],
+    )
+
+    assert ollama_embedding_model_installed(conn, "embeddinggemma:300m")
+    assert ollama_embedding_model_installed(conn, "embeddinggemma:latest")
+    assert not ollama_embedding_model_installed(conn, "unknown_embedding")
+
+    # use the ollama_model_installed for testing regular models installed, not ollama_embedding_model_installed
+    assert not ollama_embedding_model_installed(conn, "phi3.5:latest")
+    assert not ollama_embedding_model_installed(conn, "gemma2:2b")
+    assert not ollama_embedding_model_installed(conn, "llama3.1:latest")
+    assert not ollama_embedding_model_installed(conn, "scosman_net:latest")
+
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
+def test_parse_ollama_tags_with_embedding_models():
+    """Test parsing Ollama tags response that includes embedding models"""
+    json_response = """{
+        "models": [
+            {
+                "name": "phi3.5:latest",
+                "model": "phi3.5:latest"
+            },
+            {
+                "name": "embeddinggemma:300m",
+                "model": "embeddinggemma:300m"
+            },
+            {
+                "name": "embeddinggemma:latest",
+                "model": "embeddinggemma:latest"
+            },
+            {
+                "name": "unknown_embedding:latest",
+                "model": "unknown_embedding:latest"
+            }
+        ]
+    }"""
+    tags = json.loads(json_response)
+    conn = parse_ollama_tags(tags)
+
+    # Check that embedding models are properly categorized
+    assert "embeddinggemma:300m" in conn.supported_embedding_models
+    assert "embeddinggemma:latest" in conn.supported_embedding_models
+
+    # Check that regular models are still parsed correctly
+    assert "phi3.5:latest" in conn.supported_models
+
+    # Check that embedding models are NOT in the main model lists
+    assert "embeddinggemma:300m" not in conn.supported_models
+    assert "embeddinggemma:latest" not in conn.supported_models
+    assert "embeddinggemma:300m" not in conn.untested_models
+    assert "embeddinggemma:latest" not in conn.untested_models
+
+    # we assume the unknown models are normal models, not embedding models - because
+    # we don't support untested embedding models currently
+    assert "unknown_embedding:latest" not in conn.supported_embedding_models
+    assert "unknown_embedding:latest" in conn.untested_models
+
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
+def test_parse_ollama_tags_embedding_model_aliases():
+    """Test parsing Ollama tags response with embedding model aliases"""
+    json_response = """{
+        "models": [
+            {
+                "name": "embeddinggemma",
+                "model": "embeddinggemma"
+            }
+        ]
+    }"""
+    tags = json.loads(json_response)
+    conn = parse_ollama_tags(tags)
+
+    # Check that embedding model aliases are recognized
+    assert "embeddinggemma" in conn.supported_embedding_models
+
+    # Check that embedding model aliases are NOT in the main model lists
+    assert "embeddinggemma" not in conn.supported_models
+    assert "embeddinggemma" not in conn.untested_models
+
+    assert len(conn.supported_models) == 0
+    assert len(conn.untested_models) == 0
+    assert len(conn.supported_embedding_models) == 1
+
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
+def test_parse_ollama_tags_only_embedding_models():
+    """Test parsing Ollama tags response with only embedding models"""
+    json_response = """{
+        "models": [
+            {
+                "name": "embeddinggemma:300m",
+                "model": "embeddinggemma:300m"
+            }
+        ]
+    }"""
+    tags = json.loads(json_response)
+    conn = parse_ollama_tags(tags)
+
+    # Check that embedding models are found but no regular models
+    assert "embeddinggemma:300m" in conn.supported_embedding_models
+    assert conn.supported_models == []
+    assert conn.untested_models == []
+
+    # Check that embedding models are NOT in the main model lists
+    assert "embeddinggemma:300m" not in conn.supported_models
+    assert "embeddinggemma:300m" not in conn.untested_models
+
+
+def test_ollama_connection_all_embedding_models():
+    """Test OllamaConnection.all_embedding_models() method"""
+    conn = OllamaConnection(
+        message="Connected",
+        supported_models=["phi3.5:latest"],
+        untested_models=["unknown:latest"],
+        supported_embedding_models=["embeddinggemma:300m", "embeddinggemma:latest"],
+    )
+
+    embedding_models = conn.all_embedding_models()
+    assert embedding_models == ["embeddinggemma:300m", "embeddinggemma:latest"]
+
+
+def test_ollama_connection_empty_embedding_models():
+    """Test OllamaConnection.all_embedding_models() with empty list"""
+    conn = OllamaConnection(
+        message="Connected",
+        supported_models=["phi3.5:latest"],
+        untested_models=["unknown:latest"],
+        supported_embedding_models=[],
+    )
+
+    embedding_models = conn.all_embedding_models()
+    assert embedding_models == []
+
+
+@patch("kiln_ai.adapters.ollama_tools.built_in_models", MOCK_BUILT_IN_MODELS)
+@patch(
+    "kiln_ai.adapters.ollama_tools.built_in_embedding_models",
+    MOCK_BUILT_IN_EMBEDDING_MODELS,
+)
+def test_parse_ollama_tags_mixed_models_and_embeddings():
+    """Test parsing response with mix of regular models, embedding models, and unknown models"""
+    json_response = """{
+        "models": [
+            {
+                "name": "phi3.5:latest",
+                "model": "phi3.5:latest"
+            },
+            {
+                "name": "gemma2:2b",
+                "model": "gemma2:2b"
+            },
+            {
+                "name": "embeddinggemma:300m",
+                "model": "embeddinggemma:300m"
+            },
+            {
+                "name": "embeddinggemma",
+                "model": "embeddinggemma"
+            },
+            {
+                "name": "unknown_model:latest",
+                "model": "unknown_model:latest"
+            },
+            {
+                "name": "unknown_embedding:latest",
+                "model": "unknown_embedding:latest"
+            }
+        ]
+    }"""
+    tags = json.loads(json_response)
+    conn = parse_ollama_tags(tags)
+
+    # Check regular models
+    assert "phi3.5:latest" in conn.supported_models
+    assert "gemma2:2b" in conn.supported_models
+    assert "unknown_model:latest" in conn.untested_models
+
+    # Check embedding models
+    assert "embeddinggemma:300m" in conn.supported_embedding_models
+    assert "embeddinggemma" in conn.supported_embedding_models
+
+    # Check that embedding models are NOT in the main model lists
+    assert "embeddinggemma:300m" not in conn.supported_models
+    assert "embeddinggemma" not in conn.supported_models
+    assert "embeddinggemma:300m" not in conn.untested_models
+    assert "embeddinggemma" not in conn.untested_models
+
+    # Unknown embedding models should not appear in supported_embedding_models
+    assert "unknown_embedding:latest" not in conn.supported_embedding_models
+
+    # Unknown embedding models should appear in untested_models (since they're not recognized as embeddings)
+    assert "unknown_embedding:latest" not in conn.supported_models
+    assert "unknown_embedding:latest" in conn.untested_models