kiln-ai 0.6.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

kiln_ai/datamodel/__init__.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import json
+import math
+import random
 from enum import Enum, IntEnum
-from typing import TYPE_CHECKING, Dict, List, Type, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Type, Union
 
 import jsonschema
 import jsonschema.exceptions
@@ -14,6 +16,8 @@ from kiln_ai.datamodel.json_schema import JsonObjectSchema, schema_from_json_str
 from .basemodel import (
     ID_FIELD,
     ID_TYPE,
+    NAME_FIELD,
+    SHORT_NAME_FIELD,
     KilnBaseModel,
     KilnParentedModel,
     KilnParentModel,
@@ -42,26 +46,6 @@ __all__ = [
 ]
 
 
-# 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.
-
-# Filename compatible names
-NAME_REGEX = r"^[A-Za-z0-9 _-]+$"
-NAME_FIELD = Field(
-    min_length=1,
-    max_length=120,
-    pattern=NAME_REGEX,
-    description="A name for this entity.",
-)
-SHORT_NAME_FIELD = Field(
-    min_length=1,
-    max_length=32,
-    pattern=NAME_REGEX,
-    description="A name for this entity",
-)
-
-
 class Priority(IntEnum):
     """Defines priority levels for tasks and requirements, where P0 is highest priority."""
 
@@ -156,6 +140,71 @@ class TaskOutput(KilnBaseModel):
         return self
 
 
+class FineTuneStatusType(str, Enum):
+    """
+    The status type of a fine-tune (running, completed, failed, etc).
+    """
+
+    unknown = "unknown"  # server error
+    pending = "pending"
+    running = "running"
+    completed = "completed"
+    failed = "failed"
+
+
+class Finetune(KilnParentedModel):
+    name: str = NAME_FIELD
+    description: str | None = Field(
+        default=None,
+        description="A description of the fine-tune for you and your team. Not used in training.",
+    )
+    provider: str = Field(
+        description="The provider to use for the fine-tune (e.g. 'openai')."
+    )
+    base_model_id: str = Field(
+        description="The id of the base model to use for the fine-tune. This string relates to the provider's IDs for their own models, not Kiln IDs."
+    )
+    provider_id: str | None = Field(
+        default=None,
+        description="The ID of the fine-tune job on the provider's side. May not be the same as the fine_tune_model_id.",
+    )
+    fine_tune_model_id: str | None = Field(
+        default=None,
+        description="The ID of the fine-tuned model on the provider's side. May not be the same as the provider_id.",
+    )
+    dataset_split_id: str = Field(
+        description="The ID of the dataset split to use for this fine-tune.",
+    )
+    train_split_name: str = Field(
+        default="train",
+        description="The name of the training split to use for this fine-tune.",
+    )
+    validation_split_name: str | None = Field(
+        default=None,
+        description="The name of the validation split to use for this fine-tune. Optional.",
+    )
+    parameters: dict[str, str | int | float | bool] = Field(
+        default={},
+        description="The parameters to use for this fine-tune. These are provider-specific.",
+    )
+    system_message: str = Field(
+        description="The system message to use for this fine-tune.",
+    )
+    latest_status: FineTuneStatusType = Field(
+        default=FineTuneStatusType.unknown,
+        description="The latest known status of this fine-tune. Not updated in real time.",
+    )
+    properties: Dict[str, str | int | float] = Field(
+        default={},
+        description="Properties of the fine-tune. Different providers may use different properties.",
+    )
+
+    def parent_task(self) -> Task | None:
+        if not isinstance(self.parent, Task):
+            return None
+        return self.parent
+
+
 class DataSourceType(str, Enum):
     """
     The source type of a piece of data.
@@ -344,6 +393,160 @@ class TaskRun(KilnParentedModel):
         return self
 
 
+# Define the type alias for clarity
+DatasetFilter = Callable[[TaskRun], bool]
+
+
+def AllDatasetFilter(_: TaskRun) -> bool:
+    return True
+
+
+def HighRatingDatasetFilter(task_run: TaskRun) -> bool:
+    if task_run.output is None or task_run.output.rating is None:
+        return False
+    return task_run.output.rating.is_high_quality()
+
+
+class DatasetSplitDefinition(BaseModel):
+    """
+    A definition of a split in a dataset.
+
+    Example: name="train", description="The training set", percentage=0.8 (80% of the dataset)
+    """
+
+    name: str = NAME_FIELD
+    description: str | None = Field(
+        default=None,
+        description="A description of the dataset for you and your team. Not used in training.",
+    )
+    percentage: float = Field(
+        ge=0.0,
+        le=1.0,
+        description="The percentage of the dataset that this split represents (between 0 and 1).",
+    )
+
+
+AllSplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="all", percentage=1.0)
+]
+Train80Test20SplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="train", percentage=0.8),
+    DatasetSplitDefinition(name="test", percentage=0.2),
+]
+Train60Test20Val20SplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="train", percentage=0.6),
+    DatasetSplitDefinition(name="test", percentage=0.2),
+    DatasetSplitDefinition(name="val", percentage=0.2),
+]
+
+
+class DatasetSplit(KilnParentedModel):
+    """
+    A collection of task runs, with optional splits (train, test, validation).
+
+    Used to freeze a dataset into train/test/validation splits for repeatable fine-tuning or other tasks.
+
+    Maintains a list of IDs for each split, to avoid data duplication.
+    """
+
+    name: str = NAME_FIELD
+    description: str | None = Field(
+        default=None,
+        description="A description of the dataset for you and your team. Not used in training.",
+    )
+    splits: list[DatasetSplitDefinition] = Field(
+        default_factory=list,
+        description="The splits in the dataset.",
+    )
+    split_contents: dict[str, list[str]] = Field(
+        description="The contents of each split in the dataset. The key is the split name, and the value is a list of task run IDs.",
+    )
+
+    @model_validator(mode="after")
+    def validate_split_percentages(self) -> "DatasetSplit":
+        total = sum(split.percentage for split in self.splits)
+        if not math.isclose(total, 1.0, rel_tol=1e-9):
+            raise ValueError(f"The sum of split percentages must be 1.0 (got {total})")
+        return self
+
+    @classmethod
+    def from_task(
+        cls,
+        name: str,
+        task: "Task",
+        splits: list[DatasetSplitDefinition],
+        filter: DatasetFilter = AllDatasetFilter,
+        description: str | None = None,
+    ):
+        """
+        Build a dataset split from a task.
+        """
+        split_contents = cls.build_split_contents(task, splits, filter)
+        return cls(
+            parent=task,
+            name=name,
+            description=description,
+            splits=splits,
+            split_contents=split_contents,
+        )
+
+    @classmethod
+    def build_split_contents(
+        cls,
+        task: "Task",
+        splits: list[DatasetSplitDefinition],
+        filter: DatasetFilter,
+    ) -> dict[str, list[str]]:
+        valid_ids = []
+        for task_run in task.runs():
+            if filter(task_run):
+                valid_ids.append(task_run.id)
+
+        # Shuffle and split by split percentage
+        random.shuffle(valid_ids)
+        split_contents = {}
+        start_idx = 0
+        remaining_items = len(valid_ids)
+
+        # Handle all splits except the last one
+        for split in splits[:-1]:
+            split_size = round(len(valid_ids) * split.percentage)
+            split_contents[split.name] = valid_ids[start_idx : start_idx + split_size]
+            start_idx += split_size
+            remaining_items -= split_size
+
+        # Last split gets all remaining items (for rounding)
+        if splits:
+            split_contents[splits[-1].name] = valid_ids[start_idx:]
+
+        return split_contents
+
+    def parent_task(self) -> "Task | None":
+        # inline import to avoid circular import
+        from kiln_ai.datamodel import Task
+
+        if not isinstance(self.parent, Task):
+            return None
+        return self.parent
+
+    def missing_count(self) -> int:
+        """
+        Returns:
+            int: the number of task runs that have an ID persisted in this dataset split, but no longer exist in the dataset
+        """
+        parent = self.parent_task()
+        if parent is None:
+            raise ValueError("DatasetSplit has no parent task")
+
+        runs = parent.runs()
+        all_ids = set(run.id for run in runs)
+        all_ids_in_splits = set()
+        for ids in self.split_contents.values():
+            all_ids_in_splits.update(ids)
+        missing = all_ids_in_splits - all_ids
+        return len(missing)
+
+
 class TaskRequirement(BaseModel):
     """
     Defines a specific requirement that should be met by task outputs.
@@ -376,7 +579,11 @@ class TaskDeterminism(str, Enum):
 class Task(
     KilnParentedModel,
     KilnParentModel,
-    parent_of={"runs": TaskRun},
+    parent_of={
+        "runs": TaskRun,
+        "dataset_splits": DatasetSplit,
+        "finetunes": Finetune,
+    },
 ):
     """
     Represents a specific task to be performed, with associated requirements and validation rules.
@@ -416,6 +623,12 @@ class Task(
     def runs(self) -> list[TaskRun]:
         return super().runs()  # type: ignore
 
+    def dataset_splits(self) -> list[DatasetSplit]:
+        return super().dataset_splits()  # type: ignore
+
+    def finetunes(self) -> list[Finetune]:
+        return super().finetunes()  # type: ignore
+
 
 class Project(KilnParentModel, parent_of={"tasks": Task}):
     """

kiln_ai/datamodel/basemodel.py

@@ -1,4 +1,5 @@
 import json
+import re
 import shutil
 import uuid
 from abc import ABCMeta
@@ -38,6 +39,34 @@ ID_TYPE = Optional[str]
 T = TypeVar("T", bound="KilnBaseModel")
 PT = TypeVar("PT", bound="KilnParentedModel")
 
+# Naming 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.
+
+# Filename compatible names
+NAME_REGEX = r"^[A-Za-z0-9 _-]+$"
+NAME_FIELD = Field(
+    min_length=1,
+    max_length=120,
+    pattern=NAME_REGEX,
+    description="A name for this entity.",
+)
+SHORT_NAME_FIELD = Field(
+    min_length=1,
+    max_length=32,
+    pattern=NAME_REGEX,
+    description="A name for this entity",
+)
+
+
+def string_to_valid_name(name: str) -> str:
+    # Replace any character not allowed by NAME_REGEX with an underscore
+    valid_name = re.sub(r"[^A-Za-z0-9 _-]", "_", name)
+    # Replace consecutive underscores with a single underscore
+    valid_name = re.sub(r"_+", "_", valid_name)
+    # Remove leading and trailing underscores or whitespace
+    return valid_name.strip("_").strip()
+
 
 class KilnBaseModel(BaseModel):
     """Base model for all Kiln data models with common functionality for persistence and versioning.
@@ -97,6 +126,7 @@ class KilnBaseModel(BaseModel):
 
         Raises:
             ValueError: If the loaded model is not of the expected type or version
+            FileNotFoundError: If the file does not exist
         """
         with open(path, "r") as file:
             file_data = file.read()

kiln_ai/datamodel/registry.py

@@ -0,0 +1,31 @@
+from kiln_ai.datamodel import Project
+from kiln_ai.utils.config import Config
+
+
+def all_projects() -> list[Project]:
+    project_paths = Config.shared().projects
+    if project_paths is None:
+        return []
+    projects = []
+    for project_path in project_paths:
+        try:
+            projects.append(Project.load_from_file(project_path))
+        except Exception:
+            # deleted files are possible continue with the rest
+            continue
+    return projects
+
+
+def project_from_id(project_id: str) -> Project | None:
+    project_paths = Config.shared().projects
+    if project_paths is not None:
+        for project_path in project_paths:
+            try:
+                project = Project.load_from_file(project_path)
+                if project.id == project_id:
+                    return project
+            except Exception:
+                # deleted files are possible continue with the rest
+                continue
+
+    return None

kiln_ai/datamodel/test_basemodel.py

@@ -5,7 +5,11 @@ from typing import Optional
 
 import pytest
 
-from kiln_ai.datamodel.basemodel import KilnBaseModel, KilnParentedModel
+from kiln_ai.datamodel.basemodel import (
+    KilnBaseModel,
+    KilnParentedModel,
+    string_to_valid_name,
+)
 
 
 @pytest.fixture
@@ -306,3 +310,27 @@ def test_delete_no_path():
     model = KilnBaseModel()
     with pytest.raises(ValueError, match="Cannot delete model because path is not set"):
         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"
+
+    # 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"
+
+    # 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"
+
+    # Test empty string and whitespace
+    assert string_to_valid_name("") == ""
+    assert string_to_valid_name("   ") == ""

kiln_ai/datamodel/test_dataset_split.py

@@ -0,0 +1,234 @@
+import pytest
+from pydantic import ValidationError
+
+# import datamodel first or we get circular import errors
+from kiln_ai.datamodel import (
+    AllDatasetFilter,
+    AllSplitDefinition,
+    DatasetSplit,
+    DatasetSplitDefinition,
+    DataSource,
+    DataSourceType,
+    HighRatingDatasetFilter,
+    Task,
+    TaskOutput,
+    TaskOutputRating,
+    TaskOutputRatingType,
+    TaskRun,
+    Train60Test20Val20SplitDefinition,
+    Train80Test20SplitDefinition,
+)
+
+
+@pytest.fixture
+def sample_task(tmp_path):
+    task_path = tmp_path / "task.kiln"
+    task = Task(
+        name="Test Task",
+        path=task_path,
+        description="Test task for dataset splitting",
+        instruction="Test instruction",
+    )
+    task.save_to_file()
+    return task
+
+
+@pytest.fixture
+def sample_task_runs(sample_task):
+    # Create 10 task runs with different ratings
+    task_runs = []
+    for i in range(10):
+        rating = 5 if i < 6 else 1  # 6 high, 4 low ratings
+        task_run = TaskRun(
+            parent=sample_task,
+            input=f"input_{i}",
+            input_source=DataSource(
+                type=DataSourceType.human,
+                properties={"created_by": "test-user"},
+            ),
+            output=TaskOutput(
+                output=f"output_{i}",
+                source=DataSource(
+                    type=DataSourceType.human,
+                    properties={"created_by": "test-user"},
+                ),
+                rating=TaskOutputRating(
+                    value=rating, type=TaskOutputRatingType.five_star
+                ),
+            ),
+        )
+        task_run.save_to_file()
+        task_runs.append(task_run)
+    return task_runs
+
+
+@pytest.fixture
+def standard_splitstandard_splitss():
+    return [
+        DatasetSplitDefinition(name="train", percentage=0.8),
+        DatasetSplitDefinition(name="test", percentage=0.2),
+    ]
+
+
+@pytest.fixture
+def task_run():
+    return TaskRun(
+        input="test input",
+        input_source=DataSource(
+            type=DataSourceType.human,
+            properties={"created_by": "test-user"},
+        ),
+        output=TaskOutput(
+            output="test output",
+            source=DataSource(
+                type=DataSourceType.human,
+                properties={"created_by": "test-user"},
+            ),
+            rating=TaskOutputRating(rating=5, type=TaskOutputRatingType.five_star),
+        ),
+    )
+
+
+def test_dataset_split_definition():
+    split = DatasetSplitDefinition(name="train", percentage=0.8)
+    assert split.name == "train"
+    assert split.percentage == 0.8
+    assert split.description is None
+
+    # Test validation
+    with pytest.raises(ValidationError):
+        DatasetSplitDefinition(name="train", percentage=1.5)
+
+
+def test_dataset_split_validation():
+    # Test valid percentages
+    splits = [
+        DatasetSplitDefinition(name="train", percentage=0.8),
+        DatasetSplitDefinition(name="test", percentage=0.2),
+    ]
+    dataset = DatasetSplit(
+        name="test_split",
+        splits=splits,
+        split_contents={"train": [], "test": []},
+    )
+    assert dataset.splits == splits
+
+    # Test invalid percentages
+    invalid_splits = [
+        DatasetSplitDefinition(name="train", percentage=0.8),
+        DatasetSplitDefinition(name="test", percentage=0.3),
+    ]
+    with pytest.raises(ValueError, match="sum of split percentages must be 1.0"):
+        DatasetSplit(
+            name="test_split",
+            splits=invalid_splits,
+            split_contents={"train": [], "test": []},
+        )
+
+
+def test_all_dataset_filter(task_run):
+    assert AllDatasetFilter(task_run) is True
+
+
+def test_high_rating_dataset_filter(sample_task_runs):
+    for task_run in sample_task_runs:
+        assert HighRatingDatasetFilter(task_run) is (
+            task_run.output.rating.is_high_quality()
+        )
+
+
+@pytest.mark.parametrize(
+    "splits,expected_sizes",
+    [
+        (Train80Test20SplitDefinition, {"train": 8, "test": 2}),
+        (AllSplitDefinition, {"all": 10}),
+        (Train60Test20Val20SplitDefinition, {"train": 6, "test": 2, "val": 2}),
+        (
+            [
+                DatasetSplitDefinition(name="train", percentage=0.7),
+                DatasetSplitDefinition(name="validation", percentage=0.2),
+                DatasetSplitDefinition(name="test", percentage=0.1),
+            ],
+            {"train": 7, "validation": 2, "test": 1},
+        ),
+    ],
+)
+def test_dataset_split_from_task(sample_task, sample_task_runs, splits, expected_sizes):
+    assert sample_task_runs is not None
+    dataset = DatasetSplit.from_task("Split Name", sample_task, splits)
+    assert dataset.name == "Split Name"
+
+    # Check split sizes match expected
+    for split_name, expected_size in expected_sizes.items():
+        assert len(dataset.split_contents[split_name]) == expected_size
+
+    # Verify total size matches input size
+    total_size = sum(len(ids) for ids in dataset.split_contents.values())
+    assert total_size == len(sample_task_runs)
+
+
+def test_dataset_split_with_high_rating_filter(sample_task, sample_task_runs):
+    assert len(sample_task_runs) == 10
+    dataset = DatasetSplit.from_task(
+        "Split Name",
+        sample_task,
+        Train80Test20SplitDefinition,
+        filter=HighRatingDatasetFilter,
+    )
+
+    # Check that only high-rated task runs are included
+    all_ids = []
+    for ids in dataset.split_contents.values():
+        all_ids.extend(ids)
+    assert len(all_ids) == 6  # We created 6 high-rated task runs
+
+    # Check split proportions
+    train_size = len(dataset.split_contents["train"])
+    test_size = len(dataset.split_contents["test"])
+    assert train_size == 5  # ~80% of 6
+    assert test_size == 1  # ~20% of 6
+
+
+def test_dataset_split_with_single_split(sample_task, sample_task_runs):
+    splits = [DatasetSplitDefinition(name="all", percentage=1.0)]
+    dataset = DatasetSplit.from_task("Split Name", sample_task, splits)
+
+    assert len(dataset.split_contents["all"]) == len(sample_task_runs)
+
+
+def test_missing_count(sample_task, sample_task_runs):
+    assert sample_task_runs is not None
+    # Create a dataset split with all task runs
+    dataset = DatasetSplit.from_task(
+        "Split Name", sample_task, Train80Test20SplitDefinition
+    )
+
+    # Initially there should be no missing runs
+    assert dataset.missing_count() == 0
+
+    # Add some IDs to the split, that aren't on disk
+    dataset.split_contents["test"].append("1")
+    dataset.split_contents["test"].append("2")
+    dataset.split_contents["test"].append("3")
+    # shouldn't happen, but should not double count if it does
+    dataset.split_contents["train"].append("3")
+
+    # Now we should have 3 missing runs
+    assert dataset.missing_count() == 3
+
+
+def test_smaller_sample(sample_task, sample_task_runs):
+    assert sample_task_runs is not None
+    # Create a dataset split with all task runs
+    dataset = DatasetSplit.from_task(
+        "Split Name", sample_task, Train80Test20SplitDefinition
+    )
+
+    # Initially there should be no missing runs
+    assert dataset.missing_count() == 0
+
+    dataset.split_contents["test"].pop()
+    dataset.split_contents["train"].pop()
+
+    # Now we should have 0 missing runs. It's okay that dataset has newer data.
+    assert dataset.missing_count() == 0

kiln_ai/datamodel/test_example_models.py

@@ -5,8 +5,10 @@ import pytest
 from pydantic import ValidationError
 
 from kiln_ai.datamodel import (
+    DatasetSplit,
     DataSource,
     DataSourceType,
+    Finetune,
     Project,
     Task,
     TaskDeterminism,
@@ -97,6 +99,16 @@ def test_task_run_relationship(valid_task_run):
     assert valid_task_run.__class__.parent_type().__name__ == "Task"
 
 
+def test_dataset_split_relationship():
+    assert DatasetSplit.relationship_name() == "dataset_splits"
+    assert DatasetSplit.parent_type().__name__ == "Task"
+
+
+def test_base_finetune_relationship():
+    assert Finetune.relationship_name() == "finetunes"
+    assert Finetune.parent_type().__name__ == "Task"
+
+
 def test_structured_output_workflow(tmp_path):
     tmp_project_file = (
         tmp_path / "test_structured_output_runs" / Project.base_filename()