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

kiln_ai/adapters/test_docker_model_runner_tools.py

@@ -0,0 +1,305 @@
+from unittest.mock import AsyncMock, Mock, patch
+
+import httpx
+import openai
+import pytest
+
+from kiln_ai.adapters.docker_model_runner_tools import (
+    DockerModelRunnerConnection,
+    docker_model_runner_base_url,
+    parse_docker_model_runner_models,
+)
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+
+
+def test_docker_model_runner_base_url_default():
+    """Test that the default base URL is returned when no config is set."""
+    with patch("kiln_ai.adapters.docker_model_runner_tools.Config") as mock_config:
+        mock_config.shared().docker_model_runner_base_url = None
+        result = docker_model_runner_base_url()
+        assert result == "http://localhost:12434/engines/llama.cpp"
+
+
+def test_docker_model_runner_base_url_from_config():
+    """Test that the configured base URL is returned when set."""
+    with patch("kiln_ai.adapters.docker_model_runner_tools.Config") as mock_config:
+        mock_config.shared().docker_model_runner_base_url = (
+            "http://custom:8080/engines/llama.cpp"
+        )
+        result = docker_model_runner_base_url()
+        assert result == "http://custom:8080/engines/llama.cpp"
+
+
+def test_parse_docker_model_runner_models_with_supported_models():
+    """Test parsing Docker Model Runner models response with supported models."""
+    # Create mock OpenAI Model objects
+    mock_models = [
+        Mock(id="ai/llama3.2:3B-Q4_K_M"),
+        Mock(id="ai/qwen3:8B-Q4_K_M"),
+        Mock(id="ai/gemma3n:4B-Q4_K_M"),
+        Mock(id="unsupported-model"),
+    ]
+
+    with patch(
+        "kiln_ai.adapters.docker_model_runner_tools.built_in_models"
+    ) as mock_built_in_models:
+        # Mock built-in models with Docker Model Runner providers
+        mock_model = Mock()
+        mock_provider = Mock()
+        mock_provider.name = ModelProviderName.docker_model_runner
+        mock_provider.model_id = "ai/llama3.2:3B-Q4_K_M"
+        mock_model.providers = [mock_provider]
+        mock_built_in_models.__iter__ = Mock(return_value=iter([mock_model]))
+
+        result = parse_docker_model_runner_models(mock_models)  # type: ignore
+
+        assert result is not None
+        assert result.message == "Docker Model Runner connected"
+        assert "ai/llama3.2:3B-Q4_K_M" in result.supported_models
+        assert "unsupported-model" in result.untested_models
+
+
+def test_parse_docker_model_runner_models_no_models():
+    """Test parsing Docker Model Runner models response with no models."""
+    mock_models = []
+
+    result = parse_docker_model_runner_models(mock_models)
+
+    assert result is not None
+    assert "no supported models are available" in result.message
+    assert len(result.supported_models) == 0
+    assert len(result.untested_models) == 0
+
+
+def test_docker_model_runner_connection_all_models():
+    """Test that DockerModelRunnerConnection.all_models() returns both supported and untested models."""
+    connection = DockerModelRunnerConnection(
+        message="Test",
+        supported_models=["model1", "model2"],
+        untested_models=["model3", "model4"],
+    )
+
+    all_models = connection.all_models()
+    assert all_models == ["model1", "model2", "model3", "model4"]
+
+
+@pytest.mark.asyncio
+async def test_docker_model_runner_online_success():
+    """Test that docker_model_runner_online returns True when service is available."""
+    with patch(
+        "kiln_ai.adapters.docker_model_runner_tools.httpx.AsyncClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_response = Mock()
+        mock_response.raise_for_status.return_value = None
+        mock_client.get = AsyncMock(return_value=mock_response)
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        from kiln_ai.adapters.docker_model_runner_tools import (
+            docker_model_runner_online,
+        )
+
+        result = await docker_model_runner_online()
+
+        assert result is True
+        mock_client.get.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_docker_model_runner_online_failure():
+    """Test that docker_model_runner_online returns False when service is unavailable."""
+    with patch(
+        "kiln_ai.adapters.docker_model_runner_tools.httpx.AsyncClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_client.get = AsyncMock(side_effect=httpx.RequestError("Connection error"))
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        from kiln_ai.adapters.docker_model_runner_tools import (
+            docker_model_runner_online,
+        )
+
+        result = await docker_model_runner_online()
+
+        assert result is False
+
+
+@pytest.mark.asyncio
+async def test_get_docker_model_runner_connection_success():
+    """Test get_docker_model_runner_connection with successful connection."""
+    from kiln_ai.adapters.docker_model_runner_tools import (
+        get_docker_model_runner_connection,
+    )
+
+    # Mock OpenAI client and models response
+    mock_model = Mock()
+    mock_model.id = "ai/llama3.2:3B-Q4_K_M"
+    mock_models_response = [mock_model]
+
+    with (
+        patch(
+            "kiln_ai.adapters.docker_model_runner_tools.openai.OpenAI"
+        ) as mock_openai,
+        patch(
+            "kiln_ai.adapters.docker_model_runner_tools.parse_docker_model_runner_models"
+        ) as mock_parse,
+        patch(
+            "kiln_ai.adapters.docker_model_runner_tools.docker_model_runner_base_url"
+        ) as mock_base_url,
+    ):
+        mock_base_url.return_value = "http://localhost:12434/engines"
+        mock_client = Mock()
+        mock_client.models.list.return_value = mock_models_response
+        mock_openai.return_value = mock_client
+
+        expected_connection = DockerModelRunnerConnection(
+            message="Connected",
+            supported_models=["ai/llama3.2:3B-Q4_K_M"],
+            untested_models=[],
+        )
+        mock_parse.return_value = expected_connection
+
+        result = await get_docker_model_runner_connection()
+
+        assert result == expected_connection
+        mock_openai.assert_called_once_with(
+            api_key="dummy",
+            base_url="http://localhost:12434/engines/v1",
+            max_retries=0,
+        )
+        mock_parse.assert_called_once_with(mock_models_response)
+
+
+@pytest.mark.asyncio
+async def test_get_docker_model_runner_connection_with_custom_url():
+    """Test get_docker_model_runner_connection with custom URL."""
+    from kiln_ai.adapters.docker_model_runner_tools import (
+        get_docker_model_runner_connection,
+    )
+
+    # Mock OpenAI client and models response
+    mock_model = Mock()
+    mock_model.id = "ai/llama3.2:3B-Q4_K_M"
+    mock_models_response = [mock_model]
+
+    with (
+        patch(
+            "kiln_ai.adapters.docker_model_runner_tools.openai.OpenAI"
+        ) as mock_openai,
+        patch(
+            "kiln_ai.adapters.docker_model_runner_tools.parse_docker_model_runner_models"
+        ) as mock_parse,
+    ):
+        mock_client = Mock()
+        mock_client.models.list.return_value = mock_models_response
+        mock_openai.return_value = mock_client
+
+        expected_connection = DockerModelRunnerConnection(
+            message="Connected",
+            supported_models=["ai/llama3.2:3B-Q4_K_M"],
+            untested_models=[],
+        )
+        mock_parse.return_value = expected_connection
+
+        custom_url = "http://custom:8080/engines/llama.cpp"
+        result = await get_docker_model_runner_connection(custom_url)
+
+        assert result == expected_connection
+        mock_openai.assert_called_once_with(
+            api_key="dummy",
+            base_url=f"{custom_url}/v1",
+            max_retries=0,
+        )
+        mock_parse.assert_called_once_with(mock_models_response)
+
+
+@pytest.mark.asyncio
+async def test_get_docker_model_runner_connection_api_error():
+    """Test get_docker_model_runner_connection with API error."""
+    from kiln_ai.adapters.docker_model_runner_tools import (
+        get_docker_model_runner_connection,
+    )
+
+    with patch(
+        "kiln_ai.adapters.docker_model_runner_tools.openai.OpenAI"
+    ) as mock_openai:
+        mock_client = Mock()
+        mock_client.models.list.side_effect = openai.APIConnectionError(request=Mock())
+        mock_openai.return_value = mock_client
+
+        result = await get_docker_model_runner_connection()
+
+        assert result is None
+
+
+@pytest.mark.asyncio
+async def test_get_docker_model_runner_connection_connection_error():
+    """Test get_docker_model_runner_connection with connection error."""
+    from kiln_ai.adapters.docker_model_runner_tools import (
+        get_docker_model_runner_connection,
+    )
+
+    with patch(
+        "kiln_ai.adapters.docker_model_runner_tools.openai.OpenAI"
+    ) as mock_openai:
+        mock_client = Mock()
+        mock_client.models.list.side_effect = httpx.RequestError("Connection error")
+        mock_openai.return_value = mock_client
+
+        result = await get_docker_model_runner_connection()
+
+        assert result is None
+
+
+@pytest.mark.asyncio
+async def test_get_docker_model_runner_connection_http_error():
+    """Test get_docker_model_runner_connection with HTTP error."""
+    from kiln_ai.adapters.docker_model_runner_tools import (
+        get_docker_model_runner_connection,
+    )
+
+    with patch(
+        "kiln_ai.adapters.docker_model_runner_tools.openai.OpenAI"
+    ) as mock_openai:
+        mock_client = Mock()
+        mock_client.models.list.side_effect = httpx.RequestError("HTTP error")
+        mock_openai.return_value = mock_client
+
+        result = await get_docker_model_runner_connection()
+
+        assert result is None
+
+
+def test_docker_model_runner_model_installed_true():
+    """Test docker_model_runner_model_installed returns True when model is installed."""
+    from kiln_ai.adapters.docker_model_runner_tools import (
+        docker_model_runner_model_installed,
+    )
+
+    connection = DockerModelRunnerConnection(
+        message="Test",
+        supported_models=["model1", "model2"],
+        untested_models=["model3", "model4"],
+    )
+
+    # Test model in supported_models
+    assert docker_model_runner_model_installed(connection, "model1") is True
+
+    # Test model in untested_models
+    assert docker_model_runner_model_installed(connection, "model3") is True
+
+
+def test_docker_model_runner_model_installed_false():
+    """Test docker_model_runner_model_installed returns False when model is not installed."""
+    from kiln_ai.adapters.docker_model_runner_tools import (
+        docker_model_runner_model_installed,
+    )
+
+    connection = DockerModelRunnerConnection(
+        message="Test",
+        supported_models=["model1", "model2"],
+        untested_models=["model3", "model4"],
+    )
+
+    # Test model not in any list
+    assert docker_model_runner_model_installed(connection, "nonexistent_model") is False

kiln_ai/adapters/test_ml_model_list.py

@@ -1,7 +1,10 @@
+from collections import Counter
+
 import pytest
 
 from kiln_ai.adapters.ml_model_list import (
     ModelName,
+    built_in_models,
     default_structured_output_mode_for_model_provider,
     get_model_by_name,
 )
@@ -131,7 +134,7 @@ class TestDefaultStructuredOutputModeForModelProvider:
             ("llama_3_1_8b", ModelProviderName.groq, StructuredOutputMode.default),
             (
                 "qwq_32b",
-                ModelProviderName.fireworks_ai,
+                ModelProviderName.together_ai,
                 StructuredOutputMode.json_instructions,
             ),
         ],
@@ -174,3 +177,50 @@ def test_uncensored():
     for provider in model.providers:
         assert provider.uncensored
         assert provider.suggested_for_uncensored_data_gen
+
+
+def test_no_reasoning_for_structured_output():
+    """Test that no reasoning is returned for structured output"""
+    # get all models
+    for model in built_in_models:
+        for provider in model.providers:
+            if provider.reasoning_optional_for_structured_output is not None:
+                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_prompt_adaptors.py

@@ -119,14 +119,16 @@ async def test_mock_returning_run(tmp_path):
             choices=[{"message": {"content": "mock response"}}],
         )
 
+        run_config = RunConfigProperties(
+            model_name="custom_model",
+            model_provider_name="ollama",
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+        )
+
         adapter = LiteLlmAdapter(
             config=LiteLlmConfig(
-                run_config_properties=RunConfigProperties(
-                    model_name="custom_model",
-                    model_provider_name="ollama",
-                    prompt_id="simple_prompt_builder",
-                    structured_output_mode="json_schema",
-                ),
+                run_config_properties=run_config,
                 base_url="http://localhost:11434",
                 additional_body_options={"api_key": "test_key"},
             ),
@@ -140,7 +142,9 @@ async def test_mock_returning_run(tmp_path):
     assert run.id is not None
     assert run.input == "You are a mock, send me the response!"
     assert run.output.output == "mock response"
+    assert run.input_source is not None
     assert "created_by" in run.input_source.properties
+    assert run.output.source is not None
     assert run.output.source.properties == {
         "adapter_name": "kiln_openai_compatible_adapter",
         "model_name": "custom_model",
@@ -150,6 +154,9 @@ async def test_mock_returning_run(tmp_path):
         "temperature": 1.0,
         "top_p": 1.0,
     }
+    assert run.output.source.run_config is not None
+    saved_run_config = run.output.source.run_config.model_dump()
+    assert saved_run_config == run_config.model_dump()
 
 
 @pytest.mark.paid

kiln_ai/adapters/test_provider_tools.py

@@ -2,6 +2,7 @@ from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
+from kiln_ai.adapters.docker_model_runner_tools import DockerModelRunnerConnection
 from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     ModelName,
@@ -25,11 +26,7 @@ from kiln_ai.adapters.provider_tools import (
     provider_name_from_id,
     provider_warnings,
 )
-from kiln_ai.datamodel import (
-    Finetune,
-    StructuredOutputMode,
-    Task,
-)
+from kiln_ai.datamodel import Finetune, StructuredOutputMode, Task
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.task import RunConfigProperties
 
@@ -199,6 +196,7 @@ def test_provider_name_from_id_case_sensitivity():
         (ModelProviderName.ollama, "Ollama"),
         (ModelProviderName.openai, "OpenAI"),
         (ModelProviderName.fireworks_ai, "Fireworks AI"),
+        (ModelProviderName.siliconflow_cn, "SiliconFlow"),
         (ModelProviderName.kiln_fine_tune, "Fine Tuned Models"),
         (ModelProviderName.kiln_custom_registry, "Custom Models"),
     ],
@@ -420,6 +418,17 @@ async def test_builtin_model_from_invalid_provider(mock_config):
     assert provider is None
 
 
+@pytest.mark.asyncio
+async def test_builtin_model_future_proof():
+    """Test handling of a model that doesn't exist yet but could be added over the air"""
+    with patch("kiln_ai.adapters.provider_tools.built_in_models") as mock_models:
+        mock_models.__iter__.return_value = []
+
+        # should not find it, but should not raise an error
+        result = builtin_model_from("gpt_99")
+        assert result is None
+
+
 @pytest.mark.asyncio
 async def test_builtin_model_from_model_no_providers():
     """Test handling of a model with no providers"""
@@ -433,10 +442,8 @@ async def test_builtin_model_from_model_no_providers():
         )
         mock_models.__iter__.return_value = [mock_model]
 
-        with pytest.raises(ValueError) as exc_info:
-            await builtin_model_from(ModelName.phi_3_5)
-
-        assert str(exc_info.value) == f"Model {ModelName.phi_3_5} has no providers"
+        result = builtin_model_from(ModelName.phi_3_5)
+        assert result is None
 
 
 @pytest.mark.asyncio
@@ -461,7 +468,7 @@ def test_finetune_provider_model_success(mock_project, mock_task, mock_finetune)
     assert provider.model_id == "ft:gpt-3.5-turbo:custom:model-123"
     assert provider.structured_output_mode == StructuredOutputMode.json_schema
     assert provider.reasoning_capable is False
-    assert provider.parser == None
+    assert provider.parser is None
 
 
 def test_finetune_provider_model_success_final_and_intermediate(
@@ -476,7 +483,7 @@ def test_finetune_provider_model_success_final_and_intermediate(
     assert provider.model_id == "ft:gpt-3.5-turbo:custom:model-123"
     assert provider.structured_output_mode == StructuredOutputMode.json_schema
     assert provider.reasoning_capable is False
-    assert provider.parser == None
+    assert provider.parser is None
 
 
 def test_finetune_provider_model_success_r1_compatible(
@@ -590,7 +597,7 @@ def test_finetune_provider_model_structured_mode(
     assert provider.model_id == "fireworks-model-123"
     assert provider.structured_output_mode == expected_mode
     assert provider.reasoning_capable is False
-    assert provider.parser == None
+    assert provider.parser is None
 
 
 def test_openai_compatible_provider_config(mock_shared_config):
@@ -925,3 +932,57 @@ def test_finetune_provider_model_vertex_ai(mock_project, mock_task, mock_finetun
     # Verify the model_id is transformed into openai/endpoint_id format
     assert provider.model_id == "openai/456"
     assert provider.structured_output_mode == StructuredOutputMode.json_mode
+
+
+@pytest.mark.asyncio
+async def test_provider_enabled_docker_model_runner_success():
+    """Test provider_enabled for Docker Model Runner with successful connection"""
+    with patch(
+        "kiln_ai.adapters.provider_tools.get_docker_model_runner_connection",
+        new_callable=AsyncMock,
+    ) as mock_get_docker:
+        # Mock successful Docker Model Runner connection with models
+        mock_get_docker.return_value = DockerModelRunnerConnection(
+            message="Connected",
+            supported_models=["llama-3.2-3b-instruct"],
+            untested_models=[],
+        )
+
+        result = await provider_enabled(ModelProviderName.docker_model_runner)
+        assert result is True
+
+
+@pytest.mark.asyncio
+async def test_provider_enabled_docker_model_runner_no_models():
+    """Test provider_enabled for Docker Model Runner with no models"""
+    with patch(
+        "kiln_ai.adapters.provider_tools.get_docker_model_runner_connection",
+        new_callable=AsyncMock,
+    ) as mock_get_docker:
+        # Mock Docker Model Runner connection but with no models
+        mock_get_docker.return_value = DockerModelRunnerConnection(
+            message="Connected but no models", supported_models=[], untested_models=[]
+        )
+
+        result = await provider_enabled(ModelProviderName.docker_model_runner)
+        assert result is False
+
+
+@pytest.mark.asyncio
+async def test_provider_enabled_docker_model_runner_connection_error():
+    """Test provider_enabled for Docker Model Runner with connection error"""
+    with patch(
+        "kiln_ai.adapters.provider_tools.get_docker_model_runner_connection",
+        new_callable=AsyncMock,
+    ) as mock_get_docker:
+        # Mock Docker Model Runner connection failure
+        mock_get_docker.side_effect = Exception("Connection failed")
+
+        result = await provider_enabled(ModelProviderName.docker_model_runner)
+        assert result is False
+
+
+def test_provider_name_from_id_docker_model_runner():
+    """Test provider_name_from_id for Docker Model Runner"""
+    result = provider_name_from_id(ModelProviderName.docker_model_runner)
+    assert result == "Docker Model Runner"