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

kiln_ai/datamodel/test_basemodel.py

@@ -6,12 +6,16 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
+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 (
     KilnBaseModel,
     KilnParentedModel,
     string_to_valid_name,
 )
 from kiln_ai.datamodel.model_cache import ModelCache
+from kiln_ai.datamodel.task import RunConfig
 
 
 @pytest.fixture
@@ -356,7 +360,9 @@ def test_load_from_file_with_cache(test_base_file, tmp_model_cache):
     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.get_model.assert_called_once_with(
+        test_base_file, KilnBaseModel, readonly=False
+    )
     tmp_model_cache.set_model.assert_called_once()
 
     # Ensure the model is correctly loaded
@@ -407,7 +413,9 @@ def test_load_from_file_with_cached_model(test_base_file, tmp_model_cache):
         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)
+        tmp_model_cache.get_model.assert_called_once_with(
+            test_base_file, KilnBaseModel, readonly=False
+        )
         assert model is cached_model
 
         # Assert that open was not called (we used the cached model, not file)
@@ -469,3 +477,73 @@ 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
+
+
+class MockAdapter(BaseAdapter):
+    """Implementation of BaseAdapter for testing"""
+
+    async def _run(self, input):
+        return RunOutput(output="test output", intermediate_outputs=None)
+
+    def adapter_name(self) -> str:
+        return "test"
+
+
+@pytest.fixture
+def base_task():
+    return Task(name="test_task", instruction="test_instruction")
+
+
+@pytest.fixture
+def adapter(base_task):
+    return MockAdapter(
+        run_config=RunConfig(
+            task=base_task,
+            model_name="test_model",
+            model_provider_name="test_provider",
+            prompt_id="simple_prompt_builder",
+        ),
+    )
+
+
+async def test_invoke_parsing_flow(adapter):
+    # Mock dependencies
+    mock_provider = MagicMock()
+    mock_provider.parser = "test_parser"
+
+    mock_parser = MagicMock()
+    mock_parser.parse_output.return_value = RunOutput(
+        output="parsed test output", intermediate_outputs={"key": "value"}
+    )
+
+    mock_parser_class = MagicMock(return_value=mock_parser)
+
+    with (
+        patch.object(adapter, "model_provider", return_value=mock_provider),
+        patch(
+            "kiln_ai.adapters.model_adapters.base_adapter.model_parser_from_id",
+            return_value=mock_parser_class,
+        ),
+        patch("kiln_ai.adapters.model_adapters.base_adapter.Config") as mock_config,
+    ):
+        # Disable autosaving for this test
+        mock_config.shared.return_value.autosave_runs = False
+        mock_config.shared.return_value.user_id = "test_user_id"
+
+        # Execute
+        result = await adapter.invoke("test input")
+
+        # Verify parser was created correctly
+        mock_parser_class.assert_called_once_with(structured_output=False)
+
+        # Verify parsing occurred
+        mock_parser.parse_output.assert_called_once()
+        parsed_args = mock_parser.parse_output.call_args[1]
+        assert isinstance(parsed_args["original_output"], RunOutput)
+        assert parsed_args["original_output"].output == "test output"
+
+        # Verify result contains parsed output
+        assert isinstance(result, TaskRun)
+        assert result.output.output == "parsed test output"
+        assert result.intermediate_outputs == {"key": "value"}
+        assert result.input == "test input"

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,21 +3,28 @@ 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,
+)
+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
@@ -39,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}",
@@ -56,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)
@@ -131,10 +140,33 @@ def test_all_dataset_filter(task_run):
 
 
 def test_high_rating_dataset_filter(sample_task_runs):
+    num_high_quality = 0
+    num_low_quality = 0
     for task_run in sample_task_runs:
-        assert HighRatingDatasetFilter(task_run) is (
-            task_run.output.rating.is_high_quality()
+        if HighRatingDatasetFilter(task_run):
+            num_high_quality += 1
+            assert task_run.output.rating.is_high_quality() is True
+        else:
+            num_low_quality += 1
+            assert task_run.output.rating.is_high_quality() is False
+
+        # Test repaired output always considered high quality
+        task_run = task_run.model_copy(
+            update={
+                "repair_instructions": "repair instructions",
+                "repaired_output": TaskOutput(
+                    output="repaired output",
+                    source=DataSource(
+                        type=DataSourceType.human,
+                        properties={"created_by": "test-user"},
+                    ),
+                ),
+            }
         )
+        assert HighRatingDatasetFilter(task_run) is True
+
+    assert num_high_quality == 6
+    assert num_low_quality == 4
 
 
 @pytest.mark.parametrize(
@@ -173,9 +205,11 @@ def test_dataset_split_with_high_rating_filter(sample_task, sample_task_runs):
         "Split Name",
         sample_task,
         Train80Test20SplitDefinition,
-        filter=HighRatingDatasetFilter,
+        filter_id="high_rating",
     )
 
+    assert dataset.filter == "high_rating"
+
     # Check that only high-rated task runs are included
     all_ids = []
     for ids in dataset.split_contents.values():
@@ -232,3 +266,90 @@ def test_smaller_sample(sample_task, sample_task_runs):
 
     # Now we should have 0 missing runs. It's okay that dataset has newer data.
     assert dataset.missing_count() == 0
+
+
+@pytest.mark.parametrize(
+    "thinking_data,expected_result",
+    [
+        ({"reasoning": "Here's my answer"}, True),
+        ({"chain_of_thought": "Here's my answer"}, True),
+        ({"unknown": "Here's my answer"}, False),
+        ({}, False),
+        (None, False),
+    ],
+)
+def test_thinking_model_dataset_filter(
+    sample_task_runs, thinking_data, expected_result
+):
+    # Create a task run with thinking output
+    task_run = sample_task_runs[0].model_copy(
+        update={
+            "output": TaskOutput(
+                output="Let me think about this...\nHere's my answer",
+                source=DataSource(
+                    type=DataSourceType.human,
+                    properties={"created_by": "test-user"},
+                ),
+                rating=TaskOutputRating(value=5, type=TaskOutputRatingType.five_star),
+            ),
+            "intermediate_outputs": thinking_data,
+        }
+    )
+
+    assert ThinkingModelDatasetFilter(task_run) is expected_result
+
+
+@pytest.mark.parametrize(
+    "thinking_data,rating,expected_result",
+    [
+        ({"reasoning": "Here's my answer"}, 5, True),
+        ({"chain_of_thought": "Here's my answer"}, 5, True),
+        ({"unknown": "Here's my answer"}, 5, False),
+        ({}, 5, False),
+        (None, 5, False),
+        ({"reasoning": "Here's my answer"}, 1, False),
+        ({"chain_of_thought": "Here's my answer"}, 1, False),
+        ({"unknown": "Here's my answer"}, 1, False),
+        ({}, 1, False),
+        (None, 1, False),
+    ],
+)
+def test_thinking_model_dataset_filter_high_rated(
+    sample_task_runs, thinking_data, rating, expected_result
+):
+    # Create a task run with thinking output
+    task_run = sample_task_runs[0].model_copy(
+        update={
+            "output": TaskOutput(
+                output="Let me think about this...\nHere's my answer",
+                source=DataSource(
+                    type=DataSourceType.human,
+                    properties={"created_by": "test-user"},
+                ),
+                rating=TaskOutputRating(
+                    value=rating, type=TaskOutputRatingType.five_star
+                ),
+            ),
+            "intermediate_outputs": thinking_data,
+        }
+    )
+
+    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

kiln_ai/datamodel/test_datasource.py

@@ -18,14 +18,14 @@ def test_valid_synthetic_data_source():
         properties={
             "model_name": "GPT-4",
             "model_provider": "OpenAI",
-            "prompt_builder_name": "completion",
+            "prompt_id": "simple_prompt_builder",
             "adapter_name": "langchain",
         },
     )
     assert data_source.type == DataSourceType.synthetic
     assert data_source.properties["model_name"] == "GPT-4"
     assert data_source.properties["model_provider"] == "OpenAI"
-    assert data_source.properties["prompt_builder_name"] == "completion"
+    assert data_source.properties["prompt_id"] == "simple_prompt_builder"
     assert data_source.properties["adapter_name"] == "langchain"
 
 
@@ -85,6 +85,7 @@ def test_prompt_type_optional_for_synthetic():
         },
     )
     assert "prompt_builder_name" not in data_source.properties
+    assert "prompt_id" not in data_source.properties
 
 
 def test_private_data_source_properties_not_serialized():