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

kiln_ai/adapters/repair/test_repair_task.py

@@ -218,8 +218,9 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     }
 
     with patch.object(LiteLlmAdapter, "_run", new_callable=AsyncMock) as mock_run:
-        mock_run.return_value = RunOutput(
-            output=mocked_output, intermediate_outputs=None
+        mock_run.return_value = (
+            RunOutput(output=mocked_output, intermediate_outputs=None),
+            None,
         )
 
         adapter = adapter_for_task(

kiln_ai/adapters/test_prompt_builders.py

@@ -3,7 +3,7 @@ import logging
 
 import pytest
 
-from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter
+from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput
 from kiln_ai.adapters.model_adapters.test_structured_output import (
     build_structured_output_test_task,
 )
@@ -15,6 +15,7 @@ from kiln_ai.adapters.prompt_builders import (
     MultiShotPromptBuilder,
     RepairsPromptBuilder,
     SavedPromptBuilder,
+    ShortPromptBuilder,
     SimpleChainOfThoughtPromptBuilder,
     SimplePromptBuilder,
     TaskRunConfigPromptBuilder,
@@ -33,6 +34,7 @@ from kiln_ai.datamodel import (
     TaskOutput,
     TaskOutputRating,
     TaskRun,
+    Usage,
 )
 from kiln_ai.datamodel.task import RunConfigProperties, TaskRunConfig
 
@@ -58,9 +60,28 @@ def test_simple_prompt_builder(tmp_path):
     assert input not in prompt
 
 
+def test_short_prompt_builder(tmp_path):
+    task = build_test_task(tmp_path)
+    builder = ShortPromptBuilder(task=task)
+    prompt = builder.build_prompt(include_json_instructions=False)
+
+    # Should only include the instruction, not requirements
+    assert task.instruction == prompt
+    assert task.requirements[0].instruction not in prompt
+    assert task.requirements[1].instruction not in prompt
+    assert task.requirements[2].instruction not in prompt
+
+    # Should handle JSON instructions correctly
+    prompt_with_json = builder.build_prompt(include_json_instructions=True)
+    assert task.instruction in prompt_with_json
+    if task.output_schema():
+        assert "# Format Instructions" in prompt_with_json
+        assert task.output_schema() in prompt_with_json
+
+
 class MockAdapter(BaseAdapter):
-    def _run(self, input: str) -> str:
-        return "mock response"
+    async def _run(self, input: str) -> tuple[RunOutput, Usage | None]:
+        return RunOutput(output="mock response", intermediate_outputs=None), None
 
     def adapter_name(self) -> str:
         return "mock_adapter"

kiln_ai/adapters/test_provider_tools.py

@@ -5,6 +5,7 @@ import pytest
 from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     ModelName,
+    ModelParserID,
     ModelProviderName,
 )
 from kiln_ai.adapters.ollama_tools import OllamaConnection
@@ -24,7 +25,12 @@ from kiln_ai.adapters.provider_tools import (
     provider_name_from_id,
     provider_warnings,
 )
-from kiln_ai.datamodel import Finetune, StructuredOutputMode, Task
+from kiln_ai.datamodel import (
+    Finetune,
+    FinetuneDataStrategy,
+    StructuredOutputMode,
+    Task,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -65,6 +71,33 @@ def mock_finetune():
         finetune.provider = ModelProviderName.openai
         finetune.fine_tune_model_id = "ft:gpt-3.5-turbo:custom:model-123"
         finetune.structured_output_mode = StructuredOutputMode.json_schema
+        finetune.data_strategy = FinetuneDataStrategy.final_only
+        mock.return_value = finetune
+        yield mock
+
+
+@pytest.fixture
+def mock_finetune_final_and_intermediate():
+    with patch("kiln_ai.datamodel.Finetune.from_id_and_parent_path") as mock:
+        finetune = Mock(spec=Finetune)
+        finetune.provider = ModelProviderName.openai
+        finetune.fine_tune_model_id = "ft:gpt-3.5-turbo:custom:model-123"
+        finetune.structured_output_mode = StructuredOutputMode.json_schema
+        finetune.data_strategy = FinetuneDataStrategy.final_and_intermediate
+        mock.return_value = finetune
+        yield mock
+
+
+@pytest.fixture
+def mock_finetune_r1_compatible():
+    with patch("kiln_ai.datamodel.Finetune.from_id_and_parent_path") as mock:
+        finetune = Mock(spec=Finetune)
+        finetune.provider = ModelProviderName.ollama
+        finetune.fine_tune_model_id = "ft:deepseek-r1:671b:custom:model-123"
+        finetune.structured_output_mode = StructuredOutputMode.json_schema
+        finetune.data_strategy = (
+            FinetuneDataStrategy.final_and_intermediate_r1_compatible
+        )
         mock.return_value = finetune
         yield mock
 
@@ -426,6 +459,38 @@ def test_finetune_provider_model_success(mock_project, mock_task, mock_finetune)
     assert provider.name == ModelProviderName.openai
     assert provider.model_id == "ft:gpt-3.5-turbo:custom:model-123"
     assert provider.structured_output_mode == StructuredOutputMode.json_schema
+    assert provider.reasoning_capable is False
+    assert provider.parser == None
+
+
+def test_finetune_provider_model_success_final_and_intermediate(
+    mock_project, mock_task, mock_finetune_final_and_intermediate
+):
+    """Test successful creation of a fine-tuned model provider"""
+    model_id = "project-123::task-456::finetune-789"
+
+    provider = finetune_provider_model(model_id)
+
+    assert provider.name == ModelProviderName.openai
+    assert provider.model_id == "ft:gpt-3.5-turbo:custom:model-123"
+    assert provider.structured_output_mode == StructuredOutputMode.json_schema
+    assert provider.reasoning_capable is True
+    assert provider.parser == None
+
+
+def test_finetune_provider_model_success_r1_compatible(
+    mock_project, mock_task, mock_finetune_r1_compatible
+):
+    """Test successful creation of a fine-tuned model provider"""
+    model_id = "project-123::task-456::finetune-789"
+
+    provider = finetune_provider_model(model_id)
+
+    assert provider.name == ModelProviderName.ollama
+    assert provider.model_id == "ft:deepseek-r1:671b:custom:model-123"
+    assert provider.structured_output_mode == StructuredOutputMode.json_schema
+    assert provider.reasoning_capable is True
+    assert provider.parser == ModelParserID.r1_thinking
 
 
 def test_finetune_provider_model_invalid_id():
@@ -515,6 +580,7 @@ def test_finetune_provider_model_structured_mode(
     finetune.provider = provider_name
     finetune.fine_tune_model_id = "fireworks-model-123"
     finetune.structured_output_mode = structured_output_mode
+    finetune.data_strategy = FinetuneDataStrategy.final_only
     mock_finetune.return_value = finetune
 
     provider = finetune_provider_model("project-123::task-456::finetune-789")
@@ -522,6 +588,8 @@ def test_finetune_provider_model_structured_mode(
     assert provider.name == provider_name
     assert provider.model_id == "fireworks-model-123"
     assert provider.structured_output_mode == expected_mode
+    assert provider.reasoning_capable is False
+    assert provider.parser == None
 
 
 def test_openai_compatible_provider_config(mock_shared_config):
@@ -799,6 +867,7 @@ def test_finetune_provider_model_vertex_ai(mock_project, mock_task, mock_finetun
     finetune.provider = ModelProviderName.vertex
     finetune.fine_tune_model_id = "projects/123/locations/us-central1/endpoints/456"
     finetune.structured_output_mode = StructuredOutputMode.json_mode
+    finetune.data_strategy = FinetuneDataStrategy.final_only
     mock_finetune.return_value = finetune
 
     provider = finetune_provider_model("project-123::task-456::finetune-789")

kiln_ai/datamodel/__init__.py

@@ -44,6 +44,7 @@ from kiln_ai.datamodel.task_output import (
 )
 from kiln_ai.datamodel.task_run import (
     TaskRun,
+    Usage,
 )
 
 __all__ = [
@@ -74,4 +75,5 @@ __all__ = [
     "PromptId",
     "PromptGenerators",
     "prompt_generator_values",
+    "Usage",
 ]

kiln_ai/datamodel/datamodel_enums.py

@@ -56,5 +56,19 @@ class FineTuneStatusType(str, Enum):
 
 
 class FinetuneDataStrategy(str, Enum):
+    """Strategy for what data to include when fine-tuning a model."""
+
+    # Only train on the final response, ignoring any intermediate steps or chain of thought
     final_only = "final_only"
+
+    # Train on both the final response and any intermediate steps/chain of thought
     final_and_intermediate = "final_and_intermediate"
+
+    # Train using R1-style thinking format, which includes the reasoning in <think> tags in the message
+    final_and_intermediate_r1_compatible = "final_and_intermediate_r1_compatible"
+
+
+THINKING_DATA_STRATEGIES: list[FinetuneDataStrategy] = [
+    FinetuneDataStrategy.final_and_intermediate,
+    FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+]

kiln_ai/datamodel/dataset_filters.py

@@ -1,5 +1,6 @@
+import re
 from enum import Enum
-from typing import Annotated, Protocol
+from typing import Annotated, ClassVar, List, Protocol
 
 from pydantic import AfterValidator
 
@@ -59,6 +60,65 @@ class TagFilter:
         return self.tag in task_run.tags
 
 
+class MultiDatasetFilter:
+    """
+    A filter that combines multiple filters using AND logic.
+    The filters are specified in a query string format after 'multi_filter::'
+    Example: multi_filter::high_rating&thinking_model&tag::tag_name
+
+    Ampersands in filter IDs can be escaped with a backslash.
+    """
+
+    PREFIX: ClassVar[str] = "multi_filter::"
+    ESCAPED_AMPERSAND: ClassVar[str] = r"\&"
+    UNESCAPED_AMPERSAND: ClassVar[str] = "&"
+
+    @classmethod
+    def parse_filter_string(cls, filter_string: str) -> List[str]:
+        """
+        Parse a filter string into individual filter IDs, handling escaped ampersands.
+        """
+        if not filter_string.startswith(cls.PREFIX):
+            raise ValueError(f"Filter string must start with {cls.PREFIX}")
+
+        # Remove the prefix
+        content = filter_string[len(cls.PREFIX) :]
+        if not content:
+            raise ValueError("No filters specified after prefix")
+
+        # Split on unescaped ampersands
+        # This regex matches & that are not preceded by a backslash
+        parts = re.split(r"(?<!\\)&", content)
+
+        # Unescape ampersands in each part
+        filter_ids = [
+            part.replace(cls.ESCAPED_AMPERSAND, cls.UNESCAPED_AMPERSAND)
+            for part in parts
+        ]
+
+        # Validate each filter ID using the existing validation
+        for fid in filter_ids:
+            _check_dataset_filter_id(fid)
+
+        return filter_ids
+
+    @classmethod
+    def is_valid_filter_string(cls, filter_string: str) -> bool:
+        """Check if a filter string is valid."""
+        try:
+            cls.parse_filter_string(filter_string)
+            return True
+        except ValueError:
+            return False
+
+    def __init__(self, filter_id: str):
+        filter_ids = MultiDatasetFilter.parse_filter_string(filter_id)
+        self.filters = [dataset_filter_from_id(fid) for fid in filter_ids]
+
+    def __call__(self, task_run: TaskRun) -> bool:
+        return all(f(task_run) for f in self.filters)
+
+
 class StaticDatasetFilters(str, Enum):
     """Dataset filter names."""
 
@@ -98,6 +158,11 @@ def _check_dataset_filter_id(id: str) -> str:
     if id.startswith("tag::") and len(id) > 5:
         return id
 
+    if id.startswith(MultiDatasetFilter.PREFIX):
+        if not MultiDatasetFilter.is_valid_filter_string(id):
+            raise ValueError(f"Invalid multi-filter string: {id}")
+        return id
+
     raise ValueError(f"Invalid dataset filter ID: {id}")
 
 
@@ -108,6 +173,9 @@ def dataset_filter_from_id(id: DatasetFilterId) -> DatasetFilter:
     if id.startswith("tag::") and len(id) > 5:
         return TagFilter(id[5:])
 
+    if id.startswith(MultiDatasetFilter.PREFIX):
+        return MultiDatasetFilter(id)
+
     if id in static_dataset_filters:
         return static_dataset_filters[id]
 

kiln_ai/datamodel/dataset_split.py

@@ -45,6 +45,10 @@ Train80Test20SplitDefinition: list[DatasetSplitDefinition] = [
     DatasetSplitDefinition(name="train", percentage=0.8),
     DatasetSplitDefinition(name="test", percentage=0.2),
 ]
+Train80Val20SplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="train", percentage=0.8),
+    DatasetSplitDefinition(name="val", percentage=0.2),
+]
 Train60Test20Val20SplitDefinition: list[DatasetSplitDefinition] = [
     DatasetSplitDefinition(name="train", percentage=0.6),
     DatasetSplitDefinition(name="test", percentage=0.2),

kiln_ai/datamodel/eval.py

@@ -263,6 +263,10 @@ class Eval(KilnParentedModel, KilnParentModel, parent_of={"configs": EvalConfig}
         default=None,
         description="The id of the current config to use for this eval. This can be changed over time to run the same eval with different configs.",
     )
+    current_run_config_id: ID_TYPE = Field(
+        default=None,
+        description="The id of the a run config which was selected as the best run config for this eval. The run config must belong to the parent Task.",
+    )
     eval_set_filter_id: DatasetFilterId = Field(
         description="The id of the dataset filter which defines which dataset items are included when running this eval. Should be mutually exclusive with eval_configs_filter_id."
     )
@@ -272,6 +276,10 @@ class Eval(KilnParentedModel, KilnParentModel, parent_of={"configs": EvalConfig}
     output_scores: List[EvalOutputScore] = Field(
         description="The scores this evaluator should produce."
     )
+    favourite: bool = Field(
+        default=False,
+        description="Whether this eval is a favourite of the user. Rendered as a star icon in the UI.",
+    )
 
     # Workaround to return typed parent without importing Task
     def parent_task(self) -> Union["Task", None]:

kiln_ai/datamodel/finetune.py

@@ -5,6 +5,7 @@ from typing_extensions import Self
 
 from kiln_ai.datamodel.basemodel import NAME_FIELD, KilnParentedModel
 from kiln_ai.datamodel.datamodel_enums import (
+    THINKING_DATA_STRATEGIES,
     FinetuneDataStrategy,
     FineTuneStatusType,
     StructuredOutputMode,

kiln_ai/datamodel/prompt_id.py

@@ -13,6 +13,7 @@ class PromptGenerators(str, Enum):
     SIMPLE_CHAIN_OF_THOUGHT = "simple_chain_of_thought_prompt_builder"
     FEW_SHOT_CHAIN_OF_THOUGHT = "few_shot_chain_of_thought_prompt_builder"
     MULTI_SHOT_CHAIN_OF_THOUGHT = "multi_shot_chain_of_thought_prompt_builder"
+    SHORT = "short_prompt_builder"
 
 
 prompt_generator_values = [pg.value for pg in PromptGenerators]

kiln_ai/datamodel/task_output.py

@@ -64,7 +64,7 @@ class TaskOutputRating(KilnBaseModel):
     )
     requirement_ratings: Dict[ID_TYPE, RequirementRating] = Field(
         default={},
-        description="The ratings of the requirements of the task.",
+        description="The ratings of the requirements of the task. The ID can be either a task_requirement_id or a named rating for an eval_output_score name (in format 'named::<name>').",
     )
 
     # Previously we stored rating values as a dict of floats, but now we store them as RequirementRating objects.

kiln_ai/datamodel/task_run.py

@@ -3,7 +3,7 @@ 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
@@ -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]:

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,
 )
 
 
@@ -743,3 +744,56 @@ def test_task_run_validate_repaired_output_structured(tmp_path):
                 ),
             ),
         )
+
+
+@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