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

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_prompt_builders.py

@@ -359,7 +359,7 @@ def test_prompt_builder_from_id(task_with_examples):
 
     with pytest.raises(
         ValueError,
-        match="Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id'",
+        match=r"Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id'",
     ):
         prompt_builder_from_id("fine_tune_prompt::123", task)
 

kiln_ai/adapters/test_provider_tools.py

@@ -2,6 +2,8 @@ from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
+from kiln_ai.adapters.adapter_registry import litellm_core_provider_config
+from kiln_ai.adapters.docker_model_runner_tools import DockerModelRunnerConnection
 from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     ModelName,
@@ -10,6 +12,7 @@ from kiln_ai.adapters.ml_model_list import (
 )
 from kiln_ai.adapters.ollama_tools import OllamaConnection
 from kiln_ai.adapters.provider_tools import (
+    LiteLlmCoreConfig,
     builtin_model_from,
     check_provider_warnings,
     core_provider,
@@ -18,7 +21,7 @@ from kiln_ai.adapters.provider_tools import (
     finetune_provider_model,
     get_model_and_provider,
     kiln_model_provider_from,
-    lite_llm_config_for_openai_compatible,
+    lite_llm_core_config_for_provider,
     lite_llm_provider_model,
     parse_custom_model_id,
     provider_enabled,
@@ -603,7 +606,7 @@ def test_openai_compatible_provider_config(mock_shared_config):
     """Test successful creation of an OpenAI compatible provider"""
     model_id = "test_provider::gpt-4"
 
-    config = lite_llm_config_for_openai_compatible(
+    config = litellm_core_provider_config(
         RunConfigProperties(
             model_name=model_id,
             model_provider_name=ModelProviderName.openai_compatible,
@@ -638,10 +641,10 @@ def test_lite_llm_config_no_api_key(mock_shared_config):
     """Test provider creation without API key (should work as some providers don't require it, but should pass NA to LiteLLM as it requires one)"""
     model_id = "no_key_provider::gpt-4"
 
-    config = lite_llm_config_for_openai_compatible(
+    config = litellm_core_provider_config(
         RunConfigProperties(
             model_name=model_id,
-            model_provider_name=ModelProviderName.openai,
+            model_provider_name=ModelProviderName.openai_compatible,
             prompt_id="simple_prompt_builder",
             structured_output_mode="json_schema",
         )
@@ -659,7 +662,7 @@ def test_lite_llm_config_no_api_key(mock_shared_config):
 def test_lite_llm_config_invalid_id():
     """Test handling of invalid model ID format"""
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config_for_openai_compatible(
+        litellm_core_provider_config(
             RunConfigProperties(
                 model_name="invalid-id-format",
                 model_provider_name=ModelProviderName.openai_compatible,
@@ -677,7 +680,7 @@ def test_lite_llm_config_no_providers(mock_shared_config):
     mock_shared_config.return_value.openai_compatible_providers = None
 
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config_for_openai_compatible(
+        litellm_core_provider_config(
             RunConfigProperties(
                 model_name="test_provider::gpt-4",
                 model_provider_name=ModelProviderName.openai_compatible,
@@ -691,7 +694,7 @@ def test_lite_llm_config_no_providers(mock_shared_config):
 def test_lite_llm_config_provider_not_found(mock_shared_config):
     """Test handling of non-existent provider"""
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config_for_openai_compatible(
+        litellm_core_provider_config(
             RunConfigProperties(
                 model_name="unknown_provider::gpt-4",
                 model_provider_name=ModelProviderName.openai_compatible,
@@ -714,7 +717,7 @@ def test_lite_llm_config_no_base_url(mock_shared_config):
     ]
 
     with pytest.raises(ValueError) as exc_info:
-        lite_llm_config_for_openai_compatible(
+        litellm_core_provider_config(
             RunConfigProperties(
                 model_name="test_provider::gpt-4",
                 model_provider_name=ModelProviderName.openai_compatible,
@@ -931,3 +934,246 @@ 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.fixture
+def mock_config_for_lite_llm_core_config():
+    with patch("kiln_ai.adapters.provider_tools.Config") as mock:
+        config_instance = Mock()
+        mock.shared.return_value = config_instance
+
+        # Set up all the config values
+        config_instance.open_router_api_key = "test-openrouter-key"
+        config_instance.open_ai_api_key = "test-openai-key"
+        config_instance.groq_api_key = "test-groq-key"
+        config_instance.bedrock_access_key = "test-aws-access-key"
+        config_instance.bedrock_secret_key = "test-aws-secret-key"
+        config_instance.ollama_base_url = "http://test-ollama:11434"
+        config_instance.fireworks_api_key = "test-fireworks-key"
+        config_instance.anthropic_api_key = "test-anthropic-key"
+        config_instance.gemini_api_key = "test-gemini-key"
+        config_instance.vertex_project_id = "test-vertex-project"
+        config_instance.vertex_location = "us-central1"
+        config_instance.together_api_key = "test-together-key"
+        config_instance.azure_openai_api_key = "test-azure-key"
+        config_instance.azure_openai_endpoint = "https://test.openai.azure.com"
+        config_instance.huggingface_api_key = "test-hf-key"
+
+        yield mock
+
+
+@pytest.mark.parametrize(
+    "provider_name,expected_config",
+    [
+        (
+            ModelProviderName.openrouter,
+            LiteLlmCoreConfig(
+                base_url="https://openrouter.ai/api/v1",
+                additional_body_options={
+                    "api_key": "test-openrouter-key",
+                },
+                default_headers={
+                    "HTTP-Referer": "https://kiln.tech/openrouter",
+                    "X-Title": "KilnAI",
+                },
+            ),
+        ),
+        (
+            ModelProviderName.openai,
+            LiteLlmCoreConfig(additional_body_options={"api_key": "test-openai-key"}),
+        ),
+        (
+            ModelProviderName.groq,
+            LiteLlmCoreConfig(additional_body_options={"api_key": "test-groq-key"}),
+        ),
+        (
+            ModelProviderName.amazon_bedrock,
+            LiteLlmCoreConfig(
+                additional_body_options={
+                    "aws_access_key_id": "test-aws-access-key",
+                    "aws_secret_access_key": "test-aws-secret-key",
+                    "aws_region_name": "us-west-2",
+                },
+            ),
+        ),
+        (
+            ModelProviderName.ollama,
+            LiteLlmCoreConfig(
+                base_url="http://test-ollama:11434/v1",
+                additional_body_options={"api_key": "NA"},
+            ),
+        ),
+        (
+            ModelProviderName.fireworks_ai,
+            LiteLlmCoreConfig(
+                additional_body_options={"api_key": "test-fireworks-key"}
+            ),
+        ),
+        (
+            ModelProviderName.anthropic,
+            LiteLlmCoreConfig(
+                additional_body_options={"api_key": "test-anthropic-key"}
+            ),
+        ),
+        (
+            ModelProviderName.gemini_api,
+            LiteLlmCoreConfig(additional_body_options={"api_key": "test-gemini-key"}),
+        ),
+        (
+            ModelProviderName.vertex,
+            LiteLlmCoreConfig(
+                additional_body_options={
+                    "vertex_project": "test-vertex-project",
+                    "vertex_location": "us-central1",
+                },
+            ),
+        ),
+        (
+            ModelProviderName.together_ai,
+            LiteLlmCoreConfig(additional_body_options={"api_key": "test-together-key"}),
+        ),
+        (
+            ModelProviderName.azure_openai,
+            LiteLlmCoreConfig(
+                base_url="https://test.openai.azure.com",
+                additional_body_options={
+                    "api_key": "test-azure-key",
+                    "api_version": "2025-02-01-preview",
+                },
+            ),
+        ),
+        (
+            ModelProviderName.huggingface,
+            LiteLlmCoreConfig(additional_body_options={"api_key": "test-hf-key"}),
+        ),
+        (ModelProviderName.kiln_fine_tune, None),
+        (ModelProviderName.kiln_custom_registry, None),
+    ],
+)
+def test_lite_llm_core_config_for_provider(
+    mock_config_for_lite_llm_core_config, provider_name, expected_config
+):
+    config = lite_llm_core_config_for_provider(provider_name)
+    assert config == expected_config
+
+
+def test_lite_llm_core_config_for_provider_openai_compatible(
+    mock_shared_config,
+):
+    config = lite_llm_core_config_for_provider(
+        ModelProviderName.openai_compatible, "no_key_provider"
+    )
+    assert config is not None
+    assert config.base_url == "https://api.nokey.com"
+    assert config.additional_body_options == {"api_key": "NA"}
+
+
+def test_lite_llm_core_config_for_provider_openai_compatible_with_openai_compatible_provider_name(
+    mock_shared_config,
+):
+    with pytest.raises(
+        ValueError, match="OpenAI compatible provider requires a provider name"
+    ):
+        lite_llm_core_config_for_provider(ModelProviderName.openai_compatible)
+
+
+def test_lite_llm_core_config_incorrect_openai_compatible_provider_name(
+    mock_shared_config,
+):
+    with pytest.raises(
+        ValueError,
+        match="OpenAI compatible provider provider_that_does_not_exist_in_compatible_openai_providers not found",
+    ):
+        lite_llm_core_config_for_provider(
+            ModelProviderName.openai_compatible,
+            "provider_that_does_not_exist_in_compatible_openai_providers",
+        )
+
+
+def test_lite_llm_core_config_for_provider_with_string(
+    mock_config_for_lite_llm_core_config,
+):
+    # test with a string instead of an enum
+    config = lite_llm_core_config_for_provider("openai")
+    assert config == LiteLlmCoreConfig(
+        additional_body_options={"api_key": "test-openai-key"}
+    )
+
+
+def test_lite_llm_core_config_for_provider_unknown_provider():
+    with pytest.raises(ValueError, match="Unhandled enum value: unknown_provider"):
+        lite_llm_core_config_for_provider("unknown_provider")
+
+
+@patch.dict("os.environ", {"OPENROUTER_BASE_URL": "https://custom-openrouter.com"})
+def test_lite_llm_core_config_for_provider_openrouter_custom_url(
+    mock_config_for_lite_llm_core_config,
+):
+    config = lite_llm_core_config_for_provider(ModelProviderName.openrouter)
+    assert config is not None
+    assert config.base_url == "https://custom-openrouter.com"
+
+
+def test_lite_llm_core_config_for_provider_ollama_default_url(
+    mock_config_for_lite_llm_core_config,
+):
+    # Override the mock to return None for ollama_base_url
+    mock_config_for_lite_llm_core_config.shared.return_value.ollama_base_url = None
+
+    config = lite_llm_core_config_for_provider(ModelProviderName.ollama)
+    assert config is not None
+    assert config.base_url == "http://localhost:11434/v1"
+
+
+@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"