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

kiln_ai/datamodel/task_output.py

@@ -0,0 +1,321 @@
+import json
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, List, Type, Union
+
+import jsonschema
+import jsonschema.exceptions
+from pydantic import BaseModel, Field, ValidationInfo, model_validator
+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
+from kiln_ai.datamodel.strict_mode import strict_mode
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.task import Task
+
+
+class RequirementRating(BaseModel):
+    """Rating for a specific requirement within a task output."""
+
+    value: float = Field(
+        description="The rating value. Interpretation depends on rating type"
+    )
+    type: TaskOutputRatingType = Field(description="The type of rating")
+
+
+def normalize_rating(rating: float, rating_type: TaskOutputRatingType) -> float:
+    """Normalize a rating to a 0-1 scale. Simple normalization, not z-score."""
+    match rating_type:
+        case TaskOutputRatingType.five_star:
+            if rating < 1 or rating > 5:
+                raise ValueError("Five star rating must be between 1 and 5")
+            return (rating - 1) / 4
+        case TaskOutputRatingType.pass_fail:
+            if rating < 0 or rating > 1:
+                raise ValueError("Pass fail rating must 0 to 1")
+            return rating
+        case TaskOutputRatingType.pass_fail_critical:
+            if rating < -1 or rating > 1:
+                raise ValueError("Pass fail critical rating must -1 to 1")
+            return (rating + 1) / 2  # -1 to 1
+        case TaskOutputRatingType.custom:
+            raise ValueError("Custom rating type can not be normalized")
+        case _:
+            raise_exhaustive_enum_error(rating_type)
+
+
+class TaskOutputRating(KilnBaseModel):
+    """
+    A rating for a task output, including an overall rating and ratings for each requirement.
+
+    Supports:
+    - five_star: 1-5 star ratings
+    - pass_fail: boolean pass/fail (1.0 = pass, 0.0 = fail)
+    - pass_fail_critical: tri-state (1.0 = pass, 0.0 = fail, -1.0 = critical fail)
+    """
+
+    type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
+    value: float | None = Field(
+        description="The rating value. Interpretation depends on rating type:\n- five_star: 1-5 stars\n- pass_fail: 1.0 (pass) or 0.0 (fail)\n- pass_fail_critical: 1.0 (pass), 0.0 (fail), or -1.0 (critical fail)",
+        default=None,
+    )
+    requirement_ratings: Dict[ID_TYPE, RequirementRating] = Field(
+        default={},
+        description="The ratings of the requirements of the task.",
+    )
+
+    # Previously we stored rating values as a dict of floats, but now we store them as RequirementRating objects.
+    @model_validator(mode="before")
+    def upgrade_old_format(cls, data: dict) -> dict:
+        if not isinstance(data, dict):
+            return data
+
+        # Check if we have the old format (dict of floats)
+        req_ratings = data.get("requirement_ratings", {})
+        if req_ratings and all(
+            isinstance(v, (int, float)) for v in req_ratings.values()
+        ):
+            # Convert each float to a RequirementRating object
+            # all ratings are five star at the point we used this format
+            data["requirement_ratings"] = {
+                k: {"value": v, "type": TaskOutputRatingType.five_star}
+                for k, v in req_ratings.items()
+            }
+
+        return data
+
+    # Used to select high quality outputs for example selection (MultiShotPromptBuilder, etc)
+    def is_high_quality(self) -> bool:
+        if self.value is None:
+            return False
+
+        if self.type == TaskOutputRatingType.five_star:
+            return self.value >= 4
+        elif self.type == TaskOutputRatingType.pass_fail:
+            return self.value == 1.0
+        elif self.type == TaskOutputRatingType.pass_fail_critical:
+            return self.value == 1.0
+        return False
+
+    @model_validator(mode="after")
+    def validate_rating(self) -> Self:
+        if self.type not in TaskOutputRatingType:
+            raise ValueError(f"Invalid rating type: {self.type}")
+
+        # Overall rating is optional
+        if self.value is not None:
+            self._validate_rating(self.type, self.value, "overall rating")
+
+        for req_id, req_rating in self.requirement_ratings.items():
+            self._validate_rating(
+                req_rating.type,
+                req_rating.value,
+                f"requirement rating for req ID: {req_id}",
+            )
+
+        return self
+
+    def _validate_rating(
+        self, type: TaskOutputRatingType, rating: float | None, rating_name: str
+    ) -> None:
+        if type == TaskOutputRatingType.five_star:
+            self._validate_five_star(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail:
+            self._validate_pass_fail(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail_critical:
+            self._validate_pass_fail_critical(rating, rating_name)
+
+    def _validate_five_star(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type five_star must be an integer value (1-5)"
+            )
+        if rating < 1 or rating > 5:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type five_star must be between 1 and 5 stars"
+            )
+
+    def _validate_pass_fail(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be an integer value (0 or 1)"
+            )
+        if rating not in [0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be 0 (fail) or 1 (pass)"
+            )
+
+    def _validate_pass_fail_critical(
+        self, rating: float | None, rating_name: str
+    ) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be an integer value (-1, 0, or 1)"
+            )
+        if rating not in [-1, 0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be -1 (critical fail), 0 (fail), or 1 (pass)"
+            )
+
+
+class DataSourceType(str, Enum):
+    """
+    The source type of a piece of data.
+
+    Human: a human created the data
+    Synthetic: a model created the data
+    """
+
+    human = "human"
+    synthetic = "synthetic"
+
+
+class DataSourceProperty(BaseModel):
+    """
+    Defines a property that can be associated with a data source.
+
+    Includes validation rules for when properties are required or not allowed
+    based on the data source type.
+    """
+
+    name: str
+    type: Type[Union[str, int, float]]
+    required_for: List[DataSourceType] = []
+    not_allowed_for: List[DataSourceType] = []
+
+
+class DataSource(BaseModel):
+    """
+    Represents the origin of data, either human or synthetic, with associated properties.
+
+    Properties vary based on the source type - for synthetic sources this includes
+    model information, for human sources this includes creator information.
+    """
+
+    type: DataSourceType
+    properties: Dict[str, str | int | float] = Field(
+        default={},
+        description="Properties describing the data source. For synthetic things like model. For human, the human's name.",
+    )
+
+    _data_source_properties = [
+        DataSourceProperty(
+            name="created_by",
+            type=str,
+            required_for=[DataSourceType.human],
+            not_allowed_for=[DataSourceType.synthetic],
+        ),
+        DataSourceProperty(
+            name="model_name",
+            type=str,
+            required_for=[DataSourceType.synthetic],
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            name="model_provider",
+            type=str,
+            required_for=[DataSourceType.synthetic],
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            name="adapter_name",
+            type=str,
+            required_for=[DataSourceType.synthetic],
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            # Legacy field -- allow loading from old runs, but we shouldn't be setting it.
+            name="prompt_builder_name",
+            type=str,
+            not_allowed_for=[DataSourceType.human],
+        ),
+        DataSourceProperty(
+            # The PromptId of the prompt. Can be a saved prompt, fine-tune, generator name, etc. See PromptId type for more details.
+            name="prompt_id",
+            type=str,
+            not_allowed_for=[DataSourceType.human],
+        ),
+    ]
+
+    @model_validator(mode="after")
+    def validate_type(self) -> "DataSource":
+        if self.type not in DataSourceType:
+            raise ValueError(f"Invalid data source type: {self.type}")
+        return self
+
+    @model_validator(mode="after")
+    def validate_properties(self) -> "DataSource":
+        for prop in self._data_source_properties:
+            # Check the property type is correct
+            if prop.name in self.properties:
+                if not isinstance(self.properties[prop.name], prop.type):
+                    raise ValueError(
+                        f"'{prop.name}' must be of type {prop.type.__name__} for {self.type} data source"
+                    )
+            # Check the property is required for the data source type
+            if self.type in prop.required_for:
+                if prop.name not in self.properties:
+                    raise ValueError(
+                        f"'{prop.name}' is required for {self.type} data source"
+                    )
+            # Check the property is not allowed for the data source type
+            elif self.type in prop.not_allowed_for and prop.name in self.properties:
+                raise ValueError(
+                    f"'{prop.name}' is not allowed for {self.type} data source"
+                )
+        return self
+
+    @model_validator(mode="after")
+    def validate_no_empty_properties(self) -> Self:
+        for prop, value in self.properties.items():
+            if isinstance(value, str) and value == "":
+                raise ValueError(
+                    f"Property '{prop}' must be a non-empty string for {self.type} data source"
+                )
+        return self
+
+
+class TaskOutput(KilnBaseModel):
+    """
+    An output for a specific task run.
+
+    Contains the actual output content, its source (human or synthetic),
+    and optional rating information.
+    """
+
+    output: str = Field(
+        description="The output of the task. JSON formatted for structured output, plaintext for unstructured output."
+    )
+    source: DataSource | None = Field(
+        description="The source of the output: human or synthetic.",
+        default=None,
+    )
+    rating: TaskOutputRating | None = Field(
+        default=None, description="The rating of the output"
+    )
+
+    def validate_output_format(self, task: "Task") -> Self:
+        # validate output
+        if task.output_json_schema is not None:
+            try:
+                validate_schema(json.loads(self.output), task.output_json_schema)
+            except json.JSONDecodeError:
+                raise ValueError("Output is not a valid JSON object")
+            except jsonschema.exceptions.ValidationError as e:
+                raise ValueError(f"Output does not match task output schema: {e}")
+        return self
+
+    @model_validator(mode="after")
+    def validate_output_source(self, info: ValidationInfo) -> Self:
+        # On strict mode and not loaded from file, we validate output_source is not None.
+        # We want to be able to load any data, even if it's not perfect. But we want to create perfect data when adding new data.
+        if not strict_mode():
+            return self
+        if self.loaded_from_file(info):
+            return self
+        if self.source is None:
+            raise ValueError("Output source is required when strict mode is enabled")
+        return self

kiln_ai/datamodel/task_run.py

@@ -0,0 +1,164 @@
+import json
+from typing import TYPE_CHECKING, Dict, List, Union
+
+import jsonschema
+import jsonschema.exceptions
+from pydantic import 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.strict_mode import strict_mode
+from kiln_ai.datamodel.task_output import DataSource, TaskOutput
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.task import Task
+
+
+class TaskRun(KilnParentedModel):
+    """
+    Represents a single execution of a Task.
+
+    Contains the input used, its source, the output produced, and optional
+    repair information if the output needed correction.
+    """
+
+    input: str = Field(
+        description="The inputs to the task. JSON formatted for structured input, plaintext for unstructured input."
+    )
+    input_source: DataSource | None = Field(
+        default=None, description="The source of the input: human or synthetic."
+    )
+
+    output: TaskOutput = Field(description="The output of the task run.")
+    repair_instructions: str | None = Field(
+        default=None,
+        description="Instructions for fixing the output. Should define what is wrong, and how to fix it. Will be used by models for both generating a fixed output, and evaluating future models.",
+    )
+    repaired_output: TaskOutput | None = Field(
+        default=None,
+        description="An version of the output with issues fixed. This must be a 'fixed' version of the existing output, and not an entirely new output. If you wish to generate an ideal curatorial output for this task unrelated to this output, generate a new TaskOutput with type 'human' instead of using this field.",
+    )
+    intermediate_outputs: Dict[str, str] | None = Field(
+        default=None,
+        description="Intermediate outputs from the task run. Keys are the names of the intermediate output steps (cot=chain of thought, etc), values are the output data.",
+    )
+    tags: List[str] = Field(
+        default=[],
+        description="Tags for the task run. Tags are used to categorize task runs for filtering and reporting.",
+    )
+
+    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
+        )
+
+    # Workaround to return typed parent without importing Task
+    def parent_task(self) -> Union["Task", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Task":
+            return None
+        return self.parent  # type: ignore
+
+    @model_validator(mode="after")
+    def validate_input_format(self, info: ValidationInfo) -> Self:
+        # Don't validate if loading from file (not new). Too slow.
+        # We don't allow changing task schema, so this is redundant validation.
+        # Note: we still validate if editing a loaded model
+        if self.loading_from_file(info):
+            # Consider loading an existing model as validated.
+            self._last_validated_input = self.input
+            return self
+
+        # Don't validate if input has not changed. Too slow to run this every time.
+        if (
+            hasattr(self, "_last_validated_input")
+            and self.input == self._last_validated_input
+        ):
+            return self
+
+        task = self.parent_task()
+        if task is None:
+            # 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
+        if task.input_json_schema is not None:
+            try:
+                validate_schema(json.loads(self.input), task.input_json_schema)
+            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}")
+        self._last_validated_input = self.input
+        return self
+
+    @model_validator(mode="after")
+    def validate_output_format(self, info: ValidationInfo) -> Self:
+        # Don't validate if loading from file (not new). Too slow.
+        # Note: we still validate if editing a loaded model's output.
+        if self.loading_from_file(info):
+            # Consider loading an existing model as validated.
+            self._last_validated_output = self.output.output if self.output else None
+            return self
+
+        # Don't validate unless output has changed since last validation.
+        # The validator is slow and costly, don't want it running when setting other fields.
+        if (
+            hasattr(self, "_last_validated_output")
+            and self.output is not None
+            and self.output.output == self._last_validated_output
+        ):
+            return self
+
+        task = self.parent_task()
+        if task is None:
+            return self
+
+        self.output.validate_output_format(task)
+        self._last_validated_output = self.output.output if self.output else None
+        return self
+
+    @model_validator(mode="after")
+    def validate_repaired_output(self) -> Self:
+        if self.repaired_output is not None:
+            if self.repaired_output.rating is not None:
+                raise ValueError(
+                    "Repaired output rating must be None. Repaired outputs are assumed to have a perfect rating, as they have been fixed."
+                )
+        if self.repair_instructions is None and self.repaired_output is not None:
+            raise ValueError(
+                "Repair instructions are required if providing a repaired output."
+            )
+        if self.repair_instructions is not None and self.repaired_output is None:
+            raise ValueError(
+                "A repaired output is required if providing repair instructions."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def validate_input_source(self, info: ValidationInfo) -> Self:
+        # On strict mode and not loaded from file, we validate input_source is not None.
+        # We want to be able to load any data, even if it's not perfect. But we want to create perfect data when adding new data.
+        if not strict_mode():
+            return self
+        if self.loaded_from_file(info):
+            return self
+        if self.input_source is None:
+            raise ValueError("input_source is required when strict mode is enabled")
+        return self
+
+    @model_validator(mode="after")
+    def validate_tags(self) -> Self:
+        for tag in self.tags:
+            if not tag:
+                raise ValueError("Tags cannot be empty strings")
+            if " " in tag:
+                raise ValueError("Tags cannot contain spaces. Try underscores.")
+
+        return self

kiln_ai/datamodel/test_basemodel.py

@@ -6,7 +6,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from kiln_ai.adapters.model_adapters.base_adapter import AdapterInfo, BaseAdapter
+from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter
 from kiln_ai.adapters.run_output import RunOutput
 from kiln_ai.datamodel import Task, TaskRun
 from kiln_ai.datamodel.basemodel import (
@@ -15,6 +15,7 @@ from kiln_ai.datamodel.basemodel import (
     string_to_valid_name,
 )
 from kiln_ai.datamodel.model_cache import ModelCache
+from kiln_ai.datamodel.task import RunConfig
 
 
 @pytest.fixture
@@ -484,13 +485,8 @@ class MockAdapter(BaseAdapter):
     async def _run(self, input):
         return RunOutput(output="test output", intermediate_outputs=None)
 
-    def adapter_info(self) -> AdapterInfo:
-        return AdapterInfo(
-            adapter_name="test",
-            model_name=self.model_name,
-            model_provider=self.model_provider_name,
-            prompt_builder_name="test",
-        )
+    def adapter_name(self) -> str:
+        return "test"
 
 
 @pytest.fixture
@@ -501,9 +497,12 @@ def base_task():
 @pytest.fixture
 def adapter(base_task):
     return MockAdapter(
-        kiln_task=base_task,
-        model_name="test_model",
-        model_provider_name="test_provider",
+        run_config=RunConfig(
+            task=base_task,
+            model_name="test_model",
+            model_provider_name="test_provider",
+            prompt_id="simple_prompt_builder",
+        ),
     )
 
 

kiln_ai/datamodel/test_dataset_filters.py

@@ -0,0 +1,71 @@
+import pytest
+from pydantic import BaseModel
+
+from kiln_ai.datamodel.dataset_filters import (
+    AllDatasetFilter,
+    DatasetFilterId,
+    HighRatingDatasetFilter,
+    StaticDatasetFilters,
+    TagFilter,
+    ThinkingModelDatasetFilter,
+    ThinkingModelHighRatedFilter,
+    dataset_filter_from_id,
+)
+
+# Note: Many more filter tests in test_dataset_split.py
+
+
+def test_all_dataset_filter_from_id():
+    assert dataset_filter_from_id("all") == AllDatasetFilter
+
+
+def test_high_rating_dataset_filter_from_id():
+    assert dataset_filter_from_id("high_rating") == HighRatingDatasetFilter
+
+
+def test_thinking_model_dataset_filter_from_id():
+    assert dataset_filter_from_id("thinking_model") == ThinkingModelDatasetFilter
+
+
+def test_thinking_model_high_rated_dataset_filter_from_id():
+    assert (
+        dataset_filter_from_id("thinking_model_high_rated")
+        == ThinkingModelHighRatedFilter
+    )
+
+
+def test_all_static_dataset_filters():
+    for filter_id in StaticDatasetFilters:
+        assert dataset_filter_from_id(filter_id) is not None
+
+
+class ModelTester(BaseModel):
+    dsid: DatasetFilterId
+
+
+@pytest.mark.parametrize(
+    "tag,expected_error,expected_tag",
+    [
+        ("tag::test", False, "test"),
+        ("tag::other", False, "other"),
+        ("tag::", True, None),
+        ("tag", True, None),
+        ("", True, None),
+    ],
+)
+def test_tag_filter(tag, expected_error, expected_tag):
+    # Check our model validators
+    if expected_error:
+        with pytest.raises(ValueError):
+            ModelTester(dsid=tag)
+    else:
+        ModelTester(dsid=tag)
+
+    # Check the constructor
+    if expected_tag is None:
+        with pytest.raises(ValueError, match="Invalid dataset filter ID:"):
+            dataset_filter_from_id(tag)
+    else:
+        filter = dataset_filter_from_id(tag)
+        assert isinstance(filter, TagFilter)
+        assert filter.tag == expected_tag

kiln_ai/datamodel/test_dataset_split.py

@@ -3,24 +3,28 @@ from pydantic import ValidationError
 
 # import datamodel first or we get circular import errors
 from kiln_ai.datamodel import (
-    AllDatasetFilter,
-    AllSplitDefinition,
-    DatasetFilterType,
     DatasetSplit,
     DatasetSplitDefinition,
     DataSource,
     DataSourceType,
-    HighRatingDatasetFilter,
     Task,
     TaskOutput,
     TaskOutputRating,
     TaskOutputRatingType,
     TaskRun,
-    ThinkingModelDatasetFilter,
-    ThinkingModelHighRatedFilter,
+)
+from kiln_ai.datamodel.dataset_split import (
+    AllSplitDefinition,
     Train60Test20Val20SplitDefinition,
     Train80Test20SplitDefinition,
 )
+from kiln_ai.datamodel.test_dataset_filters import (
+    AllDatasetFilter,
+    HighRatingDatasetFilter,
+    TagFilter,
+    ThinkingModelDatasetFilter,
+    ThinkingModelHighRatedFilter,
+)
 
 
 @pytest.fixture
@@ -42,6 +46,7 @@ def sample_task_runs(sample_task):
     task_runs = []
     for i in range(10):
         rating = 5 if i < 6 else 1  # 6 high, 4 low ratings
+        tags = ["tag1"] if i < 6 else []
         task_run = TaskRun(
             parent=sample_task,
             input=f"input_{i}",
@@ -59,6 +64,7 @@ def sample_task_runs(sample_task):
                     value=rating, type=TaskOutputRatingType.five_star
                 ),
             ),
+            tags=tags,
         )
         task_run.save_to_file()
         task_runs.append(task_run)
@@ -199,10 +205,10 @@ def test_dataset_split_with_high_rating_filter(sample_task, sample_task_runs):
         "Split Name",
         sample_task,
         Train80Test20SplitDefinition,
-        filter_type=DatasetFilterType.HIGH_RATING,
+        filter_id="high_rating",
     )
 
-    assert dataset.filter == DatasetFilterType.HIGH_RATING
+    assert dataset.filter == "high_rating"
 
     # Check that only high-rated task runs are included
     all_ids = []
@@ -329,3 +335,21 @@ def test_thinking_model_dataset_filter_high_rated(
     )
 
     assert ThinkingModelHighRatedFilter(task_run) is expected_result
+
+
+def test_tag_dataset_filter(sample_task_runs):
+    num_tagged = 0
+    num_untagged = 0
+    filter = TagFilter("tag1")
+    for task_run in sample_task_runs:
+        if "tag1" in task_run.tags:
+            num_tagged += 1
+            assert "tag1" in task_run.tags
+            assert filter(task_run) is True
+        else:
+            num_untagged += 1
+            assert "tag1" not in task_run.tags
+            assert filter(task_run) is False
+
+    assert num_tagged == 6
+    assert num_untagged == 4