kiln-ai 0.20.1__py3-none-any.whl → 0.22.0__py3-none-any.whl

kiln_ai/datamodel/test_basemodel.py

@@ -1,5 +1,7 @@
 import datetime
 import json
+import logging
+import time
 import uuid
 from pathlib import Path
 from typing import Optional
@@ -342,12 +344,12 @@ def test_delete_no_path():
         ("Hello 👍", "Hello 👍"),
         # Invalid characters are replaced
         ("Hello@World!", "Hello@World!"),
-        ("File.name.txt", "File_name_txt"),
-        ("Special%%%Chars", "Special_Chars"),
-        ("Special#$%Chars", "Special#$_Chars"),
+        ("File.name.txt", "File name txt"),
+        ("Special%%%Chars", "Special Chars"),
+        ("Special#$%Chars", "Special#$ Chars"),
         # Consecutive invalid characters are replaced
-        ("Special%%%Chars", "Special_Chars"),
-        ("path/to/file", "path_to_file"),
+        ("Special%%%Chars", "Special Chars"),
+        ("path/to/file", "path to file"),
         # Leading/trailing special characters are removed
         ("__test__", "test"),
         ("...test...", "test"),
@@ -360,14 +362,14 @@ def test_delete_no_path():
         ("你好_世界", "你好_世界"),
         ("你好_世界_你好", "你好_世界_你好"),
         # Newlines, tabs, and other control characters are replaced
-        ("Hello\nworld", "Hello_world"),
-        ("Hello\tworld", "Hello_world"),
-        ("Hello\rworld", "Hello_world"),
-        ("Hello\fworld", "Hello_world"),
-        ("Hello\bworld", "Hello_world"),
-        ("Hello\vworld", "Hello_world"),
-        ("Hello\0world", "Hello_world"),
-        ("Hello\x00world", "Hello_world"),
+        ("Hello\nworld", "Hello world"),
+        ("Hello\tworld", "Hello world"),
+        ("Hello\rworld", "Hello world"),
+        ("Hello\fworld", "Hello world"),
+        ("Hello\bworld", "Hello world"),
+        ("Hello\vworld", "Hello world"),
+        ("Hello\0world", "Hello world"),
+        ("Hello\x00world", "Hello world"),
     ],
 )
 def test_string_to_valid_name(tmp_path, name, expected):
@@ -491,6 +493,7 @@ def test_from_id_and_parent_path(test_base_parented_file, tmp_model_cache):
     child3.save_to_file()
 
     # Test finding existing child by ID
+    assert child2.id is not None  # Type safety
     found_child = DefaultParentedModel.from_id_and_parent_path(
         child2.id, test_base_parented_file
     )
@@ -513,6 +516,7 @@ def test_from_id_and_parent_path_with_cache(test_base_parented_file, tmp_model_c
     child.save_to_file()
 
     # First load to populate cache
+    assert child.id is not None  # Type safety
     _ = DefaultParentedModel.from_id_and_parent_path(child.id, test_base_parented_file)
 
     # Mock cache to verify it's used
@@ -534,6 +538,258 @@ def test_from_id_and_parent_path_without_parent():
     assert not_found is None
 
 
+def test_from_ids_and_parent_path_basic(test_base_parented_file, tmp_model_cache):
+    """Test basic functionality of from_ids_and_parent_path method"""
+    # 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 multiple children by IDs
+    assert child1.id is not None and child2.id is not None  # Type safety
+    target_ids = {child1.id, child3.id}
+    found_children = DefaultParentedModel.from_ids_and_parent_path(
+        target_ids, test_base_parented_file
+    )
+
+    # Verify correct children were found
+    assert len(found_children) == 2
+    assert child1.id in found_children
+    assert child3.id in found_children
+    assert child2.id not in found_children
+
+    # Verify the returned models have correct data
+    assert found_children[child1.id].name == "Child1"
+    assert found_children[child3.id].name == "Child3"
+
+    # Verify they are not the same instances (deep copies)
+    assert found_children[child1.id] is not child1
+    assert found_children[child3.id] is not child3
+
+
+def test_from_ids_and_parent_path_empty_list(test_base_parented_file):
+    """Test from_ids_and_parent_path with empty ID list"""
+    found_children = DefaultParentedModel.from_ids_and_parent_path(
+        set(), test_base_parented_file
+    )
+    assert found_children == {}
+
+
+def test_from_ids_and_parent_path_none_parent():
+    """Test from_ids_and_parent_path with None parent_path"""
+    found_children = DefaultParentedModel.from_ids_and_parent_path({"any-id"}, None)
+    assert found_children == {}
+
+
+def test_from_ids_and_parent_path_no_matches(test_base_parented_file, tmp_model_cache):
+    """Test from_ids_and_parent_path when no IDs match existing children"""
+    # Set up parent and children models
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+
+    child1 = DefaultParentedModel(parent=parent, name="Child1")
+    child1.save_to_file()
+
+    # Test with non-existent IDs
+    found_children = DefaultParentedModel.from_ids_and_parent_path(
+        {"nonexistent1", "nonexistent2"}, test_base_parented_file
+    )
+    assert found_children == {}
+
+
+def test_from_ids_and_parent_path_partial_matches(
+    test_base_parented_file, tmp_model_cache
+):
+    """Test from_ids_and_parent_path when only some IDs match"""
+    # 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")
+
+    # Save children
+    child1.save_to_file()
+    child2.save_to_file()
+
+    # Test with mix of existing and non-existent IDs
+    assert child1.id is not None and child2.id is not None  # Type safety
+    target_ids = {child1.id, "nonexistent", child2.id, "another_nonexistent"}
+    found_children = DefaultParentedModel.from_ids_and_parent_path(
+        target_ids, test_base_parented_file
+    )
+
+    # Should only find the existing children
+    assert len(found_children) == 2
+    assert child1.id in found_children
+    assert child2.id in found_children
+    assert "nonexistent" not in found_children
+    assert "another_nonexistent" not in found_children
+
+
+def test_from_ids_and_parent_path_with_cache_fallback(
+    test_base_parented_file, tmp_model_cache
+):
+    """Test from_ids_and_parent_path when cache returns None and needs to load file"""
+    # Set up parent and child
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+    child = DefaultParentedModel(parent=parent, name="Child")
+    child.save_to_file()
+
+    # Mock cache to return None for get_model_id, forcing file load
+    tmp_model_cache.get_model_id = MagicMock(return_value=None)
+
+    # Test should still work by loading the file
+    assert child.id is not None  # Type safety
+    found_children = DefaultParentedModel.from_ids_and_parent_path(
+        {child.id}, test_base_parented_file
+    )
+
+    assert len(found_children) == 1
+    assert child.id in found_children
+    assert found_children[child.id].name == "Child"
+
+    # Verify cache was checked
+    tmp_model_cache.get_model_id.assert_called()
+
+
+def test_from_ids_and_parent_path_equivalent_to_individual_lookups(
+    test_base_parented_file, tmp_model_cache
+):
+    """Test that from_ids_and_parent_path returns the same results as individual lookups"""
+    # Set up parent and multiple children
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+
+    children = []
+    for i in range(10):
+        child = DefaultParentedModel(parent=parent, name=f"Child{i}")
+        child.save_to_file()
+        children.append(child)
+
+    # Select 5 children to lookup
+    target_ids = {
+        child.id for child in children[::2] if child.id is not None
+    }  # Every other child
+    assert len(target_ids) == 5  # Ensure we have 5 children to test
+
+    # Test bulk method
+    bulk_results = DefaultParentedModel.from_ids_and_parent_path(
+        target_ids, test_base_parented_file
+    )
+
+    # Test individual method
+    individual_results = {}
+    for target_id in target_ids:
+        result = DefaultParentedModel.from_id_and_parent_path(
+            target_id, test_base_parented_file
+        )
+        if result:
+            individual_results[target_id] = result
+
+    # Results should be equivalent
+    assert len(bulk_results) == len(individual_results) == 5
+
+    for target_id in target_ids:
+        assert target_id in bulk_results
+        assert target_id in individual_results
+
+        # Compare the key attributes
+        bulk_child = bulk_results[target_id]
+        individual_child = individual_results[target_id]
+
+        assert bulk_child.id == individual_child.id
+        assert bulk_child.name == individual_child.name
+        assert bulk_child.model_type == individual_child.model_type
+
+
+# Not actually paid, but we want the "must be run manually" feature of the paid marker as this is very slow
+@pytest.mark.paid
+@pytest.mark.parametrize("num_children", [100, 1000, 2500, 5000])
+def test_from_ids_and_parent_path_benchmark(
+    test_base_parented_file, tmp_model_cache, num_children
+):
+    """Benchmark test for from_ids_and_parent_path method performance at scale"""
+    # Set up parent and many children
+    parent = BaseParentExample.load_from_file(test_base_parented_file)
+
+    children = []
+    for i in range(num_children):
+        child = DefaultParentedModel(parent=parent, name=f"Child{i:05d}")
+        child.save_to_file()
+        children.append(child)
+
+    # look up all children
+    lookup_count = num_children
+    target_ids = {child.id for child in children[:lookup_count] if child.id is not None}
+    assert len(target_ids) == lookup_count
+
+    # Benchmark the bulk method using manual timing
+    def bulk_lookup():
+        return DefaultParentedModel.from_ids_and_parent_path(
+            target_ids, test_base_parented_file
+        )
+
+    # Run bulk method once and time it
+    start_time = time.perf_counter()
+    bulk_result = bulk_lookup()
+    end_time = time.perf_counter()
+    bulk_time = end_time - start_time
+
+    # Verify we got the expected results
+    assert len(bulk_result) == lookup_count
+
+    # Calculate bulk method stats
+    bulk_ops_per_second = lookup_count / bulk_time
+
+    # Benchmark the individual lookup method using manual timing
+    def individual_lookups():
+        results = {}
+        for target_id in target_ids:
+            result = DefaultParentedModel.from_id_and_parent_path(
+                target_id, test_base_parented_file
+            )
+            if result:
+                results[target_id] = result
+        return results
+
+    # Run individual lookup method
+    start_time = time.perf_counter()
+    individual_result = individual_lookups()
+    end_time = time.perf_counter()
+    individual_time = end_time - start_time
+
+    assert len(individual_result) == lookup_count
+    individual_ops_per_second = lookup_count / individual_time
+
+    # Calculate performance comparison
+    speedup = individual_time / bulk_time
+    time_savings_pct = (individual_time - bulk_time) / individual_time * 100
+
+    # Use logging to display results (will show with -s flag or --log-cli-level=INFO)
+    logger = logging.getLogger(__name__)
+    logger.info(
+        f"Benchmark results for {num_children} children, {lookup_count} lookups:"
+    )
+    logger.info(f"  Bulk method: {bulk_time:.4f}s ({bulk_ops_per_second:.2f} ops/sec)")
+    logger.info(
+        f"  Individual method: {individual_time:.4f}s ({individual_ops_per_second:.2f} ops/sec)"
+    )
+    logger.info(
+        f"  Speedup: {speedup:.2f}x faster, {time_savings_pct:.1f}% time savings"
+    )
+
+    assert bulk_time > 0, "Bulk method should complete successfully"
+    assert individual_time > 0, "Individual method should complete successfully"
+    assert speedup >= 1.0, (
+        f"Expected bulk method to be faster, but got {speedup:.2f}x speedup"
+    )
+
+
 class MockAdapter(BaseAdapter):
     """Implementation of BaseAdapter for testing"""
 
@@ -605,7 +861,7 @@ async def test_invoke_parsing_flow(adapter):
         mock_provider.reasoning_capable = True
         with pytest.raises(
             RuntimeError,
-            match="Reasoning is required for this model, but no reasoning was returned.",
+            match=r"^Reasoning is required for this model, but no reasoning was returned.$",
         ):
             await adapter.invoke("test input")
 

kiln_ai/datamodel/test_chunk_models.py

@@ -0,0 +1,317 @@
+import tempfile
+import uuid
+from enum import Enum
+from pathlib import Path
+
+import pytest
+
+from kiln_ai.datamodel.basemodel import KilnAttachmentModel
+from kiln_ai.datamodel.chunk import Chunk, ChunkedDocument, ChunkerConfig, ChunkerType
+from kiln_ai.datamodel.project import Project
+
+
+@pytest.fixture
+def mock_project(tmp_path):
+    project_root = tmp_path / str(uuid.uuid4())
+    project_root.mkdir()
+    project = Project(
+        name="Test Project",
+        description="Test description",
+        path=project_root / "project.kiln",
+    )
+    project.save_to_file()
+    return project
+
+
+class TestFixedWindowChunkerProperties:
+    """Test the FixedWindowChunkerProperties class."""
+
+    def test_required_fields(self):
+        """Test that required fields are set correctly."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={},
+            )
+
+    def test_custom_values(self):
+        """Test that custom values can be set."""
+        config = ChunkerConfig(
+            name="test-chunker",
+            chunker_type=ChunkerType.FIXED_WINDOW,
+            properties={"chunk_size": 512, "chunk_overlap": 20},
+        )
+        assert config.properties == {
+            "chunk_size": 512,
+            "chunk_overlap": 20,
+        }
+
+        assert config.chunk_size() == 512
+        assert config.chunk_overlap() == 20
+
+    def test_validation_positive_values(self):
+        """Test that positive values are accepted."""
+        config = ChunkerConfig(
+            name="test-chunker",
+            chunker_type=ChunkerType.FIXED_WINDOW,
+            properties={"chunk_size": 1, "chunk_overlap": 0},
+        )
+        assert config.properties == {
+            "chunk_size": 1,
+            "chunk_overlap": 0,
+        }
+
+        assert config.chunk_size() == 1
+        assert config.chunk_overlap() == 0
+
+    def test_validation_negative_values(self):
+        """Test that negative values are rejected."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_size": -1, "chunk_overlap": -1},
+            )
+
+    def test_validation_zero_chunk_size(self):
+        """Test that zero chunk size is rejected."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_size": 0, "chunk_overlap": 0},
+            )
+
+    def test_validation_overlap_greater_than_chunk_size(self):
+        """Test that overlap is greater than chunk size."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_size": 100, "chunk_overlap": 101},
+            )
+
+    def test_validation_overlap_less_than_zero(self):
+        """Test that overlap is less than zero."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_size": 100, "chunk_overlap": -1},
+            )
+
+    def test_validation_overlap_without_chunk_size(self):
+        """Test that overlap without chunk size is rejected."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_overlap": 10},
+            )
+
+    def test_validation_chunk_size_without_overlap(self):
+        """Test that chunk size without overlap will raise an error."""
+        with pytest.raises(ValueError, match=r"Chunk overlap is required."):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_size": 100},
+            )
+
+    def test_validation_wrong_type(self):
+        """Test that wrong type is rejected."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_size": "100", "chunk_overlap": 10},
+            )
+
+    def test_validation_none_values(self):
+        """Reject none values - we prefer not to have the properties rather than a None."""
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="test-chunker",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={"chunk_size": None, "chunk_overlap": 15},
+            )
+
+
+class TestChunkerType:
+    """Test the ChunkerType enum."""
+
+    def test_enum_values(self):
+        """Test that enum has the expected values."""
+        assert ChunkerType.FIXED_WINDOW == "fixed_window"
+
+    def test_enum_inheritance(self):
+        """Test that ChunkerType inherits from str and Enum."""
+        assert issubclass(ChunkerType, str)
+        assert issubclass(ChunkerType, Enum)
+
+    def test_enum_comparison(self):
+        """Test enum comparison operations."""
+        assert ChunkerType.FIXED_WINDOW == "fixed_window"
+        assert ChunkerType.FIXED_WINDOW.value == "fixed_window"
+
+
+class TestChunkerConfig:
+    """Test the ChunkerConfig class."""
+
+    def test_optional_description(self):
+        """Test that description is optional."""
+        config = ChunkerConfig(
+            name="test-chunker",
+            chunker_type=ChunkerType.FIXED_WINDOW,
+            properties={
+                "chunk_size": 100,
+                "chunk_overlap": 10,
+            },
+        )
+        assert config.description is None
+
+        config_with_desc = ChunkerConfig(
+            name="test-chunker",
+            description="A test chunker",
+            chunker_type=ChunkerType.FIXED_WINDOW,
+            properties={
+                "chunk_size": 100,
+                "chunk_overlap": 10,
+            },
+        )
+        assert config_with_desc.description == "A test chunker"
+
+    def test_name_validation(self):
+        """Test name field validation."""
+        # Test valid name
+        config = ChunkerConfig(
+            name="valid-name_123",
+            chunker_type=ChunkerType.FIXED_WINDOW,
+            properties={
+                "chunk_size": 100,
+                "chunk_overlap": 10,
+            },
+        )
+        assert config.name == "valid-name_123"
+
+        # Test invalid name (contains special characters)
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="invalid@name",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={},
+            )
+
+        # Test empty name
+        with pytest.raises(ValueError):
+            ChunkerConfig(
+                name="",
+                chunker_type=ChunkerType.FIXED_WINDOW,
+                properties={},
+            )
+
+    def test_parent_project_method_no_parent(self):
+        """Test parent_project method when no parent is set."""
+        config = ChunkerConfig(
+            name="test-chunker",
+            chunker_type=ChunkerType.FIXED_WINDOW,
+            properties={
+                "chunk_size": 100,
+                "chunk_overlap": 10,
+            },
+        )
+        assert config.parent_project() is None
+
+
+class TestChunk:
+    """Test the Chunk class."""
+
+    def test_required_fields(self):
+        """Test that required fields are properly validated."""
+        # Create a temporary file for the content
+        with tempfile.NamedTemporaryFile(delete=True) as tmp_file:
+            tmp_file.write(b"test content")
+            tmp_path = Path(tmp_file.name)
+
+            attachment = KilnAttachmentModel.from_file(tmp_path)
+            chunk = Chunk(content=attachment)
+            assert chunk.content == attachment
+
+    def test_content_validation(self):
+        """Test that content field is properly validated."""
+        # Create a temporary file for the attachment
+        with tempfile.NamedTemporaryFile(delete=True) as tmp_file:
+            tmp_file.write(b"test content")
+            tmp_path = Path(tmp_file.name)
+
+            # Test with valid attachment
+            attachment = KilnAttachmentModel.from_file(tmp_path)
+            chunk = Chunk(content=attachment)
+            assert chunk.content == attachment
+
+            # Test that attachment is required
+            with pytest.raises(ValueError):
+                Chunk(content=None)
+
+
+class TestChunkedDocument:
+    """Test the ChunkedDocument class."""
+
+    def test_required_fields(self):
+        """Test that required fields are properly validated."""
+        chunks = []
+        doc = ChunkedDocument(chunks=chunks, chunker_config_id="fake-id")
+        assert doc.chunks == chunks
+
+    def test_with_chunks(self):
+        """Test with actual chunks."""
+        # Create a temporary file for the attachment
+        with tempfile.NamedTemporaryFile(delete=True) as tmp_file:
+            tmp_file.write(b"test content")
+            tmp_path = Path(tmp_file.name)
+
+            attachment = KilnAttachmentModel.from_file(tmp_path)
+            chunk1 = Chunk(content=attachment)
+            chunk2 = Chunk(content=attachment)
+
+            chunks = [chunk1, chunk2]
+            doc = ChunkedDocument(chunks=chunks, chunker_config_id="fake-id")
+            assert doc.chunks == chunks
+            assert len(doc.chunks) == 2
+
+    def test_parent_extraction_method_no_parent(self):
+        """Test parent_extraction method when no parent is set."""
+        doc = ChunkedDocument(chunks=[], chunker_config_id="fake-id")
+        assert doc.parent_extraction() is None
+
+    def test_empty_chunks_list(self):
+        """Test that empty chunks list is valid."""
+        doc = ChunkedDocument(chunks=[], chunker_config_id="fake-id")
+        assert doc.chunks == []
+        assert len(doc.chunks) == 0
+
+    def test_chunks_validation(self):
+        """Test that chunks field validation works correctly."""
+        # Create a temporary file for the attachment
+        with tempfile.NamedTemporaryFile(delete=True) as tmp_file:
+            tmp_file.write(b"test content")
+            tmp_path = Path(tmp_file.name)
+
+            # Test with valid list of chunks
+            attachment = KilnAttachmentModel.from_file(tmp_path)
+            chunk = Chunk(content=attachment)
+            chunks = [chunk]
+
+            doc = ChunkedDocument(
+                chunks=chunks,
+                chunker_config_id="fake-id",
+            )
+            assert doc.chunks == chunks
+
+            # Test that chunks must be a list
+            with pytest.raises(ValueError):
+                ChunkedDocument(
+                    chunks=chunk,
+                    chunker_config_id="fake-id",
+                )

kiln_ai/datamodel/test_dataset_split.py

@@ -120,7 +120,7 @@ def test_dataset_split_validation():
         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"):
+    with pytest.raises(ValueError, match=r"sum of split percentages must be 1.0"):
         DatasetSplit(
             name="test_split",
             splits=invalid_splits,