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

kiln_ai/adapters/model_adapters/test_saving_adapter_results.py

@@ -13,7 +13,7 @@ from kiln_ai.datamodel import (
     Task,
     Usage,
 )
-from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.datamodel.task import RunConfigProperties
 from kiln_ai.utils.config import Config
 
 
@@ -41,8 +41,8 @@ def test_task(tmp_path):
 @pytest.fixture
 def adapter(test_task):
     return MockAdapter(
-        run_config=RunConfig(
-            task=test_task,
+        task=test_task,
+        run_config=RunConfigProperties(
             model_name="phi_3_5",
             model_provider_name="ollama",
             prompt_id="simple_chain_of_thought_prompt_builder",
@@ -240,8 +240,8 @@ async def test_autosave_true(test_task, adapter):
 def test_properties_for_task_output_custom_values(test_task):
     """Test that _properties_for_task_output includes custom temperature, top_p, and structured_output_mode"""
     adapter = MockAdapter(
-        run_config=RunConfig(
-            task=test_task,
+        task=test_task,
+        run_config=RunConfigProperties(
             model_name="gpt-4",
             model_provider_name="openai",
             prompt_id="simple_prompt_builder",

kiln_ai/adapters/model_adapters/test_structured_output.py

@@ -1,8 +1,10 @@
 import json
 from pathlib import Path
 from typing import Dict
+from unittest.mock import Mock, patch
 
 import pytest
+from litellm.types.utils import ModelResponse
 
 import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.adapter_registry import adapter_for_task
@@ -11,7 +13,7 @@ from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput,
 from kiln_ai.adapters.ollama_tools import ollama_online
 from kiln_ai.adapters.test_prompt_adaptors import get_all_models_and_providers
 from kiln_ai.datamodel import PromptId
-from kiln_ai.datamodel.task import RunConfig, RunConfigProperties
+from kiln_ai.datamodel.task import RunConfigProperties
 from kiln_ai.datamodel.test_json_schema import json_joke_schema, json_triangle_schema
 
 
@@ -40,8 +42,8 @@ async def test_structured_output_ollama(tmp_path, model_name):
 class MockAdapter(BaseAdapter):
     def __init__(self, kiln_task: datamodel.Task, response: Dict | str | None):
         super().__init__(
-            run_config=RunConfig(
-                task=kiln_task,
+            task=kiln_task,
+            run_config=RunConfigProperties(
                 model_name="phi_3_5",
                 model_provider_name="ollama",
                 prompt_id="simple_chain_of_thought_prompt_builder",
@@ -259,6 +261,7 @@ async def run_structured_input_task(
     model_name: str,
     provider: str,
     prompt_id: PromptId,
+    verify_trace_cot: bool = False,
 ):
     response, a, run = await run_structured_input_task_no_validation(
         task, model_name, provider, prompt_id
@@ -282,6 +285,32 @@ async def run_structured_input_task(
         assert "reasoning" in run.intermediate_outputs
         assert isinstance(run.intermediate_outputs["reasoning"], str)
 
+    # Check the trace
+    trace = run.trace
+    assert trace is not None
+    if verify_trace_cot:
+        assert len(trace) == 5
+        assert trace[0]["role"] == "system"
+        assert "You are an assistant which classifies a triangle" in trace[0]["content"]
+        assert trace[1]["role"] == "user"
+        assert trace[2]["role"] == "assistant"
+        assert trace[2].get("tool_calls") is None
+        assert trace[3]["role"] == "user"
+        assert trace[4]["role"] == "assistant"
+        assert trace[4].get("tool_calls") is None
+    else:
+        assert len(trace) == 3
+        assert trace[0]["role"] == "system"
+        assert "You are an assistant which classifies a triangle" in trace[0]["content"]
+        assert trace[1]["role"] == "user"
+        json_content = json.loads(trace[1]["content"])
+        assert json_content["a"] == 2
+        assert json_content["b"] == 2
+        assert json_content["c"] == 2
+        assert trace[2]["role"] == "assistant"
+        assert trace[2].get("tool_calls") is None
+        assert "[[equilateral]]" in trace[2]["content"]
+
 
 @pytest.mark.paid
 async def test_structured_input_gpt_4o_mini(tmp_path):
@@ -299,15 +328,92 @@ async def test_all_built_in_models_structured_input(
     )
 
 
+async def test_all_built_in_models_structured_input_mocked(tmp_path):
+    mock_response = ModelResponse(
+        model="gpt-4o-mini",
+        choices=[
+            {
+                "message": {
+                    "content": "The answer is [[equilateral]]",
+                }
+            }
+        ],
+    )
+
+    # Mock the Config.shared() method to return a mock config with required attributes
+    mock_config = Mock()
+    mock_config.open_ai_api_key = "mock_api_key"
+    mock_config.user_id = "test_user"
+
+    with (
+        patch(
+            "litellm.acompletion",
+            side_effect=[mock_response],
+        ),
+        patch("kiln_ai.utils.config.Config.shared", return_value=mock_config),
+    ):
+        await run_structured_input_test(
+            tmp_path, "llama_3_1_8b", "groq", "simple_prompt_builder"
+        )
+
+
 @pytest.mark.paid
 @pytest.mark.ollama
 @pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
 async def test_structured_input_cot_prompt_builder(tmp_path, model_name, provider_name):
     task = build_structured_input_test_task(tmp_path)
     await run_structured_input_task(
-        task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
+        task,
+        model_name,
+        provider_name,
+        "simple_chain_of_thought_prompt_builder",
+        verify_trace_cot=True,
+    )
+
+
+async def test_structured_input_cot_prompt_builder_mocked(tmp_path):
+    task = build_structured_input_test_task(tmp_path)
+    mock_response_1 = ModelResponse(
+        model="gpt-4o-mini",
+        choices=[
+            {
+                "message": {
+                    "content": "I'm thinking real hard... oh!",
+                }
+            }
+        ],
+    )
+    mock_response_2 = ModelResponse(
+        model="gpt-4o-mini",
+        choices=[
+            {
+                "message": {
+                    "content": "After thinking, I've decided the answer is [[equilateral]]",
+                }
+            }
+        ],
     )
 
+    # Mock the Config.shared() method to return a mock config with required attributes
+    mock_config = Mock()
+    mock_config.open_ai_api_key = "mock_api_key"
+    mock_config.user_id = "test_user"
+
+    with (
+        patch(
+            "litellm.acompletion",
+            side_effect=[mock_response_1, mock_response_2],
+        ),
+        patch("kiln_ai.utils.config.Config.shared", return_value=mock_config),
+    ):
+        await run_structured_input_task(
+            task,
+            "llama_3_1_8b",
+            "groq",
+            "simple_chain_of_thought_prompt_builder",
+            verify_trace_cot=True,
+        )
+
 
 @pytest.mark.paid
 @pytest.mark.ollama

kiln_ai/adapters/parsers/__init__.py

@@ -7,4 +7,4 @@ Parsing utilities for JSON and models with custom output formats (R1, etc.)
 
 from . import base_parser, json_parser, r1_parser
 
-__all__ = ["r1_parser", "base_parser", "json_parser"]
+__all__ = ["base_parser", "json_parser", "r1_parser"]

kiln_ai/adapters/provider_tools.py

@@ -2,6 +2,9 @@ import logging
 from dataclasses import dataclass
 from typing import Dict, List
 
+from kiln_ai.adapters.docker_model_runner_tools import (
+    get_docker_model_runner_connection,
+)
 from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     KilnModelProvider,
@@ -14,10 +17,10 @@ from kiln_ai.adapters.model_adapters.litellm_config import LiteLlmConfig
 from kiln_ai.adapters.ollama_tools import get_ollama_connection
 from kiln_ai.datamodel import Finetune, Task
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
-from kiln_ai.datamodel.registry import project_from_id
 from kiln_ai.datamodel.task import RunConfigProperties
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+from kiln_ai.utils.project_utils import project_from_id
 
 logger = logging.getLogger(__name__)
 
@@ -32,6 +35,15 @@ async def provider_enabled(provider_name: ModelProviderName) -> bool:
         except Exception:
             return False
 
+    if provider_name == ModelProviderName.docker_model_runner:
+        try:
+            conn = await get_docker_model_runner_connection()
+            return conn is not None and (
+                len(conn.supported_models) > 0 or len(conn.untested_models) > 0
+            )
+        except Exception:
+            return False
+
     provider_warning = provider_warnings.get(provider_name)
     if provider_warning is None:
         return False
@@ -377,6 +389,8 @@ def provider_name_from_id(id: str) -> str:
                 return "SiliconFlow"
             case ModelProviderName.cerebras:
                 return "Cerebras"
+            case ModelProviderName.docker_model_runner:
+                return "Docker Model Runner"
             case _:
                 # triggers pyright warning if I miss a case
                 raise_exhaustive_enum_error(enum_id)

kiln_ai/adapters/repair/test_repair_task.py

@@ -229,21 +229,20 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         "rating": 8,
     }
 
+    run_config = RunConfigProperties(
+        model_name="llama_3_1_8b",
+        model_provider_name="ollama",
+        prompt_id="simple_prompt_builder",
+        structured_output_mode="json_schema",
+    )
+
     with patch.object(LiteLlmAdapter, "_run", new_callable=AsyncMock) as mock_run:
         mock_run.return_value = (
             RunOutput(output=mocked_output, intermediate_outputs=None),
             None,
         )
 
-        adapter = adapter_for_task(
-            repair_task,
-            RunConfigProperties(
-                model_name="llama_3_1_8b",
-                model_provider_name="ollama",
-                prompt_id="simple_prompt_builder",
-                structured_output_mode="json_schema",
-            ),
-        )
+        adapter = adapter_for_task(repair_task, run_config)
 
         run = await adapter.invoke(repair_task_input.model_dump())
 
@@ -264,6 +263,10 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     }
     assert run.input_source.type == DataSourceType.human
     assert "created_by" in run.input_source.properties
+    assert run.output.source is not None
+    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()
 
     # Verify that the mock was called
     mock_run.assert_called_once()

kiln_ai/adapters/run_output.py

@@ -3,9 +3,12 @@ from typing import Dict
 
 from litellm.types.utils import ChoiceLogprobs
 
+from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam
+
 
 @dataclass
 class RunOutput:
     output: Dict | str
     intermediate_outputs: Dict[str, str] | None
     output_logprobs: ChoiceLogprobs | None = None
+    trace: list[ChatCompletionMessageParam] | None = None

kiln_ai/adapters/test_adapter_registry.py

@@ -8,6 +8,7 @@ from kiln_ai.adapters.ml_model_list import ModelProviderName
 from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
 from kiln_ai.adapters.model_adapters.litellm_adapter import LiteLlmAdapter
 from kiln_ai.adapters.provider_tools import kiln_model_provider_from
+from kiln_ai.datamodel.datamodel_enums import StructuredOutputMode
 from kiln_ai.datamodel.task import RunConfigProperties
 
 
@@ -17,6 +18,9 @@ def mock_config():
         mock.shared.return_value.open_ai_api_key = "test-openai-key"
         mock.shared.return_value.open_router_api_key = "test-openrouter-key"
         mock.shared.return_value.siliconflow_cn_api_key = "test-siliconflow-key"
+        mock.shared.return_value.docker_model_runner_base_url = (
+            "http://localhost:12434/engines/llama.cpp"
+        )
         yield mock
 
 
@@ -108,7 +112,7 @@ def test_siliconflow_adapter_creation(mock_config, basic_task):
         == ModelProviderName.siliconflow_cn
     )
     assert adapter.config.default_headers == {
-        "HTTP-Referer": "https://getkiln.ai/siliconflow",
+        "HTTP-Referer": "https://kiln.tech/siliconflow",
         "X-Title": "KilnAI",
     }
 
@@ -260,3 +264,78 @@ async def test_fine_tune_provider(mock_config, basic_task, mock_finetune_from_id
     )
     # The actual model name from the fine tune object
     assert provider.model_id == "test-model"
+
+
+def test_docker_model_runner_adapter_creation(mock_config, basic_task):
+    """Test Docker Model Runner adapter creation with default and custom base URL."""
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        run_config_properties=RunConfigProperties(
+            model_name="llama_3_2_3b",
+            model_provider_name=ModelProviderName.docker_model_runner,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode=StructuredOutputMode.json_schema,
+        ),
+    )
+
+    assert isinstance(adapter, LiteLlmAdapter)
+    assert adapter.config.run_config_properties.model_name == "llama_3_2_3b"
+    assert adapter.config.additional_body_options == {"api_key": "DMR"}
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.docker_model_runner
+    )
+    assert adapter.config.base_url == "http://localhost:12434/engines/llama.cpp/v1"
+    assert adapter.config.default_headers is None
+
+
+def test_docker_model_runner_adapter_creation_with_custom_url(mock_config, basic_task):
+    """Test Docker Model Runner adapter creation with custom base URL."""
+    mock_config.shared.return_value.docker_model_runner_base_url = (
+        "http://custom:8080/engines/llama.cpp"
+    )
+
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        run_config_properties=RunConfigProperties(
+            model_name="llama_3_2_3b",
+            model_provider_name=ModelProviderName.docker_model_runner,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode=StructuredOutputMode.json_schema,
+        ),
+    )
+
+    assert isinstance(adapter, LiteLlmAdapter)
+    assert adapter.config.run_config_properties.model_name == "llama_3_2_3b"
+    assert adapter.config.additional_body_options == {"api_key": "DMR"}
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.docker_model_runner
+    )
+    assert adapter.config.base_url == "http://custom:8080/engines/llama.cpp/v1"
+    assert adapter.config.default_headers is None
+
+
+def test_docker_model_runner_adapter_creation_with_none_url(mock_config, basic_task):
+    """Test Docker Model Runner adapter creation when config URL is None."""
+    mock_config.shared.return_value.docker_model_runner_base_url = None
+
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        run_config_properties=RunConfigProperties(
+            model_name="llama_3_2_3b",
+            model_provider_name=ModelProviderName.docker_model_runner,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode=StructuredOutputMode.json_schema,
+        ),
+    )
+
+    assert isinstance(adapter, LiteLlmAdapter)
+    assert adapter.config.run_config_properties.model_name == "llama_3_2_3b"
+    assert adapter.config.additional_body_options == {"api_key": "DMR"}
+    assert (
+        adapter.config.run_config_properties.model_provider_name
+        == ModelProviderName.docker_model_runner
+    )
+    assert adapter.config.base_url == "http://localhost:12434/engines/llama.cpp/v1"
+    assert adapter.config.default_headers is None