kiln-ai 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

kiln_ai/datamodel/__init__.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import json
 from enum import Enum, IntEnum
-from typing import TYPE_CHECKING, Dict, List, Self, Type, Union
+from typing import TYPE_CHECKING, Dict, List, Type, Union
 
 import jsonschema
 import jsonschema.exceptions
-from kiln_ai.datamodel.json_schema import JsonObjectSchema, schema_from_json_str
 from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
+
+from kiln_ai.datamodel.json_schema import JsonObjectSchema, schema_from_json_str
 
 from .basemodel import (
     ID_FIELD,
@@ -21,6 +23,25 @@ from .json_schema import validate_schema
 if TYPE_CHECKING:
     from . import Task
 
+
+__all__ = [
+    "basemodel",
+    "json_schema",
+    "Task",
+    "Project",
+    "TaskRun",
+    "TaskOutput",
+    "TaskOutputRating",
+    "Priority",
+    "DataSource",
+    "DataSourceType",
+    "DataSourceProperty",
+    "TaskOutputRatingType",
+    "TaskRequirement",
+    "TaskDeterminism",
+]
+
+
 # Conventions:
 # 1) Names are filename safe as they may be used as file names. They are informational and not to be used in prompts/training/validation.
 # 2) Descrptions are for Kiln users to describe/understanding the purpose of this object. They must never be used in prompts/training/validation. Use "instruction/requirements" instead.
@@ -32,6 +53,8 @@ SHORT_NAME_FIELD = Field(min_length=1, max_length=20, pattern=NAME_REGEX)
 
 
 class Priority(IntEnum):
+    """Defines priority levels for tasks and requirements, where P0 is highest priority."""
+
     p0 = 0
     p1 = 1
     p2 = 2
@@ -40,6 +63,8 @@ class Priority(IntEnum):
 
 # Only one rating type for now, but this allows for extensibility if we want to add more in the future
 class TaskOutputRatingType(str, Enum):
+    """Defines the types of rating systems available for task outputs."""
+
     five_star = "five_star"
     custom = "custom"
 
@@ -90,22 +115,13 @@ class TaskOutputRating(KilnBaseModel):
                 f"{rating_name.capitalize()} of type five_star must be between 1 and 5 stars"
             )
 
-    def validate_requirement_rating_keys(self, task: Task) -> Self:
-        if len(self.requirement_ratings) == 0:
-            return self
-
-        valid_requirement_ids = {req.id for req in task.requirements}
-        for key in self.requirement_ratings.keys():
-            if key not in valid_requirement_ids:
-                raise ValueError(
-                    f"Requirement ID '{key}' is not a valid requirement ID for this task"
-                )
-        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(
@@ -132,7 +148,10 @@ class TaskOutput(KilnBaseModel):
 
 class DataSourceType(str, Enum):
     """
-    The source of a piece of data.
+    The source type of a piece of data.
+
+    Human: a human created the data
+    Synthetic: a model created the data
     """
 
     human = "human"
@@ -140,6 +159,13 @@ class DataSourceType(str, Enum):
 
 
 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] = []
@@ -147,6 +173,13 @@ class DataSourceProperty(BaseModel):
 
 
 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={},
@@ -225,7 +258,10 @@ class DataSource(BaseModel):
 
 class TaskRun(KilnParentedModel):
     """
-    An run of a specific Task, including the input and output.
+    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(
@@ -276,19 +312,6 @@ class TaskRun(KilnParentedModel):
         self.output.validate_output_format(task)
         return self
 
-    @model_validator(mode="after")
-    def validate_requirement_ratings(self) -> Self:
-        task = self.parent_task()
-        if task is None:
-            return self
-
-        if self.output.rating is not None:
-            self.output.rating.validate_requirement_rating_keys(task)
-        if self.repaired_output is not None and self.repaired_output.rating is not None:
-            self.repaired_output.rating.validate_requirement_rating_keys(task)
-
-        return self
-
     @model_validator(mode="after")
     def validate_repaired_output(self) -> Self:
         if self.repaired_output is not None:
@@ -308,6 +331,13 @@ class TaskRun(KilnParentedModel):
 
 
 class TaskRequirement(BaseModel):
+    """
+    Defines a specific requirement that should be met by task outputs.
+
+    Includes an identifier, name, description, instruction for meeting the requirement,
+    and priority level.
+    """
+
     id: ID_TYPE = ID_FIELD
     name: str = SHORT_NAME_FIELD
     description: str | None = Field(default=None)
@@ -316,6 +346,14 @@ class TaskRequirement(BaseModel):
 
 
 class TaskDeterminism(str, Enum):
+    """
+    Defines how strictly task outputs should match expected results.
+
+    - deterministic: Requires exact matches
+    - semantic_match: Allows different wording with same meaning
+    - flexible: Allows variation in both wording and meaning within requirements
+    """
+
     deterministic = "deterministic"  # Expect exact match
     semantic_match = "semantic_match"  # Expect same meaning, but flexible on expression of the meaning
     flexible = "flexible"  # Flexible on semantic output. Eval should be custom based on parsing requirements.
@@ -326,6 +364,13 @@ class Task(
     KilnParentModel,
     parent_of={"runs": TaskRun},
 ):
+    """
+    Represents a specific task to be performed, with associated requirements and validation rules.
+
+    Contains the task definition, requirements, input/output schemas, and maintains
+    a collection of task runs.
+    """
+
     name: str = NAME_FIELD
     description: str = Field(default="")
     priority: Priority = Field(default=Priority.p2)
@@ -352,6 +397,13 @@ class Task(
 
 
 class Project(KilnParentModel, parent_of={"tasks": Task}):
+    """
+    A collection of related tasks.
+
+    Projects organize tasks into logical groups and provide high-level descriptions
+    of the overall goals.
+    """
+
     name: str = NAME_FIELD
     description: str | None = Field(
         default=None,

kiln_ai/datamodel/basemodel.py

@@ -1,4 +1,5 @@
 import json
+import shutil
 import uuid
 from abc import ABCMeta
 from builtins import classmethod
@@ -10,13 +11,10 @@ from typing import (
     Dict,
     List,
     Optional,
-    Self,
     Type,
     TypeVar,
 )
 
-from kiln_ai.utils.config import Config
-from kiln_ai.utils.formatting import snake_case
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -26,6 +24,10 @@ from pydantic import (
     model_validator,
 )
 from pydantic_core import ErrorDetails
+from typing_extensions import Self
+
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.formatting import snake_case
 
 # ID is a 12 digit random integer string.
 # Should be unique per item, at least inside the context of a parent/child relationship.
@@ -38,6 +40,16 @@ PT = TypeVar("PT", bound="KilnParentedModel")
 
 
 class KilnBaseModel(BaseModel):
+    """Base model for all Kiln data models with common functionality for persistence and versioning.
+
+    Attributes:
+        v (int): Schema version number for migration support
+        id (str): Unique identifier for the model instance
+        path (Path): File system path where the model is stored
+        created_at (datetime): Timestamp when the model was created
+        created_by (str): User ID of the creator
+    """
+
     model_config = ConfigDict(validate_assignment=True)
 
     v: int = Field(default=1)  # schema_version
@@ -62,11 +74,30 @@ class KilnBaseModel(BaseModel):
 
     @classmethod
     def load_from_folder(cls: Type[T], folderPath: Path) -> T:
+        """Load a model instance from a folder using the default filename.
+
+        Args:
+            folderPath (Path): Directory path containing the model file
+
+        Returns:
+            T: Instance of the model
+        """
         path = folderPath / cls.base_filename()
         return cls.load_from_file(path)
 
     @classmethod
     def load_from_file(cls: Type[T], path: Path) -> T:
+        """Load a model instance from a specific file path.
+
+        Args:
+            path (Path): Path to the model file
+
+        Returns:
+            T: Instance of the model
+
+        Raises:
+            ValueError: If the loaded model is not of the expected type or version
+        """
         with open(path, "r") as file:
             file_data = file.read()
             # TODO P2 perf: parsing the JSON twice here.
@@ -92,6 +123,11 @@ class KilnBaseModel(BaseModel):
         return m
 
     def save_to_file(self) -> None:
+        """Save the model instance to a file.
+
+        Raises:
+            ValueError: If the path is not set
+        """
         path = self.build_path()
         if path is None:
             raise ValueError(
@@ -105,6 +141,15 @@ class KilnBaseModel(BaseModel):
         # save the path so even if something like name changes, the file doesn't move
         self.path = path
 
+    def delete(self) -> None:
+        if self.path is None:
+            raise ValueError("Cannot delete model because path is not set")
+        dir_path = self.path.parent if self.path.is_file() else self.path
+        if dir_path is None:
+            raise ValueError("Cannot delete model because path is not set")
+        shutil.rmtree(dir_path)
+        self.path = None
+
     def build_path(self) -> Path | None:
         if self.path is not None:
             return self.path
@@ -116,6 +161,15 @@ class KilnBaseModel(BaseModel):
 
 
 class KilnParentedModel(KilnBaseModel, metaclass=ABCMeta):
+    """Base model for Kiln models that have a parent-child relationship. This base class is for child models.
+
+    This class provides functionality for managing hierarchical relationships between models,
+    including parent reference handling and file system organization.
+
+    Attributes:
+        _parent (KilnBaseModel): Reference to the parent model instance
+    """
+
     _parent: KilnBaseModel | None = None
 
     # workaround to tell typechecker that we support the parent property, even though it's not a stock property
@@ -129,6 +183,11 @@ class KilnParentedModel(KilnBaseModel, metaclass=ABCMeta):
 
     @property
     def parent(self) -> Optional[KilnBaseModel]:
+        """Get the parent model instance, loading it from disk if necessary.
+
+        Returns:
+            Optional[KilnBaseModel]: The parent model instance or None if not set
+        """
         if self._parent is not None:
             return self._parent
         # lazy load parent from path
@@ -244,6 +303,15 @@ class KilnParentedModel(KilnBaseModel, metaclass=ABCMeta):
 # Parent create methods for all child relationships
 # You must pass in parent_of in the subclass definition, defining the child relationships
 class KilnParentModel(KilnBaseModel, metaclass=ABCMeta):
+    """Base model for Kiln models that can have child models.
+
+    This class provides functionality for managing collections of child models and their persistence.
+    Child relationships must be defined using the parent_of parameter in the class definition.
+
+    Args:
+        parent_of (Dict[str, Type[KilnParentedModel]]): Mapping of relationship names to child model types
+    """
+
     @classmethod
     def _create_child_method(
         cls, relationship_name: str, child_class: Type[KilnParentedModel]
@@ -288,6 +356,19 @@ class KilnParentModel(KilnBaseModel, metaclass=ABCMeta):
         path: Path | None = None,
         parent: KilnBaseModel | None = None,
     ):
+        """Validate and save a model instance along with all its nested child relationships.
+
+        Args:
+            data (Dict[str, Any]): Model data including child relationships
+            path (Path, optional): Path where the model should be saved
+            parent (KilnBaseModel, optional): Parent model instance for parented models
+
+        Returns:
+            KilnParentModel: The validated and saved model instance
+
+        Raises:
+            ValidationError: If validation fails for the model or any of its children
+        """
         # Validate first, then save. Don't want error half way through, and partly persisted
         # TODO P2: save to tmp dir, then move atomically. But need to merge directories so later.
         cls._validate_nested(data, save=False, path=path, parent=parent)

kiln_ai/datamodel/json_schema.py

@@ -10,21 +10,55 @@ JsonObjectSchema = Annotated[
     str,
     AfterValidator(lambda v: _check_json_schema(v)),
 ]
+"""A pydantic type that validates strings containing JSON schema definitions.
+Must be a valid JSON schema object with 'type': 'object' and 'properties' defined.
+"""
 
 
 def _check_json_schema(v: str) -> str:
-    # parsing returns needed errors
+    """Internal validation function for JSON schema strings.
+
+    Args:
+        v: String containing a JSON schema definition
+
+    Returns:
+        The input string if valid
+
+    Raises:
+        ValueError: If the schema is invalid
+    """
     schema_from_json_str(v)
     return v
 
 
 def validate_schema(instance: Dict, schema_str: str) -> None:
+    """Validate a dictionary against a JSON schema.
+
+    Args:
+        instance: Dictionary to validate
+        schema_str: JSON schema string to validate against
+
+    Raises:
+        jsonschema.exceptions.ValidationError: If validation fails
+        ValueError: If the schema is invalid
+    """
     schema = schema_from_json_str(schema_str)
     v = jsonschema.Draft202012Validator(schema)
     return v.validate(instance)
 
 
 def schema_from_json_str(v: str) -> Dict:
+    """Parse and validate a JSON schema string.
+
+    Args:
+        v: String containing a JSON schema definition
+
+    Returns:
+        Dict containing the parsed JSON schema
+
+    Raises:
+        ValueError: If the input is not a valid JSON schema object with required properties
+    """
     try:
         parsed = json.loads(v)
         jsonschema.Draft202012Validator.check_schema(parsed)

kiln_ai/datamodel/test_basemodel.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Optional
 
 import pytest
+
 from kiln_ai.datamodel.basemodel import KilnBaseModel, KilnParentedModel
 
 
@@ -275,3 +276,33 @@ def test_lazy_load_parent(tmp_path):
     # Verify that the _parent attribute is now set
     assert hasattr(loaded_child, "_parent")
     assert loaded_child._parent is loaded_parent
+
+
+def test_delete(tmp_path):
+    # Test deleting a file
+    file_path = tmp_path / "test.kiln"
+    model = KilnBaseModel(path=file_path)
+    model.save_to_file()
+    assert file_path.exists()
+    model.delete()
+    assert not file_path.exists()
+    assert not file_path.parent.exists()
+    assert model.path is None
+
+
+def test_delete_dir(tmp_path):
+    # Test deleting a directory
+    dir_path = tmp_path / "test_dir"
+    dir_path.mkdir(parents=True)
+    model = KilnBaseModel(path=dir_path)
+    assert dir_path.exists()
+    model.delete()
+    assert not dir_path.exists()
+    assert model.path is None
+
+
+def test_delete_no_path():
+    # Test deleting with no path
+    model = KilnBaseModel()
+    with pytest.raises(ValueError, match="Cannot delete model because path is not set"):
+        model.delete()

kiln_ai/datamodel/test_datasource.py

@@ -1,7 +1,8 @@
 import pytest
-from kiln_ai.datamodel import DataSource, DataSourceType
 from pydantic import ValidationError
 
+from kiln_ai.datamodel import DataSource, DataSourceType
+
 
 def test_valid_human_data_source():
     data_source = DataSource(
@@ -29,16 +30,14 @@ def test_valid_synthetic_data_source():
 
 
 def test_missing_required_property():
-    with pytest.raises(
-        ValidationError, match="'created_by' is required for DataSourceType.human data"
-    ):
+    with pytest.raises(ValidationError, match="'created_by' is required for"):
         DataSource(type=DataSourceType.human)
 
 
 def test_wrong_property_type():
     with pytest.raises(
         ValidationError,
-        match="'model_name' must be of type str for DataSourceType.synthetic data",
+        match="'model_name' must be of type str for",
     ):
         DataSource(
             type=DataSourceType.synthetic,
@@ -49,7 +48,7 @@ def test_wrong_property_type():
 def test_not_allowed_property():
     with pytest.raises(
         ValidationError,
-        match="'created_by' is not allowed for DataSourceType.synthetic data",
+        match="'created_by' is not allowed for",
     ):
         DataSource(
             type=DataSourceType.synthetic,

kiln_ai/datamodel/test_example_models.py

@@ -1,6 +1,9 @@
 import json
+import sys
 
 import pytest
+from pydantic import ValidationError
+
 from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
@@ -13,7 +16,6 @@ from kiln_ai.datamodel import (
     TaskRequirement,
     TaskRun,
 )
-from pydantic import ValidationError
 
 
 @pytest.fixture
@@ -55,10 +57,11 @@ def test_task_model_validation(valid_task_run):
     with pytest.raises(ValidationError, match="Input should be"):
         DataSource(type="invalid")
 
-    with pytest.raises(ValidationError, match="Invalid data source type"):
-        task_run = valid_task_run.model_copy(deep=True)
-        task_run.input_source.type = "invalid"
-        DataSource.model_validate(task_run.input_source, strict=True)
+    if sys.version_info >= (3, 12):
+        with pytest.raises(ValidationError, match="Invalid data source type"):
+            task_run = valid_task_run.model_copy(deep=True)
+            task_run.input_source.type = "invalid"
+            DataSource.model_validate(task_run.input_source, strict=True)
 
     # Missing required field
     with pytest.raises(ValidationError, match="Input should be a valid string"):
@@ -72,9 +75,7 @@ def test_task_model_validation(valid_task_run):
         DataSource.model_validate(task_run.input_source, strict=True)
 
     # Test we catch nested validation errors
-    with pytest.raises(
-        ValidationError, match="'created_by' is required for DataSourceType.human"
-    ):
+    with pytest.raises(ValidationError, match="'created_by' is required for"):
         task_run = TaskRun(
             input="Test input",
             input_source=DataSource(
@@ -270,34 +271,6 @@ def test_task_output_requirement_rating_keys(tmp_path):
     task_run.save_to_file()
     assert task_run.output.rating.requirement_ratings is not None
 
-    # Invalid case: unknown requirement ID
-    with pytest.raises(
-        ValueError,
-        match="Requirement ID .* is not a valid requirement ID for this task",
-    ):
-        task_run = TaskRun(
-            input="Test input",
-            input_source=DataSource(
-                type=DataSourceType.human,
-                properties={"created_by": "john_doe"},
-            ),
-            parent=task,
-            output=TaskOutput(
-                output="Test output",
-                source=DataSource(
-                    type=DataSourceType.human,
-                    properties={"created_by": "john_doe"},
-                ),
-                rating=TaskOutputRating(
-                    value=4,
-                    requirement_ratings={
-                        "unknown_id": 5,
-                    },
-                ),
-            ),
-        )
-        task_run.save_to_file()
-
 
 def test_task_output_schema_validation(tmp_path):
     # Create a project, task, and example hierarchy
@@ -434,9 +407,7 @@ def test_valid_human_task_output():
 
 
 def test_invalid_human_task_output_missing_created_by():
-    with pytest.raises(
-        ValidationError, match="'created_by' is required for DataSourceType.human"
-    ):
+    with pytest.raises(ValidationError, match="'created_by' is required for"):
         TaskOutput(
             output="Test output",
             source=DataSource(
@@ -482,7 +453,7 @@ def test_valid_synthetic_task_output():
 def test_invalid_synthetic_task_output_missing_keys():
     with pytest.raises(
         ValidationError,
-        match="'model_provider' is required for DataSourceType.synthetic",
+        match="'model_provider' is required for",
     ):
         TaskOutput(
             output="Test output",

kiln_ai/datamodel/test_json_schema.py

@@ -1,10 +1,11 @@
 import pytest
+from pydantic import BaseModel
+
 from kiln_ai.datamodel.json_schema import (
     JsonObjectSchema,
     schema_from_json_str,
     validate_schema,
 )
-from pydantic import BaseModel
 
 
 class ExampleModel(BaseModel):

kiln_ai/datamodel/test_models.py

@@ -1,9 +1,10 @@
 import json
 
 import pytest
+from pydantic import ValidationError
+
 from kiln_ai.datamodel import Priority, Project, Task, TaskDeterminism
 from kiln_ai.datamodel.test_json_schema import json_joke_schema
-from pydantic import ValidationError
 
 
 @pytest.fixture

kiln_ai/datamodel/test_nested_save.py

@@ -1,7 +1,8 @@
 import pytest
-from kiln_ai.datamodel.basemodel import KilnParentedModel, KilnParentModel
 from pydantic import Field, ValidationError
 
+from kiln_ai.datamodel.basemodel import KilnParentedModel, KilnParentModel
+
 
 class ModelC(KilnParentedModel):
     code: str = Field(..., pattern=r"^[A-Z]{3}$")

kiln_ai/datamodel/test_output_rating.py

@@ -1,7 +1,8 @@
 import pytest
-from kiln_ai.datamodel import TaskOutputRating, TaskOutputRatingType
 from pydantic import ValidationError
 
+from kiln_ai.datamodel import TaskOutputRating, TaskOutputRatingType
+
 
 def test_valid_task_output_rating():
     rating = TaskOutputRating(value=4.0, requirement_ratings={"req1": 5.0, "req2": 3.0})

kiln_ai/utils/__init__.py

@@ -0,0 +1,12 @@
+"""
+# Utils
+
+Misc utilities used in the kiln_ai library.
+"""
+
+from . import config, formatting
+
+__all__ = [
+    "config",
+    "formatting",
+]

kiln_ai/utils/test_config.py

@@ -4,6 +4,7 @@ from unittest.mock import patch
 
 import pytest
 import yaml
+
 from kiln_ai.utils.config import Config, ConfigProperty, _get_user_id