kiln-ai 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/datamodel/test_embedding_models.py

@@ -0,0 +1,448 @@
+import uuid
+from pathlib import Path
+
+import pytest
+
+from kiln_ai.datamodel.basemodel import KilnAttachmentModel
+from kiln_ai.datamodel.chunk import Chunk, ChunkedDocument
+from kiln_ai.datamodel.embedding import ChunkEmbeddings, Embedding, EmbeddingConfig
+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
+
+
+@pytest.fixture
+def mock_chunked_document(tmp_path):
+    # Create a temporary file for the attachment
+    tmp_dir = tmp_path / str(uuid.uuid4())
+    tmp_dir.mkdir()
+
+    tmp_path_file = Path(tmp_dir) / f"{uuid.uuid4()}.txt"
+    tmp_path_file.write_text("test content")
+
+    attachment = KilnAttachmentModel.from_file(tmp_path_file)
+    chunks = [Chunk(content=attachment) for _ in range(3)]
+
+    doc = ChunkedDocument(
+        chunks=chunks,
+        chunker_config_id="fake-chunker-id",
+        path=Path(tmp_dir) / "chunked_document.kiln",
+    )
+    doc.save_to_file()
+
+    return doc
+
+
+class TestEmbeddingConfig:
+    """Test the EmbeddingConfig class."""
+
+    def test_required_fields(self):
+        """Test that required fields are set correctly."""
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+        )
+        assert config.name == "test-embedding"
+        assert config.model_provider_name == "openai"
+        assert config.model_name == "openai_text_embedding_3_small"
+        assert config.properties == {"dimensions": 1536}
+
+    def test_optional_description(self):
+        """Test that description is optional."""
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+        )
+        assert config.description is None
+
+        config_with_desc = EmbeddingConfig(
+            name="test-embedding",
+            description="A test embedding config",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+        )
+        assert config_with_desc.description == "A test embedding config"
+
+    def test_name_validation(self):
+        """Test name field validation."""
+        # Test valid name
+        config = EmbeddingConfig(
+            name="valid-name_123",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+        )
+        assert config.name == "valid-name_123"
+
+        # Test empty name
+        with pytest.raises(ValueError):
+            EmbeddingConfig(
+                name="",
+                model_provider_name="openai",
+                model_name="openai_text_embedding_3_small",
+                properties={"dimensions": 1536},
+            )
+
+    def test_properties_validation(self):
+        """Test properties field validation."""
+        # Test with valid properties
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={
+                "dimensions": 1536,
+            },
+        )
+        assert config.properties == {
+            "dimensions": 1536,
+        }
+
+        # Test with empty properties
+        config_empty = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={},
+        )
+        assert config_empty.properties == {}
+
+    def test_parent_project_method_no_parent(self):
+        """Test parent_project method when no parent is set."""
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+        )
+        assert config.parent_project() is None
+
+    def test_parent_project_method_with_project_parent(self, mock_project):
+        """Test parent_project method when parent is a Project."""
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+            parent=mock_project,
+        )
+        assert config.parent_project() == mock_project
+
+    def test_model_provider_name_validation(self, mock_project):
+        """Test model_provider_name field validation."""
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={},
+            parent=mock_project,
+        )
+        assert config.model_provider_name == "openai"
+
+        with pytest.raises(ValueError):
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name="invalid-provider",
+                model_name="openai_text_embedding_3_small",
+                parent=mock_project,
+                properties={},
+            )
+
+    def test_custom_dimensions_validation(self):
+        """Test that custom dimensions are properly validated."""
+
+        # this model supports custom dimensions
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+        )
+        assert config.properties == {"dimensions": 1536}
+
+        # dimensions is negative
+        with pytest.raises(ValueError, match="Dimensions must be a positive integer"):
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name="openai",
+                model_name="openai_text_embedding_3_small",
+                properties={"dimensions": -1},
+            )
+
+        # dimensions is not an integer
+        with pytest.raises(ValueError, match="Dimensions must be a positive integer"):
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name="openai",
+                model_name="openai_text_embedding_3_small",
+                properties={"dimensions": 1.5},
+            )
+
+        # dimensions is not a positive integer
+        with pytest.raises(ValueError, match="Dimensions must be a positive integer"):
+            EmbeddingConfig(
+                name="test-embedding",
+                model_provider_name="openai",
+                model_name="openai_text_embedding_3_small",
+                properties={"dimensions": "512"},
+            )
+
+    def test_dimensions_optional(self):
+        """Test that dimensions is optional and should be ignored if not provided."""
+        config = EmbeddingConfig(
+            name="test-embedding",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={},
+        )
+        assert config.properties == {}
+
+
+class TestEmbedding:
+    """Test the Embedding class."""
+
+    def test_required_fields(self):
+        """Test that required fields are properly validated."""
+        vector = [0.1, 0.2, 0.3, 0.4, 0.5]
+        embedding = Embedding(vector=vector)
+        assert embedding.vector == vector
+
+    def test_vector_validation(self):
+        """Test that vector field is properly validated."""
+        # Test with valid vector
+        vector = [0.1, 0.2, 0.3]
+        embedding = Embedding(vector=vector)
+        assert embedding.vector == vector
+
+        # Test with empty vector
+        empty_vector = []
+        embedding_empty = Embedding(vector=empty_vector)
+        assert embedding_empty.vector == empty_vector
+
+        # Test with large vector
+        large_vector = [0.1] * 1536
+        embedding_large = Embedding(vector=large_vector)
+        assert len(embedding_large.vector) == 1536
+
+    def test_vector_types(self):
+        """Test that vector accepts different numeric types."""
+        # Test with integers
+        int_vector = [1, 2, 3, 4, 5]
+        embedding_int = Embedding(vector=int_vector)
+        assert embedding_int.vector == int_vector
+
+        # Test with floats
+        float_vector = [1.1, 2.2, 3.3, 4.4, 5.5]
+        embedding_float = Embedding(vector=float_vector)
+        assert embedding_float.vector == float_vector
+
+        # Test with mixed types
+        mixed_vector = [1, 2.5, 3, 4.7, 5]
+        embedding_mixed = Embedding(vector=mixed_vector)
+        assert embedding_mixed.vector == mixed_vector
+
+
+class TestChunkEmbeddings:
+    """Test the ChunkEmbeddings class."""
+
+    def test_required_fields(self):
+        """Test that required fields are properly validated."""
+        embedding_config_id = "test-config-id"
+        embeddings = [
+            Embedding(vector=[0.1, 0.2, 0.3]),
+            Embedding(vector=[0.4, 0.5, 0.6]),
+        ]
+
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id=embedding_config_id,
+            embeddings=embeddings,
+        )
+        assert chunk_embeddings.embedding_config_id == embedding_config_id
+        assert chunk_embeddings.embeddings == embeddings
+
+    def test_embeddings_validation(self):
+        """Test that embeddings field validation works correctly."""
+        embedding_config_id = "test-config-id"
+
+        # Test with valid list of embeddings
+        embeddings = [Embedding(vector=[0.1, 0.2, 0.3])]
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id=embedding_config_id,
+            embeddings=embeddings,
+        )
+        assert chunk_embeddings.embeddings == embeddings
+
+        # Test with empty embeddings list
+        empty_embeddings = []
+        chunk_embeddings_empty = ChunkEmbeddings(
+            embedding_config_id=embedding_config_id,
+            embeddings=empty_embeddings,
+        )
+        assert chunk_embeddings_empty.embeddings == empty_embeddings
+
+        # Test with multiple embeddings
+        multiple_embeddings = [
+            Embedding(vector=[0.1, 0.2, 0.3]),
+            Embedding(vector=[0.4, 0.5, 0.6]),
+            Embedding(vector=[0.7, 0.8, 0.9]),
+        ]
+        chunk_embeddings_multiple = ChunkEmbeddings(
+            embedding_config_id=embedding_config_id,
+            embeddings=multiple_embeddings,
+        )
+        assert chunk_embeddings_multiple.embeddings == multiple_embeddings
+        assert len(chunk_embeddings_multiple.embeddings) == 3
+
+    def test_embedding_config_id_validation(self):
+        """Test embedding_config_id field validation."""
+        embeddings = [Embedding(vector=[0.1, 0.2, 0.3])]
+
+        # Test with valid ID
+        valid_id = "test-config-id-123"
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id=valid_id,
+            embeddings=embeddings,
+        )
+        assert chunk_embeddings.embedding_config_id == valid_id
+
+        # Test with numeric string ID
+        numeric_id = "12345"
+        chunk_embeddings_numeric = ChunkEmbeddings(
+            embedding_config_id=numeric_id,
+            embeddings=embeddings,
+        )
+        assert chunk_embeddings_numeric.embedding_config_id == numeric_id
+
+    def test_parent_chunked_document_method_no_parent(self):
+        """Test parent_chunked_document method when no parent is set."""
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id="test-config-id",
+            embeddings=[Embedding(vector=[0.1, 0.2, 0.3])],
+        )
+        assert chunk_embeddings.parent_chunked_document() is None
+
+    def test_parent_chunked_document_method_with_chunked_document_parent(
+        self, mock_chunked_document
+    ):
+        """Test parent_chunked_document method when parent is a ChunkedDocument."""
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id="test-config-id",
+            embeddings=[Embedding(vector=[0.1, 0.2, 0.3])],
+            parent=mock_chunked_document,
+        )
+        assert chunk_embeddings.parent_chunked_document() == mock_chunked_document
+
+    def test_embeddings_correspond_to_chunks(self, mock_chunked_document):
+        """Test that embeddings correspond to chunks in the parent chunked document."""
+        # Create embeddings that match the number of chunks in the parent
+        num_chunks = len(mock_chunked_document.chunks)
+        embeddings = [Embedding(vector=[0.1, 0.2, 0.3]) for _ in range(num_chunks)]
+
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id="test-config-id",
+            embeddings=embeddings,
+            parent=mock_chunked_document,
+        )
+        assert len(chunk_embeddings.embeddings) == num_chunks
+
+    def test_embeddings_with_different_vector_sizes(self):
+        """Test embeddings with different vector sizes."""
+        embedding_config_id = "test-config-id"
+        embeddings = [
+            Embedding(vector=[0.1, 0.2, 0.3]),  # 3 dimensions
+            Embedding(vector=[0.4, 0.5, 0.6, 0.7]),  # 4 dimensions
+            Embedding(vector=[0.8, 0.9]),  # 2 dimensions
+        ]
+
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id=embedding_config_id,
+            embeddings=embeddings,
+        )
+        assert len(chunk_embeddings.embeddings) == 3
+        assert len(chunk_embeddings.embeddings[0].vector) == 3
+        assert len(chunk_embeddings.embeddings[1].vector) == 4
+        assert len(chunk_embeddings.embeddings[2].vector) == 2
+
+
+class TestEmbeddingIntegration:
+    """Integration tests for embedding models."""
+
+    def test_embedding_config_with_project_parent(self, mock_project):
+        """Test EmbeddingConfig with Project parent."""
+        config = EmbeddingConfig(
+            name="test-embedding",
+            description="Test embedding configuration",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+            parent=mock_project,
+        )
+        assert config.parent_project() == mock_project
+        assert config.name == "test-embedding"
+        assert config.model_provider_name == "openai"
+        assert config.model_name == "openai_text_embedding_3_small"
+
+    def test_chunk_embeddings_with_chunked_document_parent(self, mock_chunked_document):
+        """Test ChunkEmbeddings with ChunkedDocument parent."""
+        # Create embeddings for each chunk
+        embeddings = []
+        for chunk in mock_chunked_document.chunks:
+            # Create a mock embedding (in real usage, this would be generated by the embedding model)
+            embedding = Embedding(vector=[0.1, 0.2, 0.3, 0.4, 0.5])
+            embeddings.append(embedding)
+
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id="test-config-id",
+            embeddings=embeddings,
+            parent=mock_chunked_document,
+        )
+        assert chunk_embeddings.parent_chunked_document() == mock_chunked_document
+        assert len(chunk_embeddings.embeddings) == len(mock_chunked_document.chunks)
+
+    def test_embedding_workflow(self, mock_project, mock_chunked_document):
+        """Test a complete embedding workflow."""
+        # 1. Create an embedding config
+        embedding_config = EmbeddingConfig(
+            name="test-embedding-config",
+            description="Test embedding configuration for workflow",
+            model_provider_name="openai",
+            model_name="openai_text_embedding_3_small",
+            properties={"dimensions": 1536},
+            parent=mock_project,
+        )
+
+        # 2. Create embeddings for the chunked document
+        embeddings = []
+        for chunk in mock_chunked_document.chunks:
+            # Simulate embedding generation
+            embedding = Embedding(vector=[0.1] * 1536)
+            embeddings.append(embedding)
+
+        # 3. Create chunk embeddings
+        chunk_embeddings = ChunkEmbeddings(
+            embedding_config_id=embedding_config.id,
+            embeddings=embeddings,
+            parent=mock_chunked_document,
+        )
+
+        # 4. Verify the relationships
+        assert embedding_config.parent_project() == mock_project
+        assert chunk_embeddings.parent_chunked_document() == mock_chunked_document
+        assert len(chunk_embeddings.embeddings) == len(mock_chunked_document.chunks)
+        assert chunk_embeddings.embedding_config_id == embedding_config.id

kiln_ai/datamodel/test_eval_model.py

@@ -402,13 +402,13 @@ def test_eval_run_five_star_score_validation(valid_eval_config, valid_eval_run_d
     assert run.scores["accuracy"] == 4.5
 
     # Invalid scores
-    with pytest.raises(ValueError, match="must be a float between 1.0 and 5.0"):
+    with pytest.raises(ValueError, match=r"must be a float between 1.0 and 5.0"):
         run = EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"accuracy": 0.5}},
         )
 
-    with pytest.raises(ValueError, match="must be a float between 1.0 and 5.0"):
+    with pytest.raises(ValueError, match=r"must be a float between 1.0 and 5.0"):
         run = EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"accuracy": 5.5}},
@@ -442,13 +442,13 @@ def test_eval_run_pass_fail_score_validation(valid_eval_config, valid_eval_run_d
     assert run.scores["check"] == 0.0
 
     # Invalid scores
-    with pytest.raises(ValueError, match="must be a float between 0.0 and 1.0"):
+    with pytest.raises(ValueError, match=r"must be a float between 0.0 and 1.0"):
         run = EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"check": -0.1}},
         )
 
-    with pytest.raises(ValueError, match="must be a float between 0.0 and 1.0"):
+    with pytest.raises(ValueError, match=r"must be a float between 0.0 and 1.0"):
         run = EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"check": 1.1}},
@@ -485,13 +485,13 @@ def test_eval_run_pass_fail_critical_score_validation(
     assert run.scores["critical"] == -1.0
 
     # Invalid scores
-    with pytest.raises(ValueError, match="must be a float between -1.0 and 1.0"):
+    with pytest.raises(ValueError, match=r"must be a float between -1.0 and 1.0"):
         run = EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"critical": -1.1}},
         )
 
-    with pytest.raises(ValueError, match="must be a float between -1.0 and 1.0"):
+    with pytest.raises(ValueError, match=r"must be a float between -1.0 and 1.0"):
         run = EvalRun(
             parent=valid_eval_config,
             **{**valid_eval_run_data, "scores": {"critical": 1.1}},

kiln_ai/datamodel/test_example_models.py

@@ -797,3 +797,178 @@ def test_usage_model_in_task_run(valid_task_run):
     assert task_run.usage.output_tokens == 50
     assert task_run.usage.total_tokens == 150
     assert task_run.usage.cost == 0.002
+
+
+@pytest.mark.parametrize(
+    "usage1_data,usage2_data,expected_data",
+    [
+        # None + None = None
+        (
+            {
+                "input_tokens": None,
+                "output_tokens": None,
+                "total_tokens": None,
+                "cost": None,
+            },
+            {
+                "input_tokens": None,
+                "output_tokens": None,
+                "total_tokens": None,
+                "cost": None,
+            },
+            {
+                "input_tokens": None,
+                "output_tokens": None,
+                "total_tokens": None,
+                "cost": None,
+            },
+        ),
+        # None + value = value
+        (
+            {
+                "input_tokens": None,
+                "output_tokens": None,
+                "total_tokens": None,
+                "cost": None,
+            },
+            {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "total_tokens": 150,
+                "cost": 0.005,
+            },
+            {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "total_tokens": 150,
+                "cost": 0.005,
+            },
+        ),
+        # value + None = value
+        (
+            {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "total_tokens": 150,
+                "cost": 0.005,
+            },
+            {
+                "input_tokens": None,
+                "output_tokens": None,
+                "total_tokens": None,
+                "cost": None,
+            },
+            {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "total_tokens": 150,
+                "cost": 0.005,
+            },
+        ),
+        # value1 + value2 = value1 + value2
+        (
+            {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "total_tokens": 150,
+                "cost": 0.005,
+            },
+            {
+                "input_tokens": 200,
+                "output_tokens": 75,
+                "total_tokens": 275,
+                "cost": 0.010,
+            },
+            {
+                "input_tokens": 300,
+                "output_tokens": 125,
+                "total_tokens": 425,
+                "cost": 0.015,
+            },
+        ),
+        # Mixed scenarios
+        (
+            {
+                "input_tokens": 100,
+                "output_tokens": None,
+                "total_tokens": 150,
+                "cost": None,
+            },
+            {
+                "input_tokens": None,
+                "output_tokens": 75,
+                "total_tokens": None,
+                "cost": 0.010,
+            },
+            {
+                "input_tokens": 100,
+                "output_tokens": 75,
+                "total_tokens": 150,
+                "cost": 0.010,
+            },
+        ),
+        # Edge case: zeros
+        (
+            {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "cost": 0.0},
+            {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "total_tokens": 150,
+                "cost": 0.005,
+            },
+            {
+                "input_tokens": 100,
+                "output_tokens": 50,
+                "total_tokens": 150,
+                "cost": 0.005,
+            },
+        ),
+    ],
+)
+def test_usage_addition(usage1_data, usage2_data, expected_data):
+    """Test Usage addition with various combinations of None and numeric values."""
+    usage1 = Usage(**usage1_data)
+    usage2 = Usage(**usage2_data)
+    result = usage1 + usage2
+
+    assert result.input_tokens == expected_data["input_tokens"]
+    assert result.output_tokens == expected_data["output_tokens"]
+    assert result.total_tokens == expected_data["total_tokens"]
+    assert result.cost == expected_data["cost"]
+
+
+def test_usage_addition_type_error():
+    """Test that adding Usage to non-Usage raises TypeError."""
+    usage = Usage(input_tokens=100, output_tokens=50, total_tokens=150, cost=0.005)
+
+    with pytest.raises(TypeError, match="Cannot add Usage with"):
+        usage + "not_a_usage"  # type: ignore
+
+    with pytest.raises(TypeError, match="Cannot add Usage with"):
+        usage + 42  # type: ignore
+
+    with pytest.raises(TypeError, match="Cannot add Usage with"):
+        usage + {"input_tokens": 100}  # type: ignore
+
+
+def test_usage_addition_immutability():
+    """Test that addition creates new Usage objects and doesn't mutate originals."""
+    usage1 = Usage(input_tokens=100, output_tokens=50, total_tokens=150, cost=0.005)
+    usage2 = Usage(input_tokens=200, output_tokens=75, total_tokens=275, cost=0.010)
+
+    original_usage1_data = usage1.model_dump()
+    original_usage2_data = usage2.model_dump()
+
+    result = usage1 + usage2
+
+    # Original objects should be unchanged
+    assert usage1.model_dump() == original_usage1_data
+    assert usage2.model_dump() == original_usage2_data
+
+    # Result should be a new object
+    assert result is not usage1
+    assert result is not usage2
+    assert result.input_tokens == 300
+    assert result.output_tokens == 125
+    assert result.total_tokens == 425
+    assert result.cost == 0.015