kiln-ai 0.8.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

kiln_ai/adapters/{test_saving_adapter_results.py → model_adapters/test_saving_adapter_results.py}

@@ -2,13 +2,17 @@ from unittest.mock import patch
 
 import pytest
 
-from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter, RunOutput
+from kiln_ai.adapters.model_adapters.base_adapter import (
+    BaseAdapter,
+    RunOutput,
+)
 from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
     Project,
     Task,
 )
+from kiln_ai.datamodel.task import RunConfig
 from kiln_ai.utils.config import Config
 
 
@@ -16,13 +20,8 @@ class MockAdapter(BaseAdapter):
     async def _run(self, input: dict | str) -> dict | str:
         return RunOutput(output="Test output", intermediate_outputs=None)
 
-    def adapter_info(self) -> AdapterInfo:
-        return AdapterInfo(
-            adapter_name="mock_adapter",
-            model_name="mock_model",
-            model_provider="mock_provider",
-            prompt_builder_name="mock_prompt_builder",
-        )
+    def adapter_name(self) -> str:
+        return "mock_adapter"
 
 
 @pytest.fixture
@@ -38,8 +37,19 @@ def test_task(tmp_path):
     return task
 
 
-def test_save_run_isolation(test_task):
-    adapter = MockAdapter(test_task)
+@pytest.fixture
+def adapter(test_task):
+    return MockAdapter(
+        run_config=RunConfig(
+            task=test_task,
+            model_name="phi_3_5",
+            model_provider_name="ollama",
+            prompt_id="simple_chain_of_thought_prompt_builder",
+        ),
+    )
+
+
+def test_save_run_isolation(test_task, adapter):
     input_data = "Test input"
     output_data = "Test output"
     run_output = RunOutput(
@@ -85,13 +95,12 @@ def test_save_run_isolation(test_task):
     assert reloaded_output.source.type == DataSourceType.synthetic
     assert reloaded_output.rating is None
     assert reloaded_output.source.properties["adapter_name"] == "mock_adapter"
-    assert reloaded_output.source.properties["model_name"] == "mock_model"
-    assert reloaded_output.source.properties["model_provider"] == "mock_provider"
+    assert reloaded_output.source.properties["model_name"] == "phi_3_5"
+    assert reloaded_output.source.properties["model_provider"] == "ollama"
     assert (
-        reloaded_output.source.properties["prompt_builder_name"]
-        == "mock_prompt_builder"
+        reloaded_output.source.properties["prompt_id"]
+        == "simple_chain_of_thought_prompt_builder"
     )
-
     # Run again, with same input and different output. Should create a new TaskRun.
     different_run_output = RunOutput(
         output="Different output", intermediate_outputs=None
@@ -101,13 +110,6 @@ def test_save_run_isolation(test_task):
     assert len(test_task.runs()) == 2
     assert "Different output" in set(run.output.output for run in test_task.runs())
 
-    # run again with same input and same output. Should not create a new TaskRun.
-    task_output = adapter.generate_run(input_data, None, run_output)
-    task_output.save_to_file()
-    assert len(test_task.runs()) == 2
-    assert "Different output" in set(run.output.output for run in test_task.runs())
-    assert output_data in set(run.output.output for run in test_task.runs())
-
     # run again with input of different type. Should create a new TaskRun and TaskOutput.
     task_output = adapter.generate_run(
         input_data,
@@ -116,7 +118,7 @@ def test_save_run_isolation(test_task):
             properties={
                 "model_name": "mock_model",
                 "model_provider": "mock_provider",
-                "prompt_builder_name": "mock_prompt_builder",
+                "prompt_id": "mock_prompt_builder",
                 "adapter_name": "mock_adapter",
             },
         ),
@@ -130,14 +132,41 @@ def test_save_run_isolation(test_task):
     assert output_data in set(run.output.output for run in test_task.runs())
 
 
+def test_generate_run_non_ascii(test_task, adapter):
+    input_data = {"key": "input with non-ascii character: 你好"}
+    output_data = {"key": "output with non-ascii character: 你好"}
+    run_output = RunOutput(
+        output=output_data,
+        intermediate_outputs=None,
+    )
+
+    task_run = adapter.generate_run(
+        input=input_data, input_source=None, run_output=run_output
+    )
+    task_run.save_to_file()
+
+    # as these values are saved as strings, they should properly represent the non-ascii characters
+    assert task_run.input == '{"key": "input with non-ascii character: 你好"}'
+    assert task_run.output.output == '{"key": "output with non-ascii character: 你好"}'
+
+    # check that the stringified unicode strings can be read back from the file
+    reloaded_task = Task.load_from_file(test_task.path)
+    reloaded_runs = reloaded_task.runs()
+    assert len(reloaded_runs) == 1
+    reloaded_run = reloaded_runs[0]
+    assert reloaded_run.input == '{"key": "input with non-ascii character: 你好"}'
+    assert (
+        reloaded_run.output.output == '{"key": "output with non-ascii character: 你好"}'
+    )
+
+
 @pytest.mark.asyncio
-async def test_autosave_false(test_task):
+async def test_autosave_false(test_task, adapter):
     with patch("kiln_ai.utils.config.Config.shared") as mock_shared:
         mock_config = mock_shared.return_value
         mock_config.autosave_runs = False
         mock_config.user_id = "test_user"
 
-        adapter = MockAdapter(test_task)
         input_data = "Test input"
 
         run = await adapter.invoke(input_data)
@@ -150,13 +179,31 @@ async def test_autosave_false(test_task):
 
 
 @pytest.mark.asyncio
-async def test_autosave_true(test_task):
+async def test_autosave_true_with_disabled(test_task, adapter):
+    with patch("kiln_ai.utils.config.Config.shared") as mock_shared:
+        mock_config = mock_shared.return_value
+        mock_config.autosave_runs = True
+        mock_config.user_id = "test_user"
+
+        input_data = "Test input"
+
+        adapter.base_adapter_config.allow_saving = False
+        run = await adapter.invoke(input_data)
+
+        # Check that no runs were saved
+        assert len(test_task.runs()) == 0
+
+        # Check that the run ID is not set
+        assert run.id is None
+
+
+@pytest.mark.asyncio
+async def test_autosave_true(test_task, adapter):
     with patch("kiln_ai.utils.config.Config.shared") as mock_shared:
         mock_config = mock_shared.return_value
         mock_config.autosave_runs = True
         mock_config.user_id = "test_user"
 
-        adapter = MockAdapter(test_task)
         input_data = "Test input"
 
         run = await adapter.invoke(input_data)
@@ -174,6 +221,9 @@ async def test_autosave_true(test_task):
         assert output.output == "Test output"
         assert output.source.type == DataSourceType.synthetic
         assert output.source.properties["adapter_name"] == "mock_adapter"
-        assert output.source.properties["model_name"] == "mock_model"
-        assert output.source.properties["model_provider"] == "mock_provider"
-        assert output.source.properties["prompt_builder_name"] == "mock_prompt_builder"
+        assert output.source.properties["model_name"] == "phi_3_5"
+        assert output.source.properties["model_provider"] == "ollama"
+        assert (
+            output.source.properties["prompt_id"]
+            == "simple_chain_of_thought_prompt_builder"
+        )

kiln_ai/adapters/{test_structured_output.py → model_adapters/test_structured_output.py}

@@ -1,22 +1,22 @@
+import json
 from pathlib import Path
 from typing import Dict
 
-import jsonschema
-import jsonschema.exceptions
 import pytest
 
 import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.adapter_registry import adapter_for_task
-from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter, RunOutput
 from kiln_ai.adapters.ml_model_list import (
     built_in_models,
 )
-from kiln_ai.adapters.ollama_tools import ollama_online
-from kiln_ai.adapters.prompt_builders import (
-    BasePromptBuilder,
-    SimpleChainOfThoughtPromptBuilder,
+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
 from kiln_ai.datamodel.test_json_schema import json_joke_schema, json_triangle_schema
 
 
@@ -34,9 +34,9 @@ async def test_structured_output_gpt_4o_mini(tmp_path):
     await run_structured_output_test(tmp_path, "gpt_4o_mini", "openai")
 
 
-@pytest.mark.parametrize("model_name", ["llama_3_1_8b"])
+@pytest.mark.parametrize("model_name", ["llama_3_1_8b", "gemma_2_2b"])
 @pytest.mark.ollama
-async def test_structured_output_ollama_llama(tmp_path, model_name):
+async def test_structured_output_ollama(tmp_path, model_name):
     if not await ollama_online():
         pytest.skip("Ollama API not running. Expect it running on localhost:11434")
     await run_structured_output_test(tmp_path, model_name, "ollama")
@@ -44,19 +44,21 @@ async def test_structured_output_ollama_llama(tmp_path, model_name):
 
 class MockAdapter(BaseAdapter):
     def __init__(self, kiln_task: datamodel.Task, response: Dict | str | None):
-        super().__init__(kiln_task)
+        super().__init__(
+            run_config=RunConfig(
+                task=kiln_task,
+                model_name="phi_3_5",
+                model_provider_name="ollama",
+                prompt_id="simple_chain_of_thought_prompt_builder",
+            ),
+        )
         self.response = response
 
     async def _run(self, input: str) -> RunOutput:
         return RunOutput(output=self.response, intermediate_outputs=None)
 
-    def adapter_info(self) -> AdapterInfo:
-        return AdapterInfo(
-            adapter_name="mock_adapter",
-            model_name="mock_model",
-            model_provider="mock_provider",
-            prompt_builder_name="mock_prompt_builder",
-        )
+    def adapter_name(self) -> str:
+        return "mock_adapter"
 
 
 async def test_mock_unstructred_response(tmp_path):
@@ -93,19 +95,10 @@ async def test_mock_unstructred_response(tmp_path):
         answer = await adapter.invoke("You are a mock, send me the response!")
 
 
-@pytest.mark.paid
-@pytest.mark.ollama
-@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
-async def test_all_built_in_models_structured_output(
-    tmp_path, model_name, provider_name
-):
+def check_supports_structured_output(model_name: str, provider_name: str):
     for model in built_in_models:
         if model.name != model_name:
             continue
-        if not model.supports_structured_output:
-            pytest.skip(
-                f"Skipping {model.name} because it does not support structured output"
-            )
         for provider in model.providers:
             if provider.name != provider_name:
                 continue
@@ -113,11 +106,20 @@ async def test_all_built_in_models_structured_output(
                 pytest.skip(
                     f"Skipping {model.name} {provider.name} because it does not support structured output"
                 )
-            await run_structured_output_test(tmp_path, model.name, provider.name)
             return
     raise RuntimeError(f"No model {model_name} {provider_name} found")
 
 
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_all_built_in_models_structured_output(
+    tmp_path, model_name, provider_name
+):
+    check_supports_structured_output(model_name, provider_name)
+    await run_structured_output_test(tmp_path, model_name, provider_name)
+
+
 def build_structured_output_test_task(tmp_path: Path):
     project = datamodel.Project(name="test", path=tmp_path / "test.kiln")
     project.save_to_file()
@@ -140,7 +142,14 @@ def build_structured_output_test_task(tmp_path: Path):
 async def run_structured_output_test(tmp_path: Path, model_name: str, provider: str):
     task = build_structured_output_test_task(tmp_path)
     a = adapter_for_task(task, model_name=model_name, provider=provider)
-    parsed = await a.invoke_returning_raw("Cows")  # a joke about cows
+    try:
+        parsed = await a.invoke_returning_raw("Cows")  # a joke about cows
+    except ValueError as e:
+        if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
+            pytest.skip(
+                f"Skipping {model_name} {provider} because Ollama is not running"
+            )
+        raise e
     if parsed is None or not isinstance(parsed, Dict):
         raise RuntimeError(f"structured response is not a dict: {parsed}")
     assert parsed["setup"] is not None
@@ -161,6 +170,7 @@ def build_structured_input_test_task(tmp_path: Path):
         parent=project,
         name="test task",
         instruction="You are an assistant which classifies a triangle given the lengths of its sides. If all sides are of equal length, the triangle is equilateral. If two sides are equal, the triangle is isosceles. Otherwise, it is scalene.\n\nAt the end of your response return the result in double square brackets. It should be plain text. It should be exactly one of the three following strings: '[[equilateral]]', or '[[isosceles]]', or '[[scalene]]'.",
+        thinking_prompt="Think step by step.",
     )
     task.input_json_schema = json_triangle_schema
     schema = task.input_schema()
@@ -177,37 +187,59 @@ def build_structured_input_test_task(tmp_path: Path):
 
 async def run_structured_input_test(tmp_path: Path, model_name: str, provider: str):
     task = build_structured_input_test_task(tmp_path)
-    await run_structured_input_task(task, model_name, provider)
+    try:
+        await run_structured_input_task(task, model_name, provider)
+    except ValueError as e:
+        if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
+            pytest.skip(
+                f"Skipping {model_name} {provider} because Ollama is not running"
+            )
+        raise e
 
 
 async def run_structured_input_task(
     task: datamodel.Task,
     model_name: str,
     provider: str,
-    pb: BasePromptBuilder | None = None,
+    prompt_id: PromptId | None = None,
 ):
     a = adapter_for_task(
-        task, model_name=model_name, provider=provider, prompt_builder=pb
+        task,
+        model_name=model_name,
+        provider=provider,
+        prompt_id=prompt_id,
     )
     with pytest.raises(ValueError):
         # not structured input in dictionary
         await a.invoke("a=1, b=2, c=3")
-    with pytest.raises(jsonschema.exceptions.ValidationError):
+    with pytest.raises(
+        ValueError,
+        match="This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema.",
+    ):
         # invalid structured input
         await a.invoke({"a": 1, "b": 2, "d": 3})
 
-    response = await a.invoke_returning_raw({"a": 2, "b": 2, "c": 2})
+    try:
+        response = await a.invoke_returning_raw({"a": 2, "b": 2, "c": 2})
+    except ValueError as e:
+        if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
+            pytest.skip(
+                f"Skipping {model_name} {provider} because Ollama is not running"
+            )
+        raise e
     assert response is not None
-    assert isinstance(response, str)
-    assert "[[equilateral]]" in response
-    adapter_info = a.adapter_info()
+    if isinstance(response, str):
+        assert "[[equilateral]]" in response
+    else:
+        assert response["is_equilateral"] is True
+
     expected_pb_name = "simple_prompt_builder"
-    if pb is not None:
-        expected_pb_name = pb.__class__.prompt_builder_name()
-    assert adapter_info.prompt_builder_name == expected_pb_name
-    assert adapter_info.model_name == model_name
-    assert adapter_info.model_provider == provider
-    assert adapter_info.adapter_name == "kiln_langchain_adapter"
+    if prompt_id is not None:
+        expected_pb_name = prompt_id
+    assert a.run_config.prompt_id == expected_pb_name
+
+    assert a.run_config.model_name == model_name
+    assert a.run_config.model_provider_name == provider
 
 
 @pytest.mark.paid
@@ -227,7 +259,54 @@ async def test_all_built_in_models_structured_input(
 @pytest.mark.paid
 @pytest.mark.ollama
 @pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
-async def test_structured_cot_prompt_builder(tmp_path, model_name, provider_name):
+async def test_structured_input_cot_prompt_builder(tmp_path, model_name, provider_name):
     task = build_structured_input_test_task(tmp_path)
-    pb = SimpleChainOfThoughtPromptBuilder(task)
-    await run_structured_input_task(task, model_name, provider_name, pb)
+    await run_structured_input_task(
+        task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
+    )
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_structured_output_cot_prompt_builder(
+    tmp_path, model_name, provider_name
+):
+    check_supports_structured_output(model_name, provider_name)
+    triangle_schema = {
+        "type": "object",
+        "properties": {
+            "is_equilateral": {
+                "type": "boolean",
+                "description": "True if all sides of the triangle are equal in length",
+            },
+            "is_scalene": {
+                "type": "boolean",
+                "description": "True if all sides of the triangle have different lengths",
+            },
+            "is_obtuse": {
+                "type": "boolean",
+                "description": "True if one of the angles is greater than 90 degrees",
+            },
+        },
+        "required": ["is_equilateral", "is_scalene", "is_obtuse"],
+        "additionalProperties": False,
+    }
+    task = build_structured_input_test_task(tmp_path)
+    task.instruction = """
+You are an assistant which classifies a triangle given the lengths of its sides. If all sides are of equal length, the triangle is equilateral. If two sides are equal, the triangle is isosceles. Otherwise, it is scalene.\n\n"
+
+When asked for a final result, this is the format (for an equilateral example):
+```json
+{
+    "is_equilateral": true,
+    "is_scalene": false,
+    "is_obtuse": false
+}
+```
+"""
+    task.output_json_schema = json.dumps(triangle_schema)
+    task.save_to_file()
+    await run_structured_input_task(
+        task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
+    )

kiln_ai/adapters/ollama_tools.py

@@ -1,4 +1,3 @@
-import os
 from typing import Any, List
 
 import httpx

kiln_ai/adapters/parsers/__init__.py

@@ -0,0 +1,10 @@
+"""
+# Parsers
+
+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"]

kiln_ai/adapters/parsers/base_parser.py

@@ -0,0 +1,12 @@
+from kiln_ai.adapters.run_output import RunOutput
+
+
+class BaseParser:
+    def __init__(self, structured_output: bool = False):
+        self.structured_output = structured_output
+
+    def parse_output(self, original_output: RunOutput) -> RunOutput:
+        """
+        Method for parsing the output of a model. Typically overridden by subclasses.
+        """
+        return original_output

kiln_ai/adapters/parsers/json_parser.py

@@ -0,0 +1,37 @@
+import json
+from typing import Any, Dict
+
+
+def parse_json_string(json_string: str) -> Dict[str, Any]:
+    """
+    Parse a JSON string into a dictionary. Handles multiple formats:
+    - Plain JSON
+    - JSON wrapped in ```json code blocks
+    - JSON wrapped in ``` code blocks
+
+    Args:
+        json_string: String containing JSON data, possibly wrapped in code blocks
+
+    Returns:
+        Dict containing parsed JSON data
+
+    Raises:
+        ValueError: If JSON parsing fails
+    """
+    # Remove code block markers if present
+    cleaned_string = json_string.strip()
+    if cleaned_string.startswith("```"):
+        # Split by newlines and remove first/last lines if they contain ```
+        lines = cleaned_string.split("\n")
+        if lines[0].startswith("```"):
+            lines = lines[1:]
+        if lines and lines[-1].strip() == "```":
+            lines = lines[:-1]
+        cleaned_string = "\n".join(lines)
+
+    try:
+        return json.loads(cleaned_string)
+    except json.JSONDecodeError as e:
+        raise ValueError(
+            f"This task requires JSON output but the model didn't return valid JSON. Search 'Troubleshooting Structured Data Issues' in our docs for more information. The model produced the following: {cleaned_string}"
+        ) from e

kiln_ai/adapters/parsers/parser_registry.py

@@ -0,0 +1,19 @@
+from typing import Type
+
+from kiln_ai.adapters.ml_model_list import ModelParserID
+from kiln_ai.adapters.parsers.base_parser import BaseParser
+from kiln_ai.adapters.parsers.r1_parser import R1ThinkingParser
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+
+def model_parser_from_id(parser_id: ModelParserID | None) -> Type[BaseParser]:
+    """
+    Get a model parser from its ID.
+    """
+    match parser_id:
+        case None:
+            return BaseParser
+        case ModelParserID.r1_thinking:
+            return R1ThinkingParser
+        case _:
+            raise_exhaustive_enum_error(parser_id)

kiln_ai/adapters/parsers/r1_parser.py

@@ -0,0 +1,69 @@
+from kiln_ai.adapters.parsers.base_parser import BaseParser
+from kiln_ai.adapters.parsers.json_parser import parse_json_string
+from kiln_ai.adapters.run_output import RunOutput
+
+
+class R1ThinkingParser(BaseParser):
+    START_TAG = "<think>"
+    END_TAG = "</think>"
+
+    def parse_output(self, original_output: RunOutput) -> RunOutput:
+        """
+        Parse the <think> </think> tags from the response into the intermediate and final outputs.
+
+        Args:
+            original_output: RunOutput containing the raw response string
+
+        Returns:
+            ParsedOutput containing the intermediate content (thinking content) and final result
+
+        Raises:
+            ValueError: If response format is invalid (missing tags, multiple tags, or no content after closing tag)
+        """
+        # This parser only works for strings
+        if not isinstance(original_output.output, str):
+            raise ValueError("Response must be a string for R1 parser")
+
+        # Strip whitespace and validate basic structure
+        cleaned_response = original_output.output.strip()
+        if not cleaned_response.startswith(self.START_TAG):
+            raise ValueError("Response must start with <think> tag")
+
+        # Find the thinking tags
+        think_start = cleaned_response.find(self.START_TAG)
+        think_end = cleaned_response.find(self.END_TAG)
+
+        if think_start == -1 or think_end == -1:
+            raise ValueError("Missing thinking tags")
+
+        # Check for multiple tags
+        if (
+            cleaned_response.count(self.START_TAG) > 1
+            or cleaned_response.count(self.END_TAG) > 1
+        ):
+            raise ValueError("Multiple thinking tags found")
+
+        # Extract thinking content
+        thinking_content = cleaned_response[
+            think_start + len(self.START_TAG) : think_end
+        ].strip()
+
+        # Extract result (everything after </think>)
+        result = cleaned_response[think_end + len(self.END_TAG) :].strip()
+
+        if not result or len(result) == 0:
+            raise ValueError("No content found after </think> tag")
+
+        # Parse JSON if needed
+        output = result
+        if self.structured_output:
+            output = parse_json_string(result)
+
+        # Add thinking content to intermediate outputs if it exists
+        intermediate_outputs = original_output.intermediate_outputs or {}
+        intermediate_outputs["reasoning"] = thinking_content
+
+        return RunOutput(
+            output=output,
+            intermediate_outputs=intermediate_outputs,
+        )