kiln-ai 0.15.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

kiln_ai/adapters/model_adapters/test_saving_adapter_results.py

@@ -11,14 +11,15 @@ from kiln_ai.datamodel import (
     DataSourceType,
     Project,
     Task,
+    Usage,
 )
 from kiln_ai.datamodel.task import RunConfig
 from kiln_ai.utils.config import Config
 
 
 class MockAdapter(BaseAdapter):
-    async def _run(self, input: dict | str) -> dict | str:
-        return RunOutput(output="Test output", intermediate_outputs=None)
+    async def _run(self, input: dict | str) -> tuple[RunOutput, Usage | None]:
+        return RunOutput(output="Test output", intermediate_outputs=None), None
 
     def adapter_name(self) -> str:
         return "mock_adapter"
@@ -45,6 +46,7 @@ def adapter(test_task):
             model_name="phi_3_5",
             model_provider_name="ollama",
             prompt_id="simple_chain_of_thought_prompt_builder",
+            structured_output_mode="json_schema",
         ),
     )
 
@@ -101,6 +103,9 @@ def test_save_run_isolation(test_task, adapter):
         reloaded_output.source.properties["prompt_id"]
         == "simple_chain_of_thought_prompt_builder"
     )
+    assert reloaded_output.source.properties["structured_output_mode"] == "json_schema"
+    assert reloaded_output.source.properties["temperature"] == 1.0
+    assert reloaded_output.source.properties["top_p"] == 1.0
     # Run again, with same input and different output. Should create a new TaskRun.
     different_run_output = RunOutput(
         output="Different output", intermediate_outputs=None
@@ -227,3 +232,40 @@ async def test_autosave_true(test_task, adapter):
             output.source.properties["prompt_id"]
             == "simple_chain_of_thought_prompt_builder"
         )
+        assert output.source.properties["structured_output_mode"] == "json_schema"
+        assert output.source.properties["temperature"] == 1.0
+        assert output.source.properties["top_p"] == 1.0
+
+
+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,
+            model_name="gpt-4",
+            model_provider_name="openai",
+            prompt_id="simple_prompt_builder",
+            temperature=0.7,
+            top_p=0.9,
+            structured_output_mode="json_schema",
+        ),
+    )
+
+    input_data = "Test input"
+    output_data = "Test output"
+    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()
+
+    # Verify custom values are preserved in properties
+    output = task_run.output
+    assert output.source.properties["adapter_name"] == "mock_adapter"
+    assert output.source.properties["model_name"] == "gpt-4"
+    assert output.source.properties["model_provider"] == "openai"
+    assert output.source.properties["prompt_id"] == "simple_prompt_builder"
+    assert output.source.properties["structured_output_mode"] == "json_schema"
+    assert output.source.properties["temperature"] == 0.7
+    assert output.source.properties["top_p"] == 0.9

kiln_ai/adapters/model_adapters/test_structured_output.py

@@ -12,11 +12,12 @@ from kiln_ai.adapters.ml_model_list import (
 from kiln_ai.adapters.model_adapters.base_adapter import (
     BaseAdapter,
     RunOutput,
+    Usage,
 )
 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.task import RunConfig, RunConfigProperties
 from kiln_ai.datamodel.test_json_schema import json_joke_schema, json_triangle_schema
 
 
@@ -50,12 +51,13 @@ class MockAdapter(BaseAdapter):
                 model_name="phi_3_5",
                 model_provider_name="ollama",
                 prompt_id="simple_chain_of_thought_prompt_builder",
+                structured_output_mode="json_schema",
             ),
         )
         self.response = response
 
-    async def _run(self, input: str) -> RunOutput:
-        return RunOutput(output=self.response, intermediate_outputs=None)
+    async def _run(self, input: str) -> tuple[RunOutput, Usage | None]:
+        return RunOutput(output=self.response, intermediate_outputs=None), None
 
     def adapter_name(self) -> str:
         return "mock_adapter"
@@ -145,7 +147,15 @@ 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)
+    a = adapter_for_task(
+        task,
+        run_config_properties=RunConfigProperties(
+            model_name=model_name,
+            model_provider_name=provider,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="unknown",
+        ),
+    )
     try:
         run = await a.invoke("Cows")  # a joke about cows
         parsed = json.loads(run.output.output)
@@ -196,10 +206,12 @@ def build_structured_input_test_task(tmp_path: Path):
     return task
 
 
-async def run_structured_input_test(tmp_path: Path, model_name: str, provider: str):
+async def run_structured_input_test(
+    tmp_path: Path, model_name: str, provider: str, prompt_id: PromptId
+):
     task = build_structured_input_test_task(tmp_path)
     try:
-        await run_structured_input_task(task, model_name, provider)
+        await run_structured_input_task(task, model_name, provider, prompt_id)
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
@@ -208,43 +220,54 @@ async def run_structured_input_test(tmp_path: Path, model_name: str, provider: s
         raise e
 
 
-async def run_structured_input_task(
+async def run_structured_input_task_no_validation(
     task: datamodel.Task,
     model_name: str,
     provider: str,
-    prompt_id: PromptId | None = None,
+    prompt_id: PromptId,
 ):
     a = adapter_for_task(
         task,
-        model_name=model_name,
-        provider=provider,
-        prompt_id=prompt_id,
+        run_config_properties=RunConfigProperties(
+            model_name=model_name,
+            model_provider_name=provider,
+            prompt_id=prompt_id,
+            structured_output_mode="unknown",
+        ),
     )
     with pytest.raises(ValueError):
         # not structured input in dictionary
         await a.invoke("a=1, b=2, c=3")
-    with pytest.raises(
-        ValueError,
-        match="This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema.",
-    ):
+    with pytest.raises(ValueError, match="This task requires a specific input"):
         # invalid structured input
         await a.invoke({"a": 1, "b": 2, "d": 3})
 
     try:
         run = await a.invoke({"a": 2, "b": 2, "c": 2})
         response = run.output.output
+        return response, a
     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,
+    prompt_id: PromptId,
+):
+    response, a = await run_structured_input_task_no_validation(
+        task, model_name, provider, prompt_id
+    )
     assert response is not None
     if isinstance(response, str):
         assert "[[equilateral]]" in response
     else:
         assert response["is_equilateral"] is True
-
     expected_pb_name = "simple_prompt_builder"
     if prompt_id is not None:
         expected_pb_name = prompt_id
@@ -271,7 +294,9 @@ async def test_structured_input_gpt_4o_mini(tmp_path):
 async def test_all_built_in_models_structured_input(
     tmp_path, model_name, provider_name
 ):
-    await run_structured_input_test(tmp_path, model_name, provider_name)
+    await run_structured_input_test(
+        tmp_path, model_name, provider_name, "simple_prompt_builder"
+    )
 
 
 @pytest.mark.paid
@@ -325,6 +350,11 @@ When asked for a final result, this is the format (for an equilateral example):
 """
     task.output_json_schema = json.dumps(triangle_schema)
     task.save_to_file()
-    await run_structured_input_task(
+    response, adapter = await run_structured_input_task_no_validation(
         task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
     )
+
+    formatted_response = json.loads(response)
+    assert formatted_response["is_equilateral"] is True
+    assert formatted_response["is_scalene"] is False
+    assert formatted_response["is_obtuse"] is False

kiln_ai/adapters/parsers/base_parser.py

@@ -2,9 +2,6 @@ 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.

kiln_ai/adapters/parsers/parser_registry.py

@@ -6,14 +6,16 @@ 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]:
+def model_parser_from_id(parser_id: ModelParserID | None) -> BaseParser:
     """
     Get a model parser from its ID.
     """
     match parser_id:
         case None:
-            return BaseParser
+            return BaseParser()
         case ModelParserID.r1_thinking:
-            return R1ThinkingParser
+            return R1ThinkingParser()
+        case ModelParserID.optional_r1_thinking:
+            return R1ThinkingParser(allow_missing_thinking=True)
         case _:
             raise_exhaustive_enum_error(parser_id)

kiln_ai/adapters/parsers/r1_parser.py

@@ -7,6 +7,9 @@ class R1ThinkingParser(BaseParser):
     START_TAG = "<think>"
     END_TAG = "</think>"
 
+    def __init__(self, allow_missing_thinking: bool = False):
+        self.allow_missing_thinking = allow_missing_thinking
+
     def parse_output(self, original_output: RunOutput) -> RunOutput:
         """
         Parse the <think> </think> tags from the response into the intermediate and final outputs.
@@ -27,6 +30,14 @@ class R1ThinkingParser(BaseParser):
             original_output.intermediate_outputs is not None
             and "reasoning" in original_output.intermediate_outputs
         ):
+            # sometimes the output and reasoning are wrapped in newlines
+            if isinstance(original_output.output, str):
+                original_output.output = original_output.output.strip()
+
+            original_output.intermediate_outputs["reasoning"] = (
+                original_output.intermediate_outputs["reasoning"].strip()
+            )
+
             return original_output
 
         # This parser only works for strings
@@ -39,7 +50,10 @@ class R1ThinkingParser(BaseParser):
         # Find the thinking tags
         think_end = cleaned_response.find(self.END_TAG)
         if think_end == -1:
-            raise ValueError("Missing </think> tag")
+            if self.allow_missing_thinking:
+                return original_output
+            else:
+                raise ValueError("Missing </think> tag")
 
         think_tag_start = cleaned_response.find(self.START_TAG)
         if think_tag_start == -1:
@@ -66,7 +80,8 @@ class R1ThinkingParser(BaseParser):
 
         # Add thinking content to intermediate outputs if it exists
         intermediate_outputs = original_output.intermediate_outputs or {}
-        intermediate_outputs["reasoning"] = thinking_content
+        if thinking_content is not None and len(thinking_content) > 0:
+            intermediate_outputs["reasoning"] = thinking_content
 
         return RunOutput(
             output=result,

kiln_ai/adapters/parsers/request_formatters.py

@@ -0,0 +1,40 @@
+import json
+from typing import Dict, Protocol
+
+from kiln_ai.adapters.ml_model_list import ModelFormatterID
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+
+class RequestFormatter(Protocol):
+    def format_input(self, original_input: Dict | str) -> Dict | str:
+        """
+        Method for formatting the input to a model.
+        """
+        ...
+
+
+class Qwen3StyleNoThinkFormatter:
+    def format_input(self, original_input: Dict | str) -> Dict | str:
+        """
+        Format the input to a model for Qwen3 /no_think instruction
+        """
+        formatted_input = (
+            original_input
+            if isinstance(original_input, str)
+            else json.dumps(original_input, indent=2)
+        )
+
+        return formatted_input + "\n\n/no_think"
+
+
+def request_formatter_from_id(
+    formatter_id: ModelFormatterID,
+) -> RequestFormatter:
+    """
+    Get a model parser from its ID.
+    """
+    match formatter_id:
+        case ModelFormatterID.qwen3_style_no_think:
+            return Qwen3StyleNoThinkFormatter()
+        case _:
+            raise_exhaustive_enum_error(formatter_id)

kiln_ai/adapters/parsers/test_parser_registry.py

@@ -28,5 +28,5 @@ def test_model_parser_from_id_invalid():
 )
 def test_model_parser_from_id_parametrized(parser_id, expected_class):
     """Test all valid parser IDs using parametrize."""
-    parser_class = model_parser_from_id(parser_id)
-    assert parser_class == expected_class
+    parser = model_parser_from_id(parser_id)
+    assert isinstance(parser, expected_class)

kiln_ai/adapters/parsers/test_r1_parser.py

@@ -46,6 +46,21 @@ def test_response_with_whitespace(parser):
     assert parsed.output.strip() == "This is the result"
 
 
+def test_empty_thinking_content(parser):
+    response = RunOutput(
+        output="""
+        <think>
+
+        </think>
+            This is the result
+    """,
+        intermediate_outputs=None,
+    )
+    parsed = parser.parse_output(response)
+    assert "reasoning" not in parsed.intermediate_outputs
+    assert parsed.output.strip() == "This is the result"
+
+
 def test_missing_start_tag(parser):
     parsed = parser.parse_output(
         RunOutput(output="Some content</think>result", intermediate_outputs=None)
@@ -86,7 +101,7 @@ def test_empty_thinking_content(parser):
         output="<think></think>This is the result", intermediate_outputs=None
     )
     parsed = parser.parse_output(response)
-    assert parsed.intermediate_outputs == {"reasoning": ""}
+    assert "reasoning" not in parsed.intermediate_outputs
     assert parsed.output == "This is the result"
 
 
@@ -154,3 +169,31 @@ def test_intermediate_outputs(parser):
         )
     )
     assert out.intermediate_outputs["reasoning"] == "Some content"
+
+
+def test_strip_newlines(parser):
+    # certain providers via LiteLLM for example, add newlines to the output
+    # and to the reasoning. This tests that we strip those newlines.
+    response = RunOutput(
+        output="\n\nSome content",
+        intermediate_outputs={
+            "reasoning": "\n\nSome thinking\n\n",
+        },
+    )
+    parsed = parser.parse_output(response)
+    assert parsed.output == "Some content"
+    assert parsed.intermediate_outputs["reasoning"] == "Some thinking"
+
+
+def test_strip_newlines_with_structured_output(parser):
+    # certain providers via LiteLLM for example, add newlines to the output
+    # and to the reasoning. This tests that we strip those newlines.
+    response = RunOutput(
+        output={"some_key": "Some content"},
+        intermediate_outputs={
+            "reasoning": "\n\nSome thinking\n\n",
+        },
+    )
+    parsed = parser.parse_output(response)
+    assert parsed.output == {"some_key": "Some content"}
+    assert parsed.intermediate_outputs["reasoning"] == "Some thinking"

kiln_ai/adapters/parsers/test_request_formatters.py

@@ -0,0 +1,76 @@
+import pytest
+
+from kiln_ai.adapters.ml_model_list import ModelFormatterID
+from kiln_ai.adapters.parsers.request_formatters import (
+    Qwen3StyleNoThinkFormatter,
+    request_formatter_from_id,
+)
+
+
+@pytest.fixture
+def qwen_formatter():
+    return Qwen3StyleNoThinkFormatter()
+
+
+def test_qwen_formatter_string_input(qwen_formatter):
+    input_text = "Hello world"
+    formatted = qwen_formatter.format_input(input_text)
+    assert formatted == "Hello world\n\n/no_think"
+
+
+def test_qwen_formatter_dict_input(qwen_formatter):
+    input_dict = {"key": "value", "nested": {"inner": "data"}}
+    formatted = qwen_formatter.format_input(input_dict)
+    expected = """{
+  "key": "value",
+  "nested": {
+    "inner": "data"
+  }
+}
+
+/no_think"""
+    assert formatted == expected
+
+
+def test_qwen_formatter_empty_input(qwen_formatter):
+    # Test empty string
+    assert qwen_formatter.format_input("") == "\n\n/no_think"
+
+    # Test empty dict
+    assert qwen_formatter.format_input({}) == "{}\n\n/no_think"
+
+
+def test_qwen_formatter_special_characters(qwen_formatter):
+    input_text = "Special chars: !@#$%^&*()_+思"
+    formatted = qwen_formatter.format_input(input_text)
+    assert formatted == "Special chars: !@#$%^&*()_+思\n\n/no_think"
+
+
+def test_qwen_formatter_multiline_string(qwen_formatter):
+    input_text = """Line 1
+    Line 2
+    Line 3"""
+    formatted = qwen_formatter.format_input(input_text)
+    assert (
+        formatted
+        == """Line 1
+    Line 2
+    Line 3
+
+/no_think"""
+    )
+
+
+def test_request_formatter_factory():
+    # Test valid formatter ID
+    formatter = request_formatter_from_id(ModelFormatterID.qwen3_style_no_think)
+    assert isinstance(formatter, Qwen3StyleNoThinkFormatter)
+
+    # Test that the formatter works
+    assert formatter.format_input("test") == "test\n\n/no_think"
+
+
+def test_request_formatter_factory_invalid_id():
+    # Test with an invalid enum value by using a string that doesn't exist in the enum
+    with pytest.raises(ValueError, match="Unhandled enum value"):
+        request_formatter_from_id("invalid_formatter_id")  # type: ignore

kiln_ai/adapters/prompt_builders.py

@@ -1,6 +1,4 @@
-import json
 from abc import ABCMeta, abstractmethod
-from typing import Dict
 
 from kiln_ai.datamodel import PromptGenerators, PromptId, Task, TaskRun
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
@@ -53,20 +51,6 @@ class BasePromptBuilder(metaclass=ABCMeta):
         """
         pass
 
-    def build_user_message(self, input: Dict | str) -> str:
-        """Build a user message from the input.
-
-        Args:
-            input (Union[Dict, str]): The input to format into a message.
-
-        Returns:
-            str: The formatted user message.
-        """
-        if isinstance(input, Dict):
-            return f"The input is:\n{json.dumps(input, indent=2, ensure_ascii=False)}"
-
-        return f"The input is:\n{input}"
-
     def chain_of_thought_prompt(self) -> str | None:
         """Build and return the chain of thought prompt string.
 
@@ -101,7 +85,6 @@ class SimplePromptBuilder(BasePromptBuilder):
         """
         base_prompt = self.task.instruction
 
-        # TODO: this is just a quick version. Formatting and best practices TBD
         if len(self.task.requirements) > 0:
             base_prompt += (
                 "\n\nYour response should respect the following requirements:\n"
@@ -113,6 +96,18 @@ class SimplePromptBuilder(BasePromptBuilder):
         return base_prompt
 
 
+class ShortPromptBuilder(BasePromptBuilder):
+    """A prompt builder that includes a the base prompt but excludes the requirements."""
+
+    def build_base_prompt(self) -> str:
+        """Build a short prompt with just the base prompt, no requirements.
+
+        Returns:
+            str: The constructed prompt string.
+        """
+        return self.task.instruction
+
+
 class MultiShotPromptBuilder(BasePromptBuilder):
     """A prompt builder that includes multiple examples in the prompt."""
 
@@ -414,6 +409,8 @@ def prompt_builder_from_id(prompt_id: PromptId, task: Task) -> BasePromptBuilder
     match typed_prompt_generator:
         case PromptGenerators.SIMPLE:
             return SimplePromptBuilder(task)
+        case PromptGenerators.SHORT:
+            return ShortPromptBuilder(task)
         case PromptGenerators.FEW_SHOT:
             return FewShotPromptBuilder(task)
         case PromptGenerators.MULTI_SHOT:

kiln_ai/adapters/provider_tools.py

@@ -1,3 +1,4 @@
+import logging
 from dataclasses import dataclass
 from typing import Dict, List
 
@@ -5,6 +6,7 @@ from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     KilnModelProvider,
     ModelName,
+    ModelParserID,
     ModelProviderName,
     StructuredOutputMode,
     built_in_models,
@@ -16,10 +18,14 @@ 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
 
+logger = logging.getLogger(__name__)
+
 
 async def provider_enabled(provider_name: ModelProviderName) -> bool:
     if provider_name == ModelProviderName.ollama:
@@ -162,6 +168,10 @@ def kiln_model_provider_from(
     # For custom registry, get the provider name and model name from the model id
     if provider_name == ModelProviderName.kiln_custom_registry:
         provider_name, name = parse_custom_model_id(name)
+    else:
+        logger.warning(
+            f"Unexpected model/provider pair. Will treat as custom model but check your model settings. Provider: {provider_name}/{name}"
+        )
 
     # Custom/untested model. Set untested, and build a ModelProvider at runtime
     if provider_name is None:
@@ -176,12 +186,15 @@ def kiln_model_provider_from(
         supports_data_gen=False,
         untested_model=True,
         model_id=name,
+        # We don't know the structured output mode for custom models, so we default to json_instructions which is the only one that works everywhere.
+        structured_output_mode=StructuredOutputMode.json_instructions,
     )
 
 
-def lite_llm_config(
-    model_id: str,
+def lite_llm_config_for_openai_compatible(
+    run_config_properties: RunConfigProperties,
 ) -> LiteLlmConfig:
+    model_id = run_config_properties.model_name
     try:
         openai_provider_name, model_id = model_id.split("::")
     except Exception:
@@ -205,10 +218,16 @@ def lite_llm_config(
             f"OpenAI compatible provider {openai_provider_name} has no base URL"
         )
 
+    # Update a copy of the run config properties to use the openai compatible provider
+    updated_run_config_properties = run_config_properties.model_copy(deep=True)
+    updated_run_config_properties.model_provider_name = (
+        ModelProviderName.openai_compatible
+    )
+    updated_run_config_properties.model_name = model_id
+
     return LiteLlmConfig(
         # OpenAI compatible, with a custom base URL
-        model_name=model_id,
-        provider_name=ModelProviderName.openai_compatible,
+        run_config_properties=updated_run_config_properties,
         base_url=base_url,
         additional_body_options={
             "api_key": api_key,
@@ -257,6 +276,14 @@ def finetune_from_id(model_id: str) -> Finetune:
     return fine_tune
 
 
+def parser_from_data_strategy(
+    data_strategy: ChatStrategy,
+) -> ModelParserID | None:
+    if data_strategy == ChatStrategy.single_turn_r1_thinking:
+        return ModelParserID.r1_thinking
+    return None
+
+
 def finetune_provider_model(
     model_id: str,
 ) -> KilnModelProvider:
@@ -266,6 +293,14 @@ def finetune_provider_model(
     model_provider = KilnModelProvider(
         name=provider,
         model_id=fine_tune.fine_tune_model_id,
+        parser=parser_from_data_strategy(fine_tune.data_strategy),
+        reasoning_capable=(
+            fine_tune.data_strategy
+            in [
+                ChatStrategy.single_turn_r1_thinking,
+            ]
+        ),
+        tuned_chat_strategy=fine_tune.data_strategy,
     )
 
     if provider == ModelProviderName.vertex and fine_tune.fine_tune_model_id: