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

kiln_ai/datamodel/test_model_cache.py

@@ -0,0 +1,244 @@
+from pathlib import Path
+from unittest import mock
+
+import pytest
+from pydantic import BaseModel
+
+from libs.core.kiln_ai.datamodel.model_cache import ModelCache
+
+
+# Define a simple Pydantic model for testing
+class ModelTest(BaseModel):
+    name: str
+    value: int
+
+
+@pytest.fixture
+def model_cache():
+    return ModelCache()
+
+
+def should_skip_test(model_cache):
+    return not model_cache._enabled
+
+
+@pytest.fixture
+def test_path(tmp_path):
+    # Create a temporary file path for testing
+    test_file = tmp_path / "test_model.kiln"
+    test_file.touch()  # Create the file
+    return test_file
+
+
+def test_set_and_get_model(model_cache, test_path):
+    if not model_cache._enabled:
+        pytest.skip("Cache is disabled on this fs")
+
+    model = ModelTest(name="test", value=123)
+    mtime_ns = test_path.stat().st_mtime_ns
+
+    model_cache.set_model(test_path, model, mtime_ns)
+    cached_model = model_cache.get_model(test_path, ModelTest)
+
+    assert cached_model is not None
+    assert cached_model.name == "test"
+    assert cached_model.value == 123
+
+
+def test_invalidate_model(model_cache, test_path):
+    model = ModelTest(name="test", value=123)
+    mtime = test_path.stat().st_mtime
+
+    model_cache.set_model(test_path, model, mtime)
+    model_cache.invalidate(test_path)
+    cached_model = model_cache.get_model(test_path, ModelTest)
+
+    assert cached_model is None
+
+
+def test_clear_cache(model_cache, test_path):
+    model = ModelTest(name="test", value=123)
+    mtime = test_path.stat().st_mtime
+
+    model_cache.set_model(test_path, model, mtime)
+    model_cache.clear()
+    cached_model = model_cache.get_model(test_path, ModelTest)
+
+    assert cached_model is None
+
+
+def test_cache_invalid_due_to_mtime_change(model_cache, test_path):
+    model = ModelTest(name="test", value=123)
+    mtime = test_path.stat().st_mtime
+
+    model_cache.set_model(test_path, model, mtime)
+
+    # Simulate a file modification by updating the mtime
+    test_path.touch()
+    cached_model = model_cache.get_model(test_path, ModelTest)
+
+    assert cached_model is None
+
+
+def test_get_model_wrong_type(model_cache, test_path):
+    if not model_cache._enabled:
+        pytest.skip("Cache is disabled on this fs")
+
+    class AnotherModel(BaseModel):
+        other_field: str
+
+    model = ModelTest(name="test", value=123)
+    mtime_ns = test_path.stat().st_mtime_ns
+
+    model_cache.set_model(test_path, model, mtime_ns)
+
+    with pytest.raises(ValueError):
+        model_cache.get_model(test_path, AnotherModel)
+
+    # Test that the cache invalidates
+    cached_model = model_cache.get_model(test_path, ModelTest)
+    assert cached_model is None
+
+
+def test_is_cache_valid_true(model_cache, test_path):
+    mtime_ns = test_path.stat().st_mtime_ns
+    assert model_cache._is_cache_valid(test_path, mtime_ns) is True
+
+
+def test_is_cache_valid_false_due_to_mtime_change(model_cache, test_path):
+    if not model_cache._enabled:
+        pytest.skip("Cache is disabled on this fs")
+
+    mtime_ns = test_path.stat().st_mtime_ns
+    # Simulate a file modification by updating the mtime
+    test_path.touch()
+    assert model_cache._is_cache_valid(test_path, mtime_ns) is False
+
+
+def test_is_cache_valid_false_due_to_missing_file(model_cache):
+    non_existent_path = Path("/non/existent/path")
+    assert model_cache._is_cache_valid(non_existent_path, 0) is False
+
+
+def test_benchmark_get_model(benchmark, model_cache, test_path):
+    model = ModelTest(name="test", value=123)
+    mtime = test_path.stat().st_mtime
+
+    # Set the model in the cache
+    model_cache.set_model(test_path, model, mtime)
+
+    # Benchmark the get_model method
+    def get_model():
+        return model_cache.get_model(test_path, ModelTest)
+
+    benchmark(get_model)
+    stats = benchmark.stats.stats
+
+    # 25k ops per second is the target. Getting 250k on Macbook, but CI will be slower
+    target = 1 / 25000
+    if stats.mean > target:
+        pytest.fail(
+            f"Average time per iteration: {stats.mean}, expected less than {target}"
+        )
+
+
+def test_get_model_returns_copy(model_cache, test_path):
+    if not model_cache._enabled:
+        pytest.skip("Cache is disabled on this fs")
+
+    model = ModelTest(name="test", value=123)
+    mtime_ns = test_path.stat().st_mtime_ns
+
+    # Set the model in the cache
+    model_cache.set_model(test_path, model, mtime_ns)
+
+    # Get a copy of the model from the cache
+    cached_model = model_cache.get_model(test_path, ModelTest)
+
+    # Different instance (is), same data (==)
+    assert cached_model is not model
+    assert cached_model == model
+
+    # Mutate the cached model
+    cached_model.name = "mutated"
+
+    # Get the model again from the cache
+    new_cached_model = model_cache.get_model(test_path, ModelTest)
+
+    # Assert that the new cached model has the original values
+    assert new_cached_model == model
+    assert new_cached_model.name == "test"
+
+    # Save the mutated model back to the cache
+    model_cache.set_model(test_path, cached_model, mtime_ns)
+
+    # Get the model again from the cache
+    updated_cached_model = model_cache.get_model(test_path, ModelTest)
+
+    # Assert that the updated cached model has the mutated values
+    assert updated_cached_model.name == "mutated"
+    assert updated_cached_model.value == 123
+
+
+def test_no_cache_when_no_fine_granularity(model_cache, test_path):
+    model = ModelTest(name="test", value=123)
+    mtime_ns = test_path.stat().st_mtime_ns
+
+    model_cache._enabled = False
+    model_cache.set_model(test_path, model, mtime_ns)
+    cached_model = model_cache.get_model(test_path, ModelTest)
+
+    # Assert that the model is not cached
+    assert cached_model is None
+    assert model_cache.model_cache == {}
+    assert model_cache._enabled is False
+
+
+def test_check_timestamp_granularity_macos():
+    with mock.patch("sys.platform", "darwin"):
+        cache = ModelCache()
+        assert cache._check_timestamp_granularity() is True
+        assert cache._enabled is True
+
+
+def test_check_timestamp_granularity_windows():
+    with mock.patch("sys.platform", "win32"):
+        cache = ModelCache()
+        assert cache._check_timestamp_granularity() is True
+        assert cache._enabled is True
+
+
+def test_check_timestamp_granularity_linux_good():
+    mock_stats = mock.Mock()
+    mock_stats.f_timespec = 9  # nanosecond precision
+
+    with (
+        mock.patch("sys.platform", "linux"),
+        mock.patch("os.statvfs", return_value=mock_stats),
+    ):
+        cache = ModelCache()
+        assert cache._check_timestamp_granularity() is True
+        assert cache._enabled is True
+
+
+def test_check_timestamp_granularity_linux_poor():
+    mock_stats = mock.Mock()
+    mock_stats.f_timespec = 3  # millisecond precision
+
+    with (
+        mock.patch("sys.platform", "linux"),
+        mock.patch("os.statvfs", return_value=mock_stats),
+    ):
+        cache = ModelCache()
+        assert cache._check_timestamp_granularity() is False
+        assert cache._enabled is False
+
+
+def test_check_timestamp_granularity_linux_error():
+    with (
+        mock.patch("sys.platform", "linux"),
+        mock.patch("os.statvfs", side_effect=OSError("Mock filesystem error")),
+    ):
+        cache = ModelCache()
+        assert cache._check_timestamp_granularity() is False
+        assert cache._enabled is False

kiln_ai/datamodel/test_models.py

@@ -1,4 +1,6 @@
 import json
+import os
+from unittest.mock import patch
 
 import pytest
 from pydantic import ValidationError
@@ -6,6 +8,7 @@ from pydantic import ValidationError
 from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
+    Finetune,
     Project,
     Task,
     TaskOutput,
@@ -27,7 +30,7 @@ def test_project_file(tmp_path):
 
 @pytest.fixture
 def test_task_file(tmp_path):
-    test_file_path = tmp_path / "task.json"
+    test_file_path = tmp_path / "task.kiln"
     data = {
         "v": 1,
         "name": "Test Task",
@@ -81,6 +84,7 @@ def test_task_serialization(test_project_file):
         instruction="Test Base Task Instruction",
         thinking_instruction="Test Thinking Instruction",
     )
+    assert task._loaded_from_file is False
 
     task.save_to_file()
 
@@ -89,6 +93,11 @@ def test_task_serialization(test_project_file):
     assert parsed_task.description == "Test Description"
     assert parsed_task.instruction == "Test Base Task Instruction"
     assert parsed_task.thinking_instruction == "Test Thinking Instruction"
+    assert parsed_task._loaded_from_file is True
+
+    # Confirm the local property is not persisted to disk
+    json_data = json.loads(parsed_task.path.read_text())
+    assert "_loaded_from_file" not in json_data
 
 
 def test_save_to_file_without_path():
@@ -225,3 +234,208 @@ def test_task_run_intermediate_outputs():
         "cot": "chain of thought output",
         "draft": "draft output",
     }
+
+
+def test_finetune_basic():
+    # Test basic initialization
+    finetune = Finetune(
+        name="test-finetune",
+        provider="openai",
+        base_model_id="gpt-3.5-turbo",
+        dataset_split_id="dataset-123",
+        train_split_name="train",
+        system_message="Test system message",
+    )
+    assert finetune.name == "test-finetune"
+    assert finetune.provider == "openai"
+    assert finetune.base_model_id == "gpt-3.5-turbo"
+    assert finetune.dataset_split_id == "dataset-123"
+    assert finetune.train_split_name == "train"
+    assert finetune.provider_id is None
+    assert finetune.parameters == {}
+    assert finetune.description is None
+
+
+def test_finetune_full():
+    # Test with all fields populated
+    finetune = Finetune(
+        name="test-finetune",
+        description="Test description",
+        provider="openai",
+        base_model_id="gpt-3.5-turbo",
+        provider_id="ft-abc123",
+        dataset_split_id="dataset-123",
+        train_split_name="train",
+        system_message="Test system message",
+        parameters={
+            "epochs": 3,
+            "learning_rate": 0.1,
+            "batch_size": 4,
+            "use_fp16": True,
+            "model_suffix": "-v1",
+        },
+    )
+    assert finetune.description == "Test description"
+    assert finetune.provider_id == "ft-abc123"
+    assert finetune.parameters == {
+        "epochs": 3,
+        "learning_rate": 0.1,
+        "batch_size": 4,
+        "use_fp16": True,
+        "model_suffix": "-v1",
+    }
+    assert finetune.system_message == "Test system message"
+
+
+def test_finetune_parent_task():
+    # Test parent_task() method
+    task = Task(name="Test Task", instruction="Test instruction")
+    finetune = Finetune(
+        name="test-finetune",
+        provider="openai",
+        base_model_id="gpt-3.5-turbo",
+        parent=task,
+        dataset_split_id="dataset-123",
+        train_split_name="train",
+        system_message="Test system message",
+    )
+
+    assert finetune.parent_task() == task
+
+    # Test with no parent
+    finetune_no_parent = Finetune(
+        name="test-finetune",
+        provider="openai",
+        base_model_id="gpt-3.5-turbo",
+        dataset_split_id="dataset-123",
+        train_split_name="train",
+        system_message="Test system message",
+    )
+    assert finetune_no_parent.parent_task() is None
+
+
+def test_finetune_parameters_validation():
+    # Test that parameters only accept valid types
+    with pytest.raises(ValidationError):
+        Finetune(
+            name="test-finetune",
+            provider="openai",
+            base_model_id="gpt-3.5-turbo",
+            parameters={"invalid": [1, 2, 3]},  # Lists are not allowed
+        )
+
+
+def test_task_run_input_source_validation(tmp_path):
+    # Setup basic output for TaskRun creation
+    output = TaskOutput(
+        output="test output",
+        source=DataSource(
+            type=DataSourceType.synthetic,
+            properties={
+                "model_name": "test-model",
+                "model_provider": "test-provider",
+                "adapter_name": "test-adapter",
+            },
+        ),
+    )
+
+    project_path = tmp_path / "project.kiln"
+    project = Project(name="Test Project", path=project_path)
+    project.save_to_file()
+    task = Task(name="Test Task", instruction="Test Instruction", parent=project)
+    task.save_to_file()
+
+    # Test 1: Creating without input_source should work when strict mode is off
+    task_run = TaskRun(
+        input="test input",
+        output=output,
+    )
+    task_run.parent = task
+    assert task_run.input_source is None
+
+    # Save for later usage
+    task_run.save_to_file()
+    task_missing_input_source = task_run.path
+
+    # Test 2: Creating with input_source should work when strict mode is off
+    task_run = TaskRun(
+        input="test input 2",
+        input_source=DataSource(
+            type=DataSourceType.human,
+            properties={"created_by": "test-user"},
+        ),
+        output=output,
+    )
+    assert task_run.input_source is not None
+
+    # Test 3: Creating without input_source should fail when strict mode is on
+    with patch("kiln_ai.datamodel.strict_mode", return_value=True):
+        with pytest.raises(ValueError) as exc_info:
+            task_run = TaskRun(
+                input="test input 3",
+                output=output,
+            )
+        assert "input_source is required when strict mode is enabled" in str(
+            exc_info.value
+        )
+
+        # Test 4: Loading from disk should work without input_source, even with strict mode on
+        assert os.path.exists(task_missing_input_source)
+        task_run = TaskRun.load_from_file(task_missing_input_source)
+        assert task_run.input_source is None
+
+
+def test_task_output_source_validation(tmp_path):
+    # Setup basic output source for validation
+    output_source = DataSource(
+        type=DataSourceType.synthetic,
+        properties={
+            "model_name": "test-model",
+            "model_provider": "test-provider",
+            "adapter_name": "test-adapter",
+        },
+    )
+
+    project_path = tmp_path / "project.kiln"
+    project = Project(name="Test Project", path=project_path)
+    project.save_to_file()
+    task = Task(name="Test Task", instruction="Test Instruction", parent=project)
+    task.save_to_file()
+
+    # Test 1: Creating without source should work when strict mode is off
+    task_output = TaskOutput(
+        output="test output",
+    )
+    assert task_output.source is None
+
+    # Save for later usage
+    task_run = TaskRun(
+        input="test input",
+        input_source=output_source,
+        output=task_output,
+    )
+    task_run.parent = task
+    task_run.save_to_file()
+    task_missing_output_source = task_run.path
+
+    # Test 2: Creating with source should work when strict mode is off
+    task_output = TaskOutput(
+        output="test output 2",
+        source=output_source,
+    )
+    assert task_output.source is not None
+
+    # Test 3: Creating without source should fail when strict mode is on
+    with patch("kiln_ai.datamodel.strict_mode", return_value=True):
+        with pytest.raises(ValueError) as exc_info:
+            task_output = TaskOutput(
+                output="test output 3",
+            )
+        assert "Output source is required when strict mode is enabled" in str(
+            exc_info.value
+        )
+
+        # Test 4: Loading from disk should work without source, even with strict mode on
+        assert os.path.exists(task_missing_output_source)
+        task_run = TaskRun.load_from_file(task_missing_output_source)
+        assert task_run.output.source is None

kiln_ai/datamodel/test_registry.py

@@ -0,0 +1,96 @@
+from unittest.mock import Mock, patch
+
+import pytest
+
+from kiln_ai.datamodel import Project
+from kiln_ai.datamodel.registry import all_projects, project_from_id
+
+
+@pytest.fixture
+def mock_config():
+    with patch("kiln_ai.datamodel.registry.Config") as mock:
+        config_instance = Mock()
+        mock.shared.return_value = config_instance
+        yield config_instance
+
+
+@pytest.fixture
+def mock_project():
+    def create_mock_project(project_id: str = "test-id"):
+        project = Mock(spec=Project)
+        project.id = project_id
+        return project
+
+    return create_mock_project
+
+
+def test_all_projects_empty(mock_config):
+    mock_config.projects = None
+    assert all_projects() == []
+
+
+def test_all_projects_success(mock_config, mock_project):
+    mock_config.projects = ["path1", "path2"]
+
+    project1 = mock_project("project1")
+    project2 = mock_project("project2")
+
+    with patch("kiln_ai.datamodel.Project.load_from_file") as mock_load:
+        mock_load.side_effect = [project1, project2]
+
+        result = all_projects()
+
+        assert len(result) == 2
+        assert result[0] == project1
+        assert result[1] == project2
+        mock_load.assert_any_call("path1")
+        mock_load.assert_any_call("path2")
+
+
+def test_all_projects_with_errors(mock_config, mock_project):
+    mock_config.projects = ["path1", "path2", "path3"]
+
+    project1 = mock_project("project1")
+    project3 = mock_project("project3")
+
+    with patch("kiln_ai.datamodel.Project.load_from_file") as mock_load:
+        mock_load.side_effect = [project1, Exception("File not found"), project3]
+
+        result = all_projects()
+
+        assert len(result) == 2
+        assert result[0] == project1
+        assert result[1] == project3
+
+
+def test_project_from_id_not_found(mock_config):
+    mock_config.projects = None
+    assert project_from_id("any-id") is None
+
+
+def test_project_from_id_success(mock_config, mock_project):
+    mock_config.projects = ["path1", "path2"]
+
+    project1 = mock_project("project1")
+    project2 = mock_project("project2")
+
+    with patch("kiln_ai.datamodel.Project.load_from_file") as mock_load:
+        mock_load.side_effect = [project1, project2]
+
+        result = project_from_id("project2")
+
+        assert result == project2
+
+
+def test_project_from_id_with_errors(mock_config, mock_project):
+    mock_config.projects = ["path1", "path2", "path3"]
+
+    project1 = mock_project("project1")
+    project3 = mock_project("target-id")
+
+    with patch("kiln_ai.datamodel.Project.load_from_file") as mock_load:
+        mock_load.side_effect = [project1, Exception("File not found"), project3]
+
+        result = project_from_id("target-id")
+
+        assert result == project3

kiln_ai/utils/config.py

@@ -67,10 +67,23 @@ class Config:
                 env_var="OPENROUTER_API_KEY",
                 sensitive=True,
             ),
+            "fireworks_api_key": ConfigProperty(
+                str,
+                env_var="FIREWORKS_API_KEY",
+                sensitive=True,
+            ),
+            "fireworks_account_id": ConfigProperty(
+                str,
+                env_var="FIREWORKS_ACCOUNT_ID",
+            ),
             "projects": ConfigProperty(
                 list,
                 default_lambda=lambda: [],
             ),
+            "custom_models": ConfigProperty(
+                list,
+                default_lambda=lambda: [],
+            ),
         }
         self._settings = self.load_settings()
 
@@ -136,7 +149,7 @@ class Config:
             settings = yaml.safe_load(f.read()) or {}
         return settings
 
-    def settings(self, hide_sensitive=False):
+    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
         if hide_sensitive:
             return {
                 k: "[hidden]"