kiln-ai 0.18.0__py3-none-any.whl → 0.20.1__py3-none-any.whl

kiln_ai/datamodel/run_config.py

@@ -0,0 +1,62 @@
+from typing import List
+
+from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
+
+from kiln_ai.datamodel.datamodel_enums import (
+    ModelProviderName,
+    StructuredOutputMode,
+)
+from kiln_ai.datamodel.prompt_id import PromptId
+from kiln_ai.datamodel.tool_id import ToolId
+
+
+class ToolsRunConfig(BaseModel):
+    """
+    A config describing which tools are available to a task.
+    """
+
+    tools: List[ToolId] = Field(
+        description="The IDs of the tools available to the task."
+    )
+
+
+class RunConfigProperties(BaseModel):
+    """
+    A configuration for running a task.
+
+    This includes everything needed to run a task, except the input and task ID. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
+    """
+
+    model_name: str = Field(description="The model to use for this run config.")
+    model_provider_name: ModelProviderName = Field(
+        description="The provider to use for this run config."
+    )
+    prompt_id: PromptId = Field(
+        description="The prompt to use for this run config. Defaults to building a simple prompt from the task if not provided.",
+    )
+    top_p: float = Field(
+        default=1.0,
+        description="The top-p value to use for this run config. Defaults to 1.0.",
+    )
+    temperature: float = Field(
+        default=1.0,
+        description="The temperature to use for this run config. Defaults to 1.0.",
+    )
+    structured_output_mode: StructuredOutputMode = Field(
+        description="The structured output mode to use for this run config.",
+    )
+    tools_config: ToolsRunConfig | None = Field(
+        default=None,
+        description="The tools config to use for this run config, defining which tools are available to the model.",
+    )
+
+    @model_validator(mode="after")
+    def validate_required_fields(self) -> Self:
+        if not (0 <= self.top_p <= 1):
+            raise ValueError("top_p must be between 0 and 1")
+
+        elif self.temperature < 0 or self.temperature > 2:
+            raise ValueError("temperature must be between 0 and 2")
+
+        return self

kiln_ai/datamodel/task.py

@@ -1,28 +1,26 @@
 from typing import TYPE_CHECKING, Dict, List, Union
 
 from pydantic import BaseModel, Field, ValidationInfo, model_validator
-from typing_extensions import Self
 
-from kiln_ai.datamodel import Finetune
 from kiln_ai.datamodel.basemodel import (
     ID_FIELD,
     ID_TYPE,
-    NAME_FIELD,
-    SHORT_NAME_FIELD,
+    FilenameString,
+    FilenameStringShort,
     KilnParentedModel,
     KilnParentModel,
 )
 from kiln_ai.datamodel.datamodel_enums import (
-    ModelProviderName,
     Priority,
     StructuredOutputMode,
     TaskOutputRatingType,
 )
 from kiln_ai.datamodel.dataset_split import DatasetSplit
 from kiln_ai.datamodel.eval import Eval
+from kiln_ai.datamodel.finetune import Finetune
 from kiln_ai.datamodel.json_schema import JsonObjectSchema, schema_from_json_str
 from kiln_ai.datamodel.prompt import BasePrompt, Prompt
-from kiln_ai.datamodel.prompt_id import PromptId
+from kiln_ai.datamodel.run_config import RunConfigProperties
 from kiln_ai.datamodel.task_run import TaskRun
 
 if TYPE_CHECKING:
@@ -38,62 +36,13 @@ class TaskRequirement(BaseModel):
     """
 
     id: ID_TYPE = ID_FIELD
-    name: str = SHORT_NAME_FIELD
+    name: FilenameStringShort = Field(description="The name of the task requirement.")
     description: str | None = Field(default=None)
     instruction: str = Field(min_length=1)
     priority: Priority = Field(default=Priority.p2)
     type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
 
 
-class RunConfigProperties(BaseModel):
-    """
-    A configuration for running a task.
-
-    This includes everything needed to run a task, except the input and task ID. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
-    """
-
-    model_name: str = Field(description="The model to use for this run config.")
-    model_provider_name: ModelProviderName = Field(
-        description="The provider to use for this run config."
-    )
-    prompt_id: PromptId = Field(
-        description="The prompt to use for this run config. Defaults to building a simple prompt from the task if not provided.",
-    )
-    top_p: float = Field(
-        default=1.0,
-        description="The top-p value to use for this run config. Defaults to 1.0.",
-    )
-    temperature: float = Field(
-        default=1.0,
-        description="The temperature to use for this run config. Defaults to 1.0.",
-    )
-    structured_output_mode: StructuredOutputMode = Field(
-        description="The structured output mode to use for this run config.",
-    )
-
-    @model_validator(mode="after")
-    def validate_required_fields(self) -> Self:
-        if not (0 <= self.top_p <= 1):
-            raise ValueError("top_p must be between 0 and 1")
-
-        elif self.temperature < 0 or self.temperature > 2:
-            raise ValueError("temperature must be between 0 and 2")
-
-        return self
-
-
-class RunConfig(RunConfigProperties):
-    """
-    A configuration for running a task.
-
-    This includes everything needed to run a task, except the input. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
-
-    For example: task, model, provider, prompt, etc.
-    """
-
-    task: "Task" = Field(description="The task to run.")
-
-
 class TaskRunConfig(KilnParentedModel):
     """
     A Kiln model for persisting a run config in a Kiln Project, nested under a task.
@@ -103,7 +52,7 @@ class TaskRunConfig(KilnParentedModel):
     A run config includes everything needed to run a task, except the input. Running the same RunConfig with the same input should make identical calls to the model (output may vary as models are non-deterministic).
     """
 
-    name: str = NAME_FIELD
+    name: FilenameString = Field(description="The name of the task run config.")
     description: str | None = Field(
         default=None, description="The description of the task run config."
     )
@@ -124,15 +73,6 @@ class TaskRunConfig(KilnParentedModel):
             return None
         return self.parent  # type: ignore
 
-    def run_config(self) -> RunConfig:
-        parent_task = self.parent_task()
-        if parent_task is None:
-            raise ValueError("Run config must be parented to a task")
-        return run_config_from_run_config_properties(
-            task=parent_task,
-            run_config_properties=self.run_config_properties,
-        )
-
     # Previously we didn't store structured_output_mode in the run_config_properties. Updgrade old models when loading from file.
     @model_validator(mode="before")
     def upgrade_old_entries(cls, data: dict, info: ValidationInfo) -> dict:
@@ -155,21 +95,6 @@ class TaskRunConfig(KilnParentedModel):
         return data
 
 
-def run_config_from_run_config_properties(
-    task: "Task",
-    run_config_properties: RunConfigProperties,
-) -> RunConfig:
-    return RunConfig(
-        task=task,
-        model_name=run_config_properties.model_name,
-        model_provider_name=run_config_properties.model_provider_name,
-        prompt_id=run_config_properties.prompt_id,
-        top_p=run_config_properties.top_p,
-        temperature=run_config_properties.temperature,
-        structured_output_mode=run_config_properties.structured_output_mode,
-    )
-
-
 class Task(
     KilnParentedModel,
     KilnParentModel,
@@ -189,7 +114,7 @@ class Task(
     a collection of task runs.
     """
 
-    name: str = NAME_FIELD
+    name: FilenameString = Field(description="The name of the task.")
     description: str | None = Field(
         default=None,
         description="A description of the task for you and your team. Will not be used in prompts/training/validation.",
@@ -216,7 +141,7 @@ class Task(
             return None
         return schema_from_json_str(self.input_json_schema)
 
-    # These wrappers help for typechecking. TODO P2: fix this in KilnParentModel
+    # These wrappers help for typechecking. We should fix this in KilnParentModel
     def runs(self, readonly: bool = False) -> list[TaskRun]:
         return super().runs(readonly=readonly)  # type: ignore
 

kiln_ai/datamodel/task_output.py

@@ -1,6 +1,6 @@
 import json
 from enum import Enum
-from typing import TYPE_CHECKING, Dict, List, Type, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel, Field, ValidationInfo, model_validator
 from typing_extensions import Self
@@ -8,6 +8,7 @@ from typing_extensions import Self
 from kiln_ai.datamodel.basemodel import ID_TYPE, KilnBaseModel
 from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
 from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
+from kiln_ai.datamodel.run_config import RunConfigProperties
 from kiln_ai.datamodel.strict_mode import strict_mode
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
@@ -199,6 +200,10 @@ class DataSource(BaseModel):
         default={},
         description="Properties describing the data source. For synthetic things like model. For human, the human's name.",
     )
+    run_config: Optional[RunConfigProperties] = Field(
+        default=None,
+        description="The run config used to generate the data, if generated by a running a model in Kiln (only true for type=synthetic).",
+    )
 
     _data_source_properties = [
         DataSourceProperty(
@@ -307,7 +312,7 @@ class TaskOutput(KilnBaseModel):
         if task.output_json_schema is not None:
             try:
                 output_parsed = json.loads(self.output)
-            except json.JSONDecodeError as e:
+            except json.JSONDecodeError:
                 raise ValueError("Output is not a valid JSON object")
 
             validate_schema_with_value_error(

kiln_ai/datamodel/task_run.py

@@ -8,6 +8,7 @@ from kiln_ai.datamodel.basemodel import KilnParentedModel
 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
+from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam
 
 if TYPE_CHECKING:
     from kiln_ai.datamodel.task import Task
@@ -35,6 +36,42 @@ class Usage(BaseModel):
         ge=0,
     )
 
+    def __add__(self, other: "Usage") -> "Usage":
+        """Add two Usage objects together, handling None values gracefully.
+
+        None + None = None
+        None + value = value
+        value + None = value
+        value1 + value2 = value1 + value2
+        """
+        if not isinstance(other, Usage):
+            raise TypeError(f"Cannot add Usage with {type(other).__name__}")
+
+        def _add_optional_int(a: int | None, b: int | None) -> int | None:
+            if a is None and b is None:
+                return None
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return a + b
+
+        def _add_optional_float(a: float | None, b: float | None) -> float | None:
+            if a is None and b is None:
+                return None
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return a + b
+
+        return Usage(
+            input_tokens=_add_optional_int(self.input_tokens, other.input_tokens),
+            output_tokens=_add_optional_int(self.output_tokens, other.output_tokens),
+            total_tokens=_add_optional_int(self.total_tokens, other.total_tokens),
+            cost=_add_optional_float(self.cost, other.cost),
+        )
+
 
 class TaskRun(KilnParentedModel):
     """
@@ -72,6 +109,10 @@ class TaskRun(KilnParentedModel):
         default=None,
         description="Usage information for the task run. This includes the number of input tokens, output tokens, and total tokens used.",
     )
+    trace: list[ChatCompletionMessageParam] | None = Field(
+        default=None,
+        description="The trace of the task run in OpenAI format. This is the list of messages that were sent to/from the model.",
+    )
 
     def thinking_training_data(self) -> str | None:
         """

kiln_ai/datamodel/test_basemodel.py

@@ -1,5 +1,6 @@
 import datetime
 import json
+import uuid
 from pathlib import Path
 from typing import Optional
 from unittest.mock import MagicMock, patch
@@ -12,10 +13,11 @@ from kiln_ai.datamodel import Task, TaskRun
 from kiln_ai.datamodel.basemodel import (
     KilnBaseModel,
     KilnParentedModel,
+    name_validator,
     string_to_valid_name,
 )
 from kiln_ai.datamodel.model_cache import ModelCache
-from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 @pytest.fixture
@@ -328,28 +330,81 @@ def test_delete_no_path():
         model.delete()
 
 
-def test_string_to_valid_name():
-    # Test basic valid strings remain unchanged
-    assert string_to_valid_name("Hello World") == "Hello World"
-    assert string_to_valid_name("Test-123") == "Test-123"
-    assert string_to_valid_name("my_file_name") == "my_file_name"
+@pytest.mark.parametrize(
+    "name,expected",
+    [
+        # Basic valid strings remain unchanged
+        ("Hello World", "Hello World"),
+        ("Test-123", "Test-123"),
+        ("my_file_name", "my_file_name"),
+        ("multiple!!!symbols", "multiple!!!symbols"),
+        # Emoji
+        ("Hello 👍", "Hello 👍"),
+        # Invalid characters are replaced
+        ("Hello@World!", "Hello@World!"),
+        ("File.name.txt", "File_name_txt"),
+        ("Special%%%Chars", "Special_Chars"),
+        ("Special#$%Chars", "Special#$_Chars"),
+        # Consecutive invalid characters are replaced
+        ("Special%%%Chars", "Special_Chars"),
+        ("path/to/file", "path_to_file"),
+        # Leading/trailing special characters are removed
+        ("__test__", "test"),
+        ("...test...", "test"),
+        # Whitespace is replaced
+        ("", ""),
+        ("   ", ""),
+        ("Hello   World", "Hello World"),
+        # Unicode characters are replaced
+        ("你好", "你好"),
+        ("你好_世界", "你好_世界"),
+        ("你好_世界_你好", "你好_世界_你好"),
+        # Newlines, tabs, and other control characters are replaced
+        ("Hello\nworld", "Hello_world"),
+        ("Hello\tworld", "Hello_world"),
+        ("Hello\rworld", "Hello_world"),
+        ("Hello\fworld", "Hello_world"),
+        ("Hello\bworld", "Hello_world"),
+        ("Hello\vworld", "Hello_world"),
+        ("Hello\0world", "Hello_world"),
+        ("Hello\x00world", "Hello_world"),
+    ],
+)
+def test_string_to_valid_name(tmp_path, name, expected):
+    assert string_to_valid_name(name) == expected
 
-    # Test invalid characters are replaced
-    assert string_to_valid_name("Hello@World!") == "Hello_World"
-    assert string_to_valid_name("File.name.txt") == "File_name_txt"
-    assert string_to_valid_name("Special#$%Chars") == "Special_Chars"
+    # check we can create a folder with the valid name
+    dir_path = tmp_path / str(uuid.uuid4()) / expected
+    dir_path.mkdir(parents=True)
 
-    # Test consecutive invalid characters
-    assert string_to_valid_name("multiple!!!symbols") == "multiple_symbols"
-    assert string_to_valid_name("path/to/file") == "path_to_file"
 
-    # Test leading/trailing special characters
-    assert string_to_valid_name("__test__") == "test"
-    assert string_to_valid_name("...test...") == "test"
+@pytest.mark.parametrize(
+    "name,min_length,max_length,should_pass",
+    [
+        # Valid cases
+        ("ValidName", 5, 20, True),
+        ("Short", 1, 10, True),
+        ("LongerValidName", 5, 20, True),
+        # None case (line 53)
+        (None, 5, 20, False),
+        # Too short cases (lines 57-59)
+        ("Hi", 5, 20, False),
+        ("", 1, 20, False),
+        ("a", 2, 20, False),
+        # Too long cases (lines 61-63)
+        ("ThisNameIsTooLong", 5, 10, False),
+        ("VeryVeryVeryLongName", 1, 15, False),
+    ],
+)
+def test_name_validator_error_conditions(name, min_length, max_length, should_pass):
+    validator = name_validator(min_length=min_length, max_length=max_length)
 
-    # Test empty string and whitespace
-    assert string_to_valid_name("") == ""
-    assert string_to_valid_name("   ") == ""
+    if should_pass:
+        result = validator(name)
+        assert result == name
+    else:
+        with pytest.raises(ValueError):
+            validator(name)
 
 
 def test_load_from_file_with_cache(test_base_file, tmp_model_cache):
@@ -497,8 +552,8 @@ def base_task():
 @pytest.fixture
 def adapter(base_task):
     return MockAdapter(
-        run_config=RunConfig(
-            task=base_task,
+        task=base_task,
+        run_config=RunConfigProperties(
             model_name="test_model",
             model_provider_name="openai",
             prompt_id="simple_prompt_builder",
@@ -553,3 +608,140 @@ async def test_invoke_parsing_flow(adapter):
             match="Reasoning is required for this model, but no reasoning was returned.",
         ):
             await adapter.invoke("test input")
+
+
+async def test_invoke_parsing_flow_basic_no_reasoning(adapter):
+    """Test for reasoning_optional_for_structured_output
+    when reasoning is not required.
+    This is a special case where we want to return the output as is.
+    """
+    # Mock dependencies
+    mock_provider = MagicMock()
+    mock_provider.parser = "test_parser"
+    mock_provider.formatter = None
+    mock_provider.reasoning_capable = False
+    mock_provider.reasoning_optional_for_structured_output = True
+
+    mock_parser = MagicMock()
+    mock_parser.parse_output.return_value = RunOutput(
+        output="parsed test output", intermediate_outputs={"key": "value"}
+    )
+
+    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,
+        ),
+        patch("kiln_ai.adapters.model_adapters.base_adapter.Config") as mock_config,
+    ):
+        # Disable autosaving for this test
+        mock_config.shared.return_value.autosave_runs = False
+        mock_config.shared.return_value.user_id = "test_user_id"
+
+        # Execute
+        result = await adapter.invoke("test input")
+
+        # Verify parsing occurred
+        mock_parser.parse_output.assert_called_once()
+        parsed_args = mock_parser.parse_output.call_args[1]
+        assert isinstance(parsed_args["original_output"], RunOutput)
+        assert parsed_args["original_output"].output == "test output"
+
+        # Verify result contains parsed output
+        assert isinstance(result, TaskRun)
+        assert result.output.output == "parsed test output"
+        assert result.intermediate_outputs == {"key": "value"}
+        assert result.input == "test input"
+
+
+async def test_invoke_parsing_flow_no_reasoning_with_structured_output(adapter):
+    """Test for reasoning_optional_for_structured_output
+    when reasoning is required but not provided, with structured output enabled.
+    This is a special case where we don't want to error, but we want to return the output as is.
+    """
+    # Mock dependencies
+    mock_provider = MagicMock()
+    mock_provider.parser = "test_parser"
+    mock_provider.formatter = None
+    mock_provider.reasoning_capable = True
+    mock_provider.reasoning_optional_for_structured_output = True
+
+    mock_parser = MagicMock()
+    mock_parser.parse_output.return_value = RunOutput(
+        output="parsed test output", intermediate_outputs={"key": "value"}
+    )
+
+    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,
+        ),
+        patch("kiln_ai.adapters.model_adapters.base_adapter.Config") as mock_config,
+        patch.object(adapter, "has_structured_output", return_value=True),
+    ):
+        # Disable autosaving for this test
+        mock_config.shared.return_value.autosave_runs = False
+        mock_config.shared.return_value.user_id = "test_user_id"
+
+        # Execute
+        result = await adapter.invoke("test input")
+
+        # Verify parsing occurred
+        mock_parser.parse_output.assert_called_once()
+        parsed_args = mock_parser.parse_output.call_args[1]
+        assert isinstance(parsed_args["original_output"], RunOutput)
+        assert parsed_args["original_output"].output == "test output"
+
+        # Verify result contains parsed output
+        assert isinstance(result, TaskRun)
+        assert result.output.output == "parsed test output"
+        assert result.intermediate_outputs == {"key": "value"}
+        assert result.input == "test input"
+
+
+async def test_invoke_parsing_flow_with_reasoning_and_structured_output(adapter):
+    """Test for reasoning_optional_for_structured_output
+    when reasoning is provided with structured output enabled.
+    This is a special case where we want to return the output as is.
+    """
+    # Mock dependencies
+    mock_provider = MagicMock()
+    mock_provider.parser = "test_parser"
+    mock_provider.formatter = None
+    mock_provider.reasoning_capable = True
+    mock_provider.reasoning_optional_for_structured_output = True
+
+    mock_parser = MagicMock()
+    mock_parser.parse_output.return_value = RunOutput(
+        output="parsed test output", intermediate_outputs={"reasoning": "value"}
+    )
+
+    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,
+        ),
+        patch("kiln_ai.adapters.model_adapters.base_adapter.Config") as mock_config,
+        patch.object(adapter, "has_structured_output", return_value=True),
+    ):
+        # Disable autosaving for this test
+        mock_config.shared.return_value.autosave_runs = False
+        mock_config.shared.return_value.user_id = "test_user_id"
+
+        # Execute
+        result = await adapter.invoke("test input")
+
+        # Verify parsing occurred
+        mock_parser.parse_output.assert_called_once()
+        parsed_args = mock_parser.parse_output.call_args[1]
+        assert isinstance(parsed_args["original_output"], RunOutput)
+        assert parsed_args["original_output"].output == "test output"
+
+        # Verify result contains parsed output
+        assert isinstance(result, TaskRun)
+        assert result.output.output == "parsed test output"
+        assert result.intermediate_outputs == {"reasoning": "value"}
+        assert result.input == "test input"

kiln_ai/datamodel/test_eval_model.py

@@ -517,13 +517,13 @@ def test_eval_run_score_keys_must_match(valid_eval_config, valid_eval_run_data):
     valid_eval_config.parent = eval
 
     # Correct
-    run = EvalRun(
+    EvalRun(
         parent=valid_eval_config,
         **{**valid_eval_run_data, "scores": {"accuracy": 4.5, "critical": 1.0}},
     )
 
     # Correct but wrong order still okay
-    run = EvalRun(
+    EvalRun(
         parent=valid_eval_config,
         **{**valid_eval_run_data, "scores": {"critical": 1.0, "accuracy": 4.5}},
     )
@@ -533,7 +533,7 @@ def test_eval_run_score_keys_must_match(valid_eval_config, valid_eval_run_data):
         ValueError,
         match="The scores produced by the evaluator must match the scores expected by the eval",
     ):
-        run = EvalRun(
+        EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"accuracy": 4.5}},
         )
@@ -543,7 +543,7 @@ def test_eval_run_score_keys_must_match(valid_eval_config, valid_eval_run_data):
         ValueError,
         match="The scores produced by the evaluator must match the scores expected by the eval",
     ):
-        run = EvalRun(
+        EvalRun(
             parent=valid_eval_config,
             **{
                 **valid_eval_run_data,
@@ -556,7 +556,7 @@ def test_eval_run_score_keys_must_match(valid_eval_config, valid_eval_run_data):
         ValueError,
         match="The scores produced by the evaluator must match the scores expected by the eval",
     ):
-        run = EvalRun(
+        EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"accuracy": 4.5, "wrong": 1.0}},
         )
@@ -566,7 +566,7 @@ def test_eval_run_custom_scores_not_allowed(valid_eval_config, valid_eval_run_da
     with pytest.raises(
         ValueError, match="Custom scores are not supported in evaluators"
     ):
-        eval = Eval(
+        Eval(
             name="Test Eval",
             eval_set_filter_id="tag::tag1",
             eval_configs_filter_id="tag::tag2",