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

kiln_ai/adapters/parsers/test_json_parser.py

@@ -0,0 +1,81 @@
+import pytest
+
+from kiln_ai.adapters.parsers.json_parser import parse_json_string
+
+
+def test_parse_plain_json():
+    json_str = '{"key": "value", "number": 42}'
+    result = parse_json_string(json_str)
+    assert result == {"key": "value", "number": 42}
+
+
+def test_parse_json_with_code_block():
+    json_str = """```
+    {"key": "value", "number": 42}
+    ```"""
+    result = parse_json_string(json_str)
+    assert result == {"key": "value", "number": 42}
+
+
+def test_parse_json_with_language_block():
+    json_str = """```json
+    {"key": "value", "number": 42}
+    ```"""
+    result = parse_json_string(json_str)
+    assert result == {"key": "value", "number": 42}
+
+
+def test_parse_json_with_whitespace():
+    json_str = """
+        {
+            "key": "value",
+            "number": 42
+        }
+    """
+    result = parse_json_string(json_str)
+    assert result == {"key": "value", "number": 42}
+
+
+def test_parse_invalid_json():
+    json_str = '{"key": "value", invalid}'
+    with pytest.raises(ValueError) as exc_info:
+        parse_json_string(json_str)
+    assert (
+        "This task requires JSON output but the model didn't return valid JSON."
+        in str(exc_info.value)
+    )
+
+
+def test_parse_empty_code_block():
+    json_str = """```json
+    ```"""
+    with pytest.raises(ValueError) as exc_info:
+        parse_json_string(json_str)
+    assert (
+        "This task requires JSON output but the model didn't return valid JSON."
+        in str(exc_info.value)
+    )
+
+
+def test_parse_complex_json():
+    json_str = """```json
+    {
+        "string": "hello",
+        "number": 42,
+        "bool": true,
+        "null": null,
+        "array": [1, 2, 3],
+        "nested": {
+            "inner": "value"
+        }
+    }
+    ```"""
+    result = parse_json_string(json_str)
+    assert result == {
+        "string": "hello",
+        "number": 42,
+        "bool": True,
+        "null": None,
+        "array": [1, 2, 3],
+        "nested": {"inner": "value"},
+    }

kiln_ai/adapters/parsers/test_parser_registry.py

@@ -0,0 +1,32 @@
+import pytest
+
+from kiln_ai.adapters.ml_model_list import ModelParserID
+from kiln_ai.adapters.parsers.base_parser import BaseParser
+from kiln_ai.adapters.parsers.parser_registry import model_parser_from_id
+from kiln_ai.adapters.parsers.r1_parser import R1ThinkingParser
+
+
+def test_model_parser_from_id_invalid():
+    """Test that invalid parser ID raises ValueError."""
+
+    # Create a mock enum value that isn't handled
+    class MockModelParserID:
+        mock_value = "mock_value"
+
+    with pytest.raises(ValueError) as exc_info:
+        model_parser_from_id(MockModelParserID.mock_value)  # type: ignore
+
+    assert "Unhandled enum value" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    "parser_id,expected_class",
+    [
+        (None, BaseParser),
+        (ModelParserID.r1_thinking, R1ThinkingParser),
+    ],
+)
+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

kiln_ai/adapters/parsers/test_r1_parser.py

@@ -0,0 +1,144 @@
+import pytest
+
+from kiln_ai.adapters.parsers.r1_parser import R1ThinkingParser
+from kiln_ai.adapters.run_output import RunOutput
+
+
+@pytest.fixture
+def parser():
+    return R1ThinkingParser()
+
+
+def test_valid_response(parser):
+    response = RunOutput(
+        output="<think>This is thinking content</think>This is the result",
+        intermediate_outputs=None,
+    )
+    parsed = parser.parse_output(response)
+    assert parsed.intermediate_outputs["reasoning"] == "This is thinking content"
+    assert parsed.output == "This is the result"
+
+
+def test_response_with_whitespace(parser):
+    response = RunOutput(
+        output="""
+        <think>
+            This is thinking content
+        </think>
+            This is the result
+    """,
+        intermediate_outputs=None,
+    )
+    parsed = parser.parse_output(response)
+    assert (
+        parsed.intermediate_outputs["reasoning"].strip() == "This is thinking content"
+    )
+    assert parsed.output.strip() == "This is the result"
+
+
+def test_missing_start_tag(parser):
+    with pytest.raises(ValueError, match="Response must start with <think> tag"):
+        parser.parse_output(
+            RunOutput(output="Some content</think>result", intermediate_outputs=None)
+        )
+
+
+def test_missing_end_tag(parser):
+    with pytest.raises(ValueError, match="Missing thinking tags"):
+        parser.parse_output(
+            RunOutput(output="<think>Some content", intermediate_outputs=None)
+        )
+
+
+def test_multiple_start_tags(parser):
+    with pytest.raises(ValueError, match="Multiple thinking tags found"):
+        parser.parse_output(
+            RunOutput(
+                output="<think>content1<think>content2</think>result",
+                intermediate_outputs=None,
+            )
+        )
+
+
+def test_multiple_end_tags(parser):
+    with pytest.raises(ValueError, match="Multiple thinking tags found"):
+        parser.parse_output(
+            RunOutput(
+                output="<think>content</think></think>result", intermediate_outputs=None
+            )
+        )
+
+
+def test_empty_thinking_content(parser):
+    response = RunOutput(
+        output="<think></think>This is the result", intermediate_outputs=None
+    )
+    parsed = parser.parse_output(response)
+    assert parsed.intermediate_outputs == {"reasoning": ""}
+    assert parsed.output == "This is the result"
+
+
+def test_missing_result(parser):
+    with pytest.raises(ValueError, match="No content found after </think> tag"):
+        parser.parse_output(
+            RunOutput(output="<think>Some content</think>", intermediate_outputs=None)
+        )
+
+
+def test_multiline_content(parser):
+    response = RunOutput(
+        output="""<think>Line 1
+    Line 2
+    Line 3</think>Final result""",
+        intermediate_outputs=None,
+    )
+    parsed = parser.parse_output(response)
+    assert "Line 1" in parsed.intermediate_outputs["reasoning"]
+    assert "Line 2" in parsed.intermediate_outputs["reasoning"]
+    assert "Line 3" in parsed.intermediate_outputs["reasoning"]
+    assert parsed.output == "Final result"
+
+
+def test_special_characters(parser):
+    response = RunOutput(
+        output="<think>Content with: !@#$%^&*思()</think>Result with: !@#$%^&*思()",
+        intermediate_outputs=None,
+    )
+    parsed = parser.parse_output(response)
+    assert parsed.intermediate_outputs["reasoning"] == "Content with: !@#$%^&*思()"
+    assert parsed.output == "Result with: !@#$%^&*思()"
+
+
+def test_non_string_input(parser):
+    with pytest.raises(ValueError, match="Response must be a string for R1 parser"):
+        parser.parse_output(RunOutput(output={}, intermediate_outputs=None))
+
+
+def test_intermediate_outputs(parser):
+    # append to existing intermediate outputs
+    out = parser.parse_output(
+        RunOutput(
+            output="<think>Some content</think>result",
+            intermediate_outputs={"existing": "data"},
+        )
+    )
+    assert out.intermediate_outputs["reasoning"] == "Some content"
+    assert out.intermediate_outputs["existing"] == "data"
+
+    # empty dict is allowed
+    out = parser.parse_output(
+        RunOutput(
+            output="<think>Some content</think>result",
+            intermediate_outputs={},
+        )
+    )
+    assert out.intermediate_outputs["reasoning"] == "Some content"
+
+    # None is allowed
+    out = parser.parse_output(
+        RunOutput(
+            output="<think>Some content</think>result",
+            intermediate_outputs=None,
+        )
+    )
+    assert out.intermediate_outputs["reasoning"] == "Some content"

kiln_ai/adapters/prompt_builders.py

@@ -2,8 +2,8 @@ import json
 from abc import ABCMeta, abstractmethod
 from typing import Dict
 
-from kiln_ai.datamodel import Task, TaskRun
-from kiln_ai.utils.formatting import snake_case
+from kiln_ai.datamodel import PromptGenerators, PromptId, Task, TaskRun
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
 class BasePromptBuilder(metaclass=ABCMeta):
@@ -20,25 +20,38 @@ class BasePromptBuilder(metaclass=ABCMeta):
         """
         self.task = task
 
-    @abstractmethod
-    def build_prompt(self) -> str:
+    def prompt_id(self) -> str | None:
+        """Returns the ID of the prompt, scoped to this builder.
+
+        Returns:
+            str | None: The ID of the prompt, or None if not set.
+        """
+        return None
+
+    def build_prompt(self, include_json_instructions) -> str:
         """Build and return the complete prompt string.
 
         Returns:
             str: The constructed prompt.
         """
-        pass
+        prompt = self.build_base_prompt()
 
-    @classmethod
-    def prompt_builder_name(cls) -> str:
-        """Returns the name of the prompt builder, to be used for persisting into the datastore.
+        if include_json_instructions and self.task.output_schema():
+            prompt = (
+                prompt
+                + f"\n\n# Format Instructions\n\nReturn a JSON object conforming to the following schema:\n```\n{self.task.output_schema()}\n```"
+            )
+
+        return prompt
 
-        Default implementation gets the name of the prompt builder in snake case. If you change the class name, you should override this so prior saved data is compatible.
+    @abstractmethod
+    def build_base_prompt(self) -> str:
+        """Build and return the complete prompt string.
 
         Returns:
-            str: The prompt builder name in snake_case format.
+            str: The constructed prompt.
         """
-        return snake_case(cls.__name__)
+        pass
 
     def build_user_message(self, input: Dict | str) -> str:
         """Build a user message from the input.
@@ -50,7 +63,7 @@ class BasePromptBuilder(metaclass=ABCMeta):
             str: The formatted user message.
         """
         if isinstance(input, Dict):
-            return f"The input is:\n{json.dumps(input, indent=2)}"
+            return f"The input is:\n{json.dumps(input, indent=2, ensure_ascii=False)}"
 
         return f"The input is:\n{input}"
 
@@ -70,7 +83,7 @@ class BasePromptBuilder(metaclass=ABCMeta):
         Returns:
             str: The constructed prompt string.
         """
-        base_prompt = self.build_prompt()
+        base_prompt = self.build_prompt(include_json_instructions=False)
         cot_prompt = self.chain_of_thought_prompt()
         if cot_prompt:
             base_prompt += "\n# Thinking Instructions\n\n" + cot_prompt
@@ -80,7 +93,7 @@ class BasePromptBuilder(metaclass=ABCMeta):
 class SimplePromptBuilder(BasePromptBuilder):
     """A basic prompt builder that combines task instruction with requirements."""
 
-    def build_prompt(self) -> str:
+    def build_base_prompt(self) -> str:
         """Build a simple prompt with instruction and requirements.
 
         Returns:
@@ -95,7 +108,7 @@ class SimplePromptBuilder(BasePromptBuilder):
             )
             # iterate requirements, formatting them in numbereed list like 1) task.instruction\n2)...
             for i, requirement in enumerate(self.task.requirements):
-                base_prompt += f"{i+1}) {requirement.instruction}\n"
+                base_prompt += f"{i + 1}) {requirement.instruction}\n"
 
         return base_prompt
 
@@ -112,18 +125,18 @@ class MultiShotPromptBuilder(BasePromptBuilder):
         """
         return 25
 
-    def build_prompt(self) -> str:
+    def build_base_prompt(self) -> str:
         """Build a prompt with instruction, requirements, and multiple examples.
 
         Returns:
             str: The constructed prompt string with examples.
         """
-        base_prompt = f"# Instruction\n\n{ self.task.instruction }\n\n"
+        base_prompt = f"# Instruction\n\n{self.task.instruction}\n\n"
 
         if len(self.task.requirements) > 0:
             base_prompt += "# Requirements\n\nYour response should respect the following requirements:\n"
             for i, requirement in enumerate(self.task.requirements):
-                base_prompt += f"{i+1}) {requirement.instruction}\n"
+                base_prompt += f"{i + 1}) {requirement.instruction}\n"
             base_prompt += "\n"
 
         valid_examples = self.collect_examples()
@@ -140,11 +153,11 @@ class MultiShotPromptBuilder(BasePromptBuilder):
     def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
         # Prefer repaired output if it exists, otherwise use the regular output
         output = example.repaired_output or example.output
-        return f"## Example {index+1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
+        return f"## Example {index + 1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
 
     def collect_examples(self) -> list[TaskRun]:
         valid_examples: list[TaskRun] = []
-        runs = self.task.runs()
+        runs = self.task.runs(readonly=True)
 
         # first pass, we look for repaired outputs. These are the best examples.
         for run in runs:
@@ -198,7 +211,7 @@ class RepairsPromptBuilder(MultiShotPromptBuilder):
         ):
             return super().prompt_section_for_example(index, example)
 
-        prompt_section = f"## Example {index+1}\n\nInput: {example.input}\n\n"
+        prompt_section = f"## Example {index + 1}\n\nInput: {example.input}\n\n"
         prompt_section += (
             f"Initial Output Which Was Insufficient: {example.output.output}\n\n"
         )
@@ -209,7 +222,7 @@ class RepairsPromptBuilder(MultiShotPromptBuilder):
         return prompt_section
 
 
-def chain_of_thought_prompt(task: Task) -> str | None:
+def chain_of_thought_prompt(task: Task) -> str:
     """Standard implementation to build and return the chain of thought prompt string.
 
     Returns:
@@ -244,23 +257,132 @@ class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
         return chain_of_thought_prompt(self.task)
 
 
-prompt_builder_registry = {
-    "simple_prompt_builder": SimplePromptBuilder,
-    "multi_shot_prompt_builder": MultiShotPromptBuilder,
-    "few_shot_prompt_builder": FewShotPromptBuilder,
-    "repairs_prompt_builder": RepairsPromptBuilder,
-    "simple_chain_of_thought_prompt_builder": SimpleChainOfThoughtPromptBuilder,
-    "few_shot_chain_of_thought_prompt_builder": FewShotChainOfThoughtPromptBuilder,
-    "multi_shot_chain_of_thought_prompt_builder": MultiShotChainOfThoughtPromptBuilder,
-}
+class SavedPromptBuilder(BasePromptBuilder):
+    """A prompt builder that looks up a static prompt."""
+
+    def __init__(self, task: Task, prompt_id: str):
+        super().__init__(task)
+        prompt_model = next(
+            (
+                prompt
+                for prompt in task.prompts(readonly=True)
+                if prompt.id == prompt_id
+            ),
+            None,
+        )
+        if not prompt_model:
+            raise ValueError(f"Prompt ID not found: {prompt_id}")
+        self.prompt_model = prompt_model
+
+    def prompt_id(self) -> str | None:
+        return self.prompt_model.id
+
+    def build_base_prompt(self) -> str:
+        """Returns a saved prompt.
+
+        Returns:
+            str: The prompt string.
+        """
+        return self.prompt_model.prompt
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return self.prompt_model.chain_of_thought_instructions
+
+
+class TaskRunConfigPromptBuilder(BasePromptBuilder):
+    """A prompt builder that looks up a static prompt in a task run config."""
+
+    def __init__(self, task: Task, run_config_prompt_id: str):
+        parts = run_config_prompt_id.split("::")
+        if len(parts) != 4:
+            raise ValueError(
+                f"Invalid task run config prompt ID: {run_config_prompt_id}. Expected format: 'task_run_config::[project_id]::[task_id]::[run_config_id]'."
+            )
+
+        task_id = parts[2]
+        if task_id != task.id:
+            raise ValueError(
+                f"Task run config prompt ID: {run_config_prompt_id}. Task ID mismatch. Expected: {task.id}, got: {task_id}."
+            )
+
+        run_config_id = parts[3]
+        run_config = next(
+            (
+                run_config
+                for run_config in task.run_configs(readonly=True)
+                if run_config.id == run_config_id
+            ),
+            None,
+        )
+        if not run_config:
+            raise ValueError(
+                f"Task run config ID not found: {run_config_id} for prompt id {run_config_prompt_id}"
+            )
+        if run_config.prompt is None:
+            raise ValueError(
+                f"Task run config ID {run_config_id} does not have a stored prompt. Used as prompt id {run_config_prompt_id}"
+            )
+
+        # Load the prompt from the model
+        self.prompt = run_config.prompt.prompt
+        self.cot_prompt = run_config.prompt.chain_of_thought_instructions
+        self.id = run_config_prompt_id
+
+        super().__init__(task)
+
+    def prompt_id(self) -> str | None:
+        return self.id
+
+    def build_base_prompt(self) -> str:
+        return self.prompt
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return self.cot_prompt
+
+
+class FineTunePromptBuilder(BasePromptBuilder):
+    """A prompt builder that looks up a fine-tune prompt."""
+
+    def __init__(self, task: Task, nested_fine_tune_id: str):
+        super().__init__(task)
+
+        # IDs are in project_id::task_id::fine_tune_id format
+        self.full_fine_tune_id = nested_fine_tune_id
+        parts = nested_fine_tune_id.split("::")
+        if len(parts) != 3:
+            raise ValueError(
+                f"Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id', got: {nested_fine_tune_id}"
+            )
+        fine_tune_id = parts[2]
+
+        fine_tune_model = next(
+            (
+                fine_tune
+                for fine_tune in task.finetunes(readonly=True)
+                if fine_tune.id == fine_tune_id
+            ),
+            None,
+        )
+        if not fine_tune_model:
+            raise ValueError(f"Fine-tune ID not found: {fine_tune_id}")
+        self.fine_tune_model = fine_tune_model
+
+    def prompt_id(self) -> str | None:
+        return self.full_fine_tune_id
+
+    def build_base_prompt(self) -> str:
+        return self.fine_tune_model.system_message
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return self.fine_tune_model.thinking_instructions
 
 
 # Our UI has some names that are not the same as the class names, which also hint parameters.
-def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
+def prompt_builder_from_id(prompt_id: PromptId, task: Task) -> BasePromptBuilder:
     """Convert a name used in the UI to the corresponding prompt builder class.
 
     Args:
-        ui_name (str): The UI name for the prompt builder type.
+        prompt_id (PromptId): The prompt ID.
 
     Returns:
         type[BasePromptBuilder]: The corresponding prompt builder class.
@@ -268,20 +390,42 @@ def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
     Raises:
         ValueError: If the UI name is not recognized.
     """
-    match ui_name:
-        case "basic":
-            return SimplePromptBuilder
-        case "few_shot":
-            return FewShotPromptBuilder
-        case "many_shot":
-            return MultiShotPromptBuilder
-        case "repairs":
-            return RepairsPromptBuilder
-        case "simple_chain_of_thought":
-            return SimpleChainOfThoughtPromptBuilder
-        case "few_shot_chain_of_thought":
-            return FewShotChainOfThoughtPromptBuilder
-        case "multi_shot_chain_of_thought":
-            return MultiShotChainOfThoughtPromptBuilder
+
+    # Saved prompts are prefixed with "id::"
+    if prompt_id.startswith("id::"):
+        prompt_id = prompt_id[4:]
+        return SavedPromptBuilder(task, prompt_id)
+
+    # Task run config prompts are prefixed with "task_run_config::"
+    # task_run_config::[project_id]::[task_id]::[run_config_id]
+    if prompt_id.startswith("task_run_config::"):
+        return TaskRunConfigPromptBuilder(task, prompt_id)
+
+    # Fine-tune prompts are prefixed with "fine_tune_prompt::"
+    if prompt_id.startswith("fine_tune_prompt::"):
+        prompt_id = prompt_id[18:]
+        return FineTunePromptBuilder(task, prompt_id)
+
+    # Check if the prompt_id matches any enum value
+    if prompt_id not in [member.value for member in PromptGenerators]:
+        raise ValueError(f"Unknown prompt generator: {prompt_id}")
+    typed_prompt_generator = PromptGenerators(prompt_id)
+
+    match typed_prompt_generator:
+        case PromptGenerators.SIMPLE:
+            return SimplePromptBuilder(task)
+        case PromptGenerators.FEW_SHOT:
+            return FewShotPromptBuilder(task)
+        case PromptGenerators.MULTI_SHOT:
+            return MultiShotPromptBuilder(task)
+        case PromptGenerators.REPAIRS:
+            return RepairsPromptBuilder(task)
+        case PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT:
+            return SimpleChainOfThoughtPromptBuilder(task)
+        case PromptGenerators.FEW_SHOT_CHAIN_OF_THOUGHT:
+            return FewShotChainOfThoughtPromptBuilder(task)
+        case PromptGenerators.MULTI_SHOT_CHAIN_OF_THOUGHT:
+            return MultiShotChainOfThoughtPromptBuilder(task)
         case _:
-            raise ValueError(f"Unknown prompt builder: {ui_name}")
+            # Type checking will find missing cases
+            raise_exhaustive_enum_error(typed_prompt_generator)