kiln-ai 0.14.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

kiln_ai/datamodel/task_run.py

@@ -3,11 +3,11 @@ from typing import TYPE_CHECKING, Dict, List, Union
 
 import jsonschema
 import jsonschema.exceptions
-from pydantic import Field, ValidationInfo, model_validator
+from pydantic import BaseModel, Field, ValidationInfo, model_validator
 from typing_extensions import Self
 
 from kiln_ai.datamodel.basemodel import KilnParentedModel
-from kiln_ai.datamodel.json_schema import validate_schema
+from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 from kiln_ai.datamodel.strict_mode import strict_mode
 from kiln_ai.datamodel.task_output import DataSource, TaskOutput
 
@@ -15,6 +15,29 @@ if TYPE_CHECKING:
     from kiln_ai.datamodel.task import Task
 
 
+class Usage(BaseModel):
+    input_tokens: int | None = Field(
+        default=None,
+        description="The number of input tokens used in the task run.",
+        ge=0,
+    )
+    output_tokens: int | None = Field(
+        default=None,
+        description="The number of output tokens used in the task run.",
+        ge=0,
+    )
+    total_tokens: int | None = Field(
+        default=None,
+        description="The total number of tokens used in the task run.",
+        ge=0,
+    )
+    cost: float | None = Field(
+        default=None,
+        description="The cost of the task run in US dollars, saved at runtime (prices can change over time).",
+        ge=0,
+    )
+
+
 class TaskRun(KilnParentedModel):
     """
     Represents a single execution of a Task.
@@ -47,17 +70,26 @@ class TaskRun(KilnParentedModel):
         default=[],
         description="Tags for the task run. Tags are used to categorize task runs for filtering and reporting.",
     )
+    usage: Usage | None = Field(
+        default=None,
+        description="Usage information for the task run. This includes the number of input tokens, output tokens, and total tokens used.",
+    )
+
+    def thinking_training_data(self) -> str | None:
+        """
+        Get the thinking training data from the task run.
+        """
+        if self.intermediate_outputs is None:
+            return None
+        return self.intermediate_outputs.get(
+            "reasoning"
+        ) or self.intermediate_outputs.get("chain_of_thought")
 
     def has_thinking_training_data(self) -> bool:
         """
         Does this run have thinking data that we can use to train a thinking model?
         """
-        if self.intermediate_outputs is None:
-            return False
-        return (
-            "chain_of_thought" in self.intermediate_outputs
-            or "reasoning" in self.intermediate_outputs
-        )
+        return self.thinking_training_data() is not None
 
     # Workaround to return typed parent without importing Task
     def parent_task(self) -> Union["Task", None]:
@@ -87,14 +119,19 @@ class TaskRun(KilnParentedModel):
             # don't validate this relationship until we have a path or parent. Give them time to build it (but will catch it before saving)
             return self
 
-        # validate output
+        # validate input
         if task.input_json_schema is not None:
             try:
-                validate_schema(json.loads(self.input), task.input_json_schema)
+                input_parsed = json.loads(self.input)
             except json.JSONDecodeError:
                 raise ValueError("Input is not a valid JSON object")
-            except jsonschema.exceptions.ValidationError as e:
-                raise ValueError(f"Input does not match task input schema: {e}")
+
+            validate_schema_with_value_error(
+                input_parsed,
+                task.input_json_schema,
+                "Input does not match task input schema.",
+            )
+
         self._last_validated_input = self.input
         return self
 
@@ -131,6 +168,24 @@ class TaskRun(KilnParentedModel):
                 raise ValueError(
                     "Repaired output rating must be None. Repaired outputs are assumed to have a perfect rating, as they have been fixed."
                 )
+
+            task = self.parent_task()
+            if (
+                task is not None
+                and self.repaired_output.output is not None
+                and task.output_json_schema is not None
+            ):
+                try:
+                    output_parsed = json.loads(self.repaired_output.output)
+                except json.JSONDecodeError:
+                    raise ValueError("Repaired output is not a valid JSON object")
+
+                validate_schema_with_value_error(
+                    output_parsed,
+                    task.output_json_schema,
+                    "Repaired output does not match task output schema.",
+                )
+
         if self.repair_instructions is None and self.repaired_output is not None:
             raise ValueError(
                 "Repair instructions are required if providing a repaired output."
@@ -139,6 +194,7 @@ class TaskRun(KilnParentedModel):
             raise ValueError(
                 "A repaired output is required if providing repair instructions."
             )
+
         return self
 
     @model_validator(mode="after")

kiln_ai/datamodel/test_basemodel.py

@@ -483,7 +483,7 @@ class MockAdapter(BaseAdapter):
     """Implementation of BaseAdapter for testing"""
 
     async def _run(self, input):
-        return RunOutput(output="test output", intermediate_outputs=None)
+        return RunOutput(output="test output", intermediate_outputs=None), None
 
     def adapter_name(self) -> str:
         return "test"
@@ -510,6 +510,7 @@ async def test_invoke_parsing_flow(adapter):
     # Mock dependencies
     mock_provider = MagicMock()
     mock_provider.parser = "test_parser"
+    mock_provider.formatter = None
     mock_provider.reasoning_capable = False
 
     mock_parser = MagicMock()
@@ -517,13 +518,11 @@ async def test_invoke_parsing_flow(adapter):
         output="parsed test output", intermediate_outputs={"key": "value"}
     )
 
-    mock_parser_class = MagicMock(return_value=mock_parser)
-
     with (
         patch.object(adapter, "model_provider", return_value=mock_provider),
         patch(
             "kiln_ai.adapters.model_adapters.base_adapter.model_parser_from_id",
-            return_value=mock_parser_class,
+            return_value=mock_parser,
         ),
         patch("kiln_ai.adapters.model_adapters.base_adapter.Config") as mock_config,
     ):
@@ -534,9 +533,6 @@ async def test_invoke_parsing_flow(adapter):
         # Execute
         result = await adapter.invoke("test input")
 
-        # Verify parser was created correctly
-        mock_parser_class.assert_called_once_with(structured_output=False)
-
         # Verify parsing occurred
         mock_parser.parse_output.assert_called_once()
         parsed_args = mock_parser.parse_output.call_args[1]

kiln_ai/datamodel/test_dataset_filters.py

@@ -1,3 +1,5 @@
+from unittest.mock import Mock
+
 import pytest
 from pydantic import BaseModel
 
@@ -5,12 +7,14 @@ from kiln_ai.datamodel.dataset_filters import (
     AllDatasetFilter,
     DatasetFilterId,
     HighRatingDatasetFilter,
+    MultiDatasetFilter,
     StaticDatasetFilters,
     TagFilter,
     ThinkingModelDatasetFilter,
     ThinkingModelHighRatedFilter,
     dataset_filter_from_id,
 )
+from kiln_ai.datamodel.task_run import TaskRun
 
 # Note: Many more filter tests in test_dataset_split.py
 
@@ -69,3 +73,81 @@ def test_tag_filter(tag, expected_error, expected_tag):
         filter = dataset_filter_from_id(tag)
         assert isinstance(filter, TagFilter)
         assert filter.tag == expected_tag
+
+
+class TestMultiDatasetFilter:
+    @pytest.mark.parametrize(
+        "filter_string,expected_filters",
+        [
+            ("multi_filter::high_rating", ["high_rating"]),
+            (
+                "multi_filter::high_rating&thinking_model",
+                ["high_rating", "thinking_model"],
+            ),
+            ("multi_filter::tag::test&high_rating", ["tag::test", "high_rating"]),
+            (
+                "multi_filter::high_rating&tag::tag\\&name",
+                ["high_rating", "tag::tag&name"],
+            ),
+        ],
+    )
+    def test_valid_filter_string_parsing(self, filter_string, expected_filters):
+        """Test that valid filter strings are parsed correctly."""
+        assert MultiDatasetFilter.parse_filter_string(filter_string) == expected_filters
+        assert MultiDatasetFilter.is_valid_filter_string(filter_string)
+
+    @pytest.mark.parametrize(
+        "filter_string,expected_error",
+        [
+            (
+                "not_multi_filter::high_rating",
+                "Filter string must start with multi_filter::",
+            ),
+            ("multi_filter::", "No filters specified after prefix"),
+            ("multi_filter::high_rating&", "Invalid dataset filter ID:"),
+            ("multi_filter::invalid_filter", "Invalid dataset filter ID:"),
+        ],
+    )
+    def test_invalid_filter_string_handling(self, filter_string, expected_error):
+        """Test that invalid filter strings raise appropriate errors."""
+        with pytest.raises(ValueError, match=expected_error):
+            MultiDatasetFilter.parse_filter_string(filter_string)
+        assert not MultiDatasetFilter.is_valid_filter_string(filter_string)
+
+    def test_filter_combination_logic(self):
+        """Test that multiple filters are combined with AND logic."""
+        # Create a mock task run
+        task_run = Mock(spec=TaskRun)
+        task_run.output = Mock()
+        task_run.output.rating = Mock()
+        task_run.output.rating.is_high_quality.return_value = True
+        task_run.tags = ["test_tag"]
+        task_run.has_thinking_training_data.return_value = True
+        task_run.repaired_output = None
+
+        # Test combining high_rating and tag filters
+        filter_id = "multi_filter::high_rating&tag::test_tag"
+        multi_filter = dataset_filter_from_id(filter_id)
+        assert multi_filter(task_run)
+
+        # Test that it fails if one filter fails
+        task_run.tags = ["wrong_tag"]
+        assert not multi_filter(task_run)
+        task_run.tags = ["test_tag"]
+        assert multi_filter(task_run)
+        task_run.output.rating.is_high_quality.return_value = False
+        assert not multi_filter(task_run)
+
+        # Verify the mock was called as expected
+        task_run.output.rating.is_high_quality.assert_called()
+
+    def test_filter_creation_from_id(self):
+        """Test that multi filters can be created via dataset_filter_from_id."""
+        filter_id = "multi_filter::high_rating&thinking_model"
+        filter = dataset_filter_from_id(filter_id)
+        assert isinstance(filter, MultiDatasetFilter)
+        assert len(filter.filters) == 2
+        assert any(isinstance(f, type(HighRatingDatasetFilter)) for f in filter.filters)
+        assert any(
+            isinstance(f, type(ThinkingModelDatasetFilter)) for f in filter.filters
+        )

kiln_ai/datamodel/test_dataset_split.py

@@ -17,6 +17,7 @@ from kiln_ai.datamodel.dataset_split import (
     AllSplitDefinition,
     Train60Test20Val20SplitDefinition,
     Train80Test20SplitDefinition,
+    Train80Val20SplitDefinition,
 )
 from kiln_ai.datamodel.test_dataset_filters import (
     AllDatasetFilter,
@@ -174,6 +175,7 @@ def test_high_rating_dataset_filter(sample_task_runs):
     [
         (Train80Test20SplitDefinition, {"train": 8, "test": 2}),
         (AllSplitDefinition, {"all": 10}),
+        (Train80Val20SplitDefinition, {"train": 8, "val": 2}),
         (Train60Test20Val20SplitDefinition, {"train": 6, "test": 2, "val": 2}),
         (
             [

kiln_ai/datamodel/test_example_models.py

@@ -16,6 +16,7 @@ from kiln_ai.datamodel import (
     TaskOutputRatingType,
     TaskRequirement,
     TaskRun,
+    Usage,
 )
 
 
@@ -358,6 +359,9 @@ def test_task_output_schema_validation(tmp_path):
         task_output.save_to_file()
 
 
+_input_schema_match = "Input does not match task input schema"
+
+
 def test_task_input_schema_validation(tmp_path):
     # Create a project and task hierarchy
     project = Project(name="Test Project", path=(tmp_path / "test_project"))
@@ -395,18 +399,18 @@ def test_task_input_schema_validation(tmp_path):
     valid_task_output.save_to_file()
 
     # Changing to invalid input
-    with pytest.raises(ValueError, match=_schema_match):
+    with pytest.raises(ValueError, match=_input_schema_match):
         valid_task_output.input = '{"name": "John Doe", "age": "thirty"}'
         valid_task_output.save_to_file()
 
     # loading from file, then changing to invalid input
     loaded_task_output = TaskRun.load_from_file(valid_task_output.path)
-    with pytest.raises(ValueError, match=_schema_match):
+    with pytest.raises(ValueError, match=_input_schema_match):
         loaded_task_output.input = '{"name": "John Doe", "age": "thirty"}'
         loaded_task_output.save_to_file()
 
     # Invalid case: input does not match task input schema
-    with pytest.raises(ValueError, match=_schema_match):
+    with pytest.raises(ValueError, match=_input_schema_match):
         task_output = TaskRun(
             input='{"name": "John Doe", "age": "thirty"}',
             input_source=DataSource(
@@ -642,3 +646,154 @@ def test_task_run_validate_repaired_output():
         )
 
     assert "Repaired output rating must be None" in str(exc_info.value)
+
+
+def test_task_run_validate_repaired_output_structured(tmp_path):
+    # Create a project, task, and example hierarchy
+    project = Project(name="Test Project", path=(tmp_path / "test_project"))
+    project.save_to_file()
+    task = Task(
+        name="Test Task",
+        instruction="test instruction",
+        parent=project,
+        output_json_schema=json.dumps(
+            {
+                "type": "object",
+                "properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
+                "required": ["name", "age"],
+            }
+        ),
+    )
+    task.save_to_file()
+
+    # test valid repaired output schema
+    task_run = TaskRun(
+        parent=task,
+        input="test input",
+        input_source=DataSource(
+            type=DataSourceType.human,
+            properties={"created_by": "john_doe"},
+        ),
+        output=TaskOutput(
+            output='{"name": "John Doe", "age": 30}',
+            source=DataSource(
+                type=DataSourceType.human,
+                properties={"created_by": "john_doe"},
+            ),
+        ),
+        repair_instructions="Fix the output",
+        repaired_output=TaskOutput(
+            output='{"name": "John Doe", "age": 30}',
+            source=DataSource(
+                type=DataSourceType.human, properties={"created_by": "john_doe"}
+            ),
+        ),
+    )
+
+    assert task_run.repaired_output is not None
+    assert task_run.repaired_output.rating is None
+
+    # test invalid JSON
+    with pytest.raises(ValueError):
+        TaskRun(
+            parent=task,
+            input="test input",
+            input_source=DataSource(
+                type=DataSourceType.human,
+                properties={"created_by": "john_doe"},
+            ),
+            output=TaskOutput(
+                output='{"name": "John Doe", "age": 30}',
+                source=DataSource(
+                    type=DataSourceType.human,
+                    properties={"created_by": "john_doe"},
+                ),
+            ),
+            repair_instructions="Fix the output",
+            repaired_output=TaskOutput(
+                output='{"name": "John Doe", "age": 30',  # missing closing brace
+                source=DataSource(
+                    type=DataSourceType.human,
+                    properties={"created_by": "john_doe"},
+                ),
+            ),
+        )
+
+    # test invalid repaired output schema
+    with pytest.raises(ValueError):
+        TaskRun(
+            parent=task,
+            input="test input",
+            input_source=DataSource(
+                type=DataSourceType.human,
+                properties={"created_by": "john_doe"},
+            ),
+            output=TaskOutput(
+                output='{"name": "John Doe", "age": 30}',
+                source=DataSource(
+                    type=DataSourceType.human,
+                    properties={"created_by": "john_doe"},
+                ),
+            ),
+            repair_instructions="Fix the output",
+            repaired_output=TaskOutput(
+                output='{"name": "John Doe", "age": "thirty"}',  # invalid schema
+                source=DataSource(
+                    type=DataSourceType.human,
+                    properties={"created_by": "john_doe"},
+                ),
+            ),
+        )
+
+
+@pytest.mark.parametrize(
+    "input_tokens,output_tokens,total_tokens,cost,should_raise",
+    [
+        # Valid cases
+        (100, 50, 150, 0.002, False),  # All fields
+        (None, None, None, None, False),  # All None (defaults)
+        # Invalid cases
+        (-100, 50, 150, 0.002, True),  # Negative input_tokens
+        (100, -50, 150, 0.002, True),  # Negative output_tokens
+        (100, 50, -150, 0.002, True),  # Negative total_tokens
+        (100, 50, 150, -0.002, True),  # Negative cost
+    ],
+)
+def test_usage_model(input_tokens, output_tokens, total_tokens, cost, should_raise):
+    """Test the Usage model with various input combinations."""
+    if should_raise:
+        with pytest.raises(ValidationError):
+            Usage(
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                total_tokens=total_tokens,
+                cost=cost,
+            )
+    else:
+        usage = Usage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+            cost=cost,
+        )
+        assert usage.input_tokens == input_tokens
+        assert usage.output_tokens == output_tokens
+        assert usage.total_tokens == total_tokens
+        assert usage.cost == cost
+
+
+def test_usage_model_in_task_run(valid_task_run):
+    """Test that Usage can be properly set in a TaskRun."""
+    usage = Usage(
+        input_tokens=100,
+        output_tokens=50,
+        total_tokens=150,
+        cost=0.002,
+    )
+    task_run = valid_task_run.model_copy(deep=True)
+    task_run.usage = usage
+    assert task_run.usage == usage
+    assert task_run.usage.input_tokens == 100
+    assert task_run.usage.output_tokens == 50
+    assert task_run.usage.total_tokens == 150
+    assert task_run.usage.cost == 0.002

kiln_ai/datamodel/test_json_schema.py

@@ -1,3 +1,4 @@
+import jsonschema
 import pytest
 from pydantic import BaseModel
 
@@ -6,6 +7,7 @@ from kiln_ai.datamodel.json_schema import (
     schema_from_json_str,
     string_to_json_key,
     validate_schema,
+    validate_schema_with_value_error,
 )
 
 
@@ -71,15 +73,32 @@ def test_validate_schema_content():
     o = {"setup": "asdf", "punchline": "asdf", "rating": 1}
     validate_schema(o, json_joke_schema)
     o = {"setup": "asdf"}
-    with pytest.raises(Exception):
+    with pytest.raises(jsonschema.exceptions.ValidationError):
         validate_schema(0, json_joke_schema)
     o = {"setup": "asdf", "punchline": "asdf"}
     validate_schema(o, json_joke_schema)
     o = {"setup": "asdf", "punchline": "asdf", "rating": "1"}
-    with pytest.raises(Exception):
+    with pytest.raises(jsonschema.exceptions.ValidationError):
         validate_schema(o, json_joke_schema)
 
 
+def test_validate_schema_content_with_value_error():
+    o = {"setup": "asdf", "punchline": "asdf", "rating": 1}
+    validate_schema_with_value_error(o, json_joke_schema, "PREFIX")
+    o = {"setup": "asdf"}
+    with pytest.raises(
+        ValueError, match="PREFIX The error from the schema check was: "
+    ):
+        validate_schema_with_value_error(0, json_joke_schema, "PREFIX")
+    o = {"setup": "asdf", "punchline": "asdf"}
+    validate_schema_with_value_error(o, json_joke_schema, "PREFIX")
+    o = {"setup": "asdf", "punchline": "asdf", "rating": "1"}
+    with pytest.raises(
+        ValueError, match="PREFIX The error from the schema check was: "
+    ):
+        validate_schema_with_value_error(o, json_joke_schema, "PREFIX")
+
+
 json_triangle_schema = """{
   "type": "object",
   "properties": {
@@ -122,7 +141,7 @@ def test_triangle_schema():
     assert schema["properties"]["c"]["type"] == "integer"
     assert schema["required"] == ["a", "b", "c"]
     validate_schema({"a": 1, "b": 2, "c": 3}, json_triangle_schema)
-    with pytest.raises(Exception):
+    with pytest.raises(jsonschema.exceptions.ValidationError):
         validate_schema({"a": 1, "b": 2, "c": "3"}, json_triangle_schema)
 
 

kiln_ai/datamodel/test_model_perf.py

@@ -119,7 +119,8 @@ def test_benchmark_load_from_file(benchmark, task_run):
     avg_time_per_iteration = total_time / iterations
     ops_per_second = 1.0 / avg_time_per_iteration
 
-    # I get 8k ops per second on my MBP. Lower value here for CI.
+    # I get 8k ops per second on my MBP. Lower value here for CI and parallel testing.
     # Prior to optimization was 290 ops per second.
-    if ops_per_second < 1000:
+    print(f"Ops per second: {ops_per_second:.6f}")
+    if ops_per_second < 500:
         pytest.fail(f"Ops per second: {ops_per_second:.6f}, expected more than 1k ops")

kiln_ai/datamodel/test_models.py

@@ -547,20 +547,34 @@ def test_prompt_parent_task():
             False,
             None,
         ),
-        # Test 3: Invalid case - thinking instructions with final_only
+        # Test 3: Valid case - no thinking instructions with final_and_intermediate_r1_compatible
+        (
+            None,
+            FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+            False,
+            None,
+        ),
+        # Test 4: Invalid case - thinking instructions with final_only
         (
             "Think step by step",
             FinetuneDataStrategy.final_only,
             True,
             "Thinking instructions can only be used when data_strategy is final_and_intermediate",
         ),
-        # Test 4: Invalid case - no thinking instructions with final_and_intermediate
+        # Test 5: Invalid case - no thinking instructions with final_and_intermediate
         (
             None,
             FinetuneDataStrategy.final_and_intermediate,
             True,
             "Thinking instructions are required when data_strategy is final_and_intermediate",
         ),
+        # Test 6: Invalid case - thinking instructions with final_and_intermediate_r1_compatible
+        (
+            "Think step by step",
+            FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+            True,
+            "Thinking instructions can only be used when data_strategy is final_and_intermediate",
+        ),
     ],
 )
 def test_finetune_thinking_instructions_validation(
@@ -617,3 +631,37 @@ def test_task_run_has_thinking_training_data(intermediate_outputs, expected):
         intermediate_outputs=intermediate_outputs,
     )
     assert task_run.has_thinking_training_data() == expected
+
+
+@pytest.mark.parametrize(
+    "intermediate_outputs,expected",
+    [
+        # No intermediate outputs
+        (None, None),
+        # Empty intermediate outputs
+        ({}, None),
+        # Only chain_of_thought
+        ({"chain_of_thought": "thinking process"}, "thinking process"),
+        # Only reasoning
+        ({"reasoning": "reasoning process"}, "reasoning process"),
+        # Both chain_of_thought and reasoning (should return reasoning as it's checked first)
+        (
+            {"chain_of_thought": "thinking process", "reasoning": "reasoning process"},
+            "reasoning process",
+        ),
+        # Other intermediate outputs but no thinking data
+        ({"other_output": "some data"}, None),
+        # Mixed other outputs with thinking data
+        (
+            {"chain_of_thought": "thinking process", "other_output": "some data"},
+            "thinking process",
+        ),
+    ],
+)
+def test_task_run_thinking_training_data(intermediate_outputs, expected):
+    task_run = TaskRun(
+        input="test input",
+        output=TaskOutput(output="test output"),
+        intermediate_outputs=intermediate_outputs,
+    )
+    assert task_run.thinking_training_data() == expected