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

kiln_ai/datamodel/model_cache.py

@@ -0,0 +1,116 @@
+"""
+A simple cache for our datamodel.
+
+Works at the file level, caching the pydantic model based on the file path.
+
+Keeping this really simple. Our goal is to really be "disk-backed" data model, so using disk primitives.
+
+ - Use disk mtime to determine if the cached model is stale.
+ - Still using glob for iterating over projects, just caching at the file level
+ - Use path as the cache key
+ - Cache always populated from a disk read, so we know it refects what's on disk. Even if we had a memory-constructed version, we don't cache that.
+ - Cache the parsed model, not the raw file contents. Parsing and validating is what's expensive. >99% speedup when measured.
+"""
+
+import os
+import sys
+import warnings
+from pathlib import Path
+from typing import Dict, Optional, Tuple, Type, TypeVar
+
+from pydantic import BaseModel
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class ModelCache:
+    _shared_instance = None
+
+    def __init__(self):
+        # Store both the model and the modified time of the cached file contents
+        self.model_cache: Dict[Path, Tuple[BaseModel, int]] = {}
+        self._enabled = self._check_timestamp_granularity()
+        if not self._enabled:
+            warnings.warn(
+                "File system does not support fine-grained timestamps. "
+                "Model caching has been disabled to ensure consistency."
+            )
+
+    @classmethod
+    def shared(cls):
+        if cls._shared_instance is None:
+            cls._shared_instance = cls()
+        return cls._shared_instance
+
+    def _is_cache_valid(self, path: Path, cached_mtime_ns: int) -> bool:
+        try:
+            current_mtime_ns = path.stat().st_mtime_ns
+        except Exception:
+            return False
+        return cached_mtime_ns == current_mtime_ns
+
+    def _get_model(self, path: Path, model_type: Type[T]) -> Optional[T]:
+        if path not in self.model_cache:
+            return None
+        model, cached_mtime_ns = self.model_cache[path]
+        if not self._is_cache_valid(path, cached_mtime_ns):
+            self.invalidate(path)
+            return None
+
+        if not isinstance(model, model_type):
+            self.invalidate(path)
+            raise ValueError(f"Model at {path} is not of type {model_type.__name__}")
+        return model
+
+    def get_model(self, path: Path, model_type: Type[T]) -> Optional[T]:
+        # We return a copy so in-memory edits don't impact the cache until they are saved
+        # Benchmark shows about 2x slower, but much more foolproof
+        model = self._get_model(path, model_type)
+        if model:
+            return model.model_copy(deep=True)
+        return None
+
+    def get_model_id(self, path: Path, model_type: Type[T]) -> Optional[str]:
+        model = self._get_model(path, model_type)
+        if model and hasattr(model, "id"):
+            id = model.id  # type: ignore
+            if isinstance(id, str):
+                return id
+        return None
+
+    def set_model(self, path: Path, model: BaseModel, mtime_ns: int):
+        # disable caching if the filesystem doesn't support fine-grained timestamps
+        if not self._enabled:
+            return
+        self.model_cache[path] = (model, mtime_ns)
+
+    def invalidate(self, path: Path):
+        if path in self.model_cache:
+            del self.model_cache[path]
+
+    def clear(self):
+        self.model_cache.clear()
+
+    def _check_timestamp_granularity(self) -> bool:
+        """Check if filesystem supports fine-grained timestamps (microseconds or better)."""
+
+        # MacOS and Windows support fine-grained timestamps
+        if sys.platform in ["darwin", "win32"]:
+            return True
+
+        # Linux supports fine-grained timestamps SOMETIMES. ext4 should work.
+        try:
+            # Get filesystem stats for the current directory
+            stats = os.statvfs(Path(__file__).parent)
+
+            # f_timespec was added in Linux 5.6 (2020)
+            # Returns nanoseconds precision as a power of 10
+            # e.g., 1 = decisecond, 2 = centisecond, 3 = millisecond, etc.
+            timespec = getattr(stats, "f_timespec", 0)
+
+            # Consider microsecond precision (6) or better as "fine-grained"
+            return timespec >= 6
+        except (AttributeError, OSError):
+            # If f_timespec isn't available or other errors occur,
+            # assume poor granularity to be safe
+            return False

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

@@ -2,10 +2,16 @@ import datetime
 import json
 from pathlib import Path
 from typing import Optional
+from unittest.mock import MagicMock, patch
 
 import pytest
 
-from kiln_ai.datamodel.basemodel import KilnBaseModel, KilnParentedModel
+from kiln_ai.datamodel.basemodel import (
+    KilnBaseModel,
+    KilnParentedModel,
+    string_to_valid_name,
+)
+from kiln_ai.datamodel.model_cache import ModelCache
 
 
 @pytest.fixture
@@ -41,6 +47,17 @@ def test_newer_file(tmp_path) -> Path:
     return test_file_path
 
 
+@pytest.fixture
+def tmp_model_cache():
+    temp_cache = ModelCache()
+    # We're testing integration, not cache functions, in this file
+    temp_cache._enabled = True
+    with (
+        patch("kiln_ai.datamodel.basemodel.ModelCache.shared", return_value=temp_cache),
+    ):
+        yield temp_cache
+
+
 def test_load_from_file(test_base_file):
     model = KilnBaseModel.load_from_file(test_base_file)
     assert model.v == 1
@@ -273,9 +290,8 @@ def test_lazy_load_parent(tmp_path):
     assert loaded_parent.name == "Parent"
     assert loaded_parent.path == parent.path
 
-    # Verify that the _parent attribute is now set
-    assert hasattr(loaded_child, "_parent")
-    assert loaded_child._parent is loaded_parent
+    # Verify that the parent is cached
+    assert loaded_child.cached_parent() is loaded_parent
 
 
 def test_delete(tmp_path):
@@ -306,3 +322,150 @@ 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("   ") == ""
+
+
+def test_load_from_file_with_cache(test_base_file, tmp_model_cache):
+    tmp_model_cache.get_model = MagicMock(return_value=None)
+    tmp_model_cache.set_model = MagicMock()
+
+    # Load the model
+    model = KilnBaseModel.load_from_file(test_base_file)
+
+    # Check that the cache was checked and set
+    tmp_model_cache.get_model.assert_called_once_with(test_base_file, KilnBaseModel)
+    tmp_model_cache.set_model.assert_called_once()
+
+    # Ensure the model is correctly loaded
+    assert model.v == 1
+    assert model.path == test_base_file
+
+
+def test_save_to_file_invalidates_cache(test_base_file, tmp_model_cache):
+    # Create and save the model
+    model = KilnBaseModel(path=test_base_file)
+
+    # Set mock after to ignore any previous calls, we want to see save calls it
+    tmp_model_cache.invalidate = MagicMock()
+    model.save_to_file()
+
+    # Check that the cache was invalidated. Might be called multiple times for setting props like path. but must be called at least once.
+    tmp_model_cache.invalidate.assert_called_with(test_base_file)
+
+
+def test_delete_invalidates_cache(tmp_path, tmp_model_cache):
+    # Create and save the model
+    file_path = tmp_path / "test.kiln"
+    model = KilnBaseModel(path=file_path)
+    model.save_to_file()
+
+    # populate and check cache
+    model = KilnBaseModel.load_from_file(file_path)
+    cached_model = tmp_model_cache.get_model(file_path, KilnBaseModel)
+    assert cached_model.id == model.id
+
+    tmp_model_cache.invalidate = MagicMock()
+
+    # Delete the model
+    model.delete()
+
+    # Check that the cache was invalidated
+    tmp_model_cache.invalidate.assert_called_with(file_path)
+    assert tmp_model_cache.get_model(file_path, KilnBaseModel) is None
+
+
+def test_load_from_file_with_cached_model(test_base_file, tmp_model_cache):
+    # Set up the mock to return a cached model
+    cached_model = KilnBaseModel(v=1, path=test_base_file)
+    tmp_model_cache.get_model = MagicMock(return_value=cached_model)
+
+    with patch("builtins.open", create=True) as mock_open:
+        # Load the model
+        model = KilnBaseModel.load_from_file(test_base_file)
+
+        # Check that the cache was checked and the cached model was returned
+        tmp_model_cache.get_model.assert_called_once_with(test_base_file, KilnBaseModel)
+        assert model is cached_model
+
+        # Assert that open was not called (we used the cached model, not file)
+        mock_open.assert_not_called()
+
+
+def test_from_id_and_parent_path(test_base_parented_file, tmp_model_cache):
+    # Set up parent and children models
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+
+    child1 = DefaultParentedModel(parent=parent, name="Child1")
+    child2 = DefaultParentedModel(parent=parent, name="Child2")
+    child3 = DefaultParentedModel(parent=parent, name="Child3")
+
+    # Save all children
+    child1.save_to_file()
+    child2.save_to_file()
+    child3.save_to_file()
+
+    # Test finding existing child by ID
+    found_child = DefaultParentedModel.from_id_and_parent_path(
+        child2.id, test_base_parented_file
+    )
+    assert found_child is not None
+    assert found_child.id == child2.id
+    assert found_child.name == "Child2"
+    assert found_child is not child2  # not same instance (deep copy)
+
+    # Test non-existent ID returns None
+    not_found = DefaultParentedModel.from_id_and_parent_path(
+        "nonexistent", test_base_parented_file
+    )
+    assert not_found is None
+
+
+def test_from_id_and_parent_path_with_cache(test_base_parented_file, tmp_model_cache):
+    # Set up parent and child
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+    child = DefaultParentedModel(parent=parent, name="Child")
+    child.save_to_file()
+
+    # First load to populate cache
+    _ = DefaultParentedModel.from_id_and_parent_path(child.id, test_base_parented_file)
+
+    # Mock cache to verify it's used
+    tmp_model_cache.get_model_id = MagicMock(return_value=child.id)
+
+    # Load again - should use cache
+    found_child = DefaultParentedModel.from_id_and_parent_path(
+        child.id, test_base_parented_file
+    )
+
+    assert found_child is not None
+    assert found_child.id == child.id
+    tmp_model_cache.get_model_id.assert_called()
+
+
+def test_from_id_and_parent_path_without_parent():
+    # Test with None parent_path
+    not_found = DefaultParentedModel.from_id_and_parent_path("any-id", None)
+    assert not_found is None

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()