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

kiln_ai/datamodel/test_datasource.py

@@ -38,6 +38,15 @@ def test_valid_file_import_data_source():
     assert data_source.properties["file_name"] == "test.txt"
 
 
+def test_empty_valid_tool_call_data_source():
+    data_source = DataSource(
+        type=DataSourceType.tool_call,
+        properties={},
+    )
+    assert data_source.type == DataSourceType.tool_call
+    assert data_source.properties == {}
+
+
 def test_missing_required_property():
     with pytest.raises(ValidationError, match="'created_by' is required for"):
         DataSource(type=DataSourceType.human)
@@ -79,6 +88,35 @@ def test_not_allowed_property():
         )
 
 
+def test_not_allowed_property_tool_call():
+    with pytest.raises(
+        ValidationError,
+        match="'created_by' is not allowed for",
+    ):
+        DataSource(
+            type=DataSourceType.tool_call,
+            properties={
+                "model_name": "GPT-4",
+                "model_provider": "OpenAI",
+                "adapter_name": "langchain",
+                "created_by": "John Doe",
+            },
+        )
+
+
+def test_not_allowed_file_name_tool_call():
+    with pytest.raises(
+        ValidationError,
+        match="'file_name' is not allowed for",
+    ):
+        DataSource(
+            type=DataSourceType.tool_call,
+            properties={
+                "file_name": "test.txt",
+            },
+        )
+
+
 def test_extra_properties():
     data_source = DataSource(
         type=DataSourceType.synthetic,
@@ -94,6 +132,18 @@ def test_extra_properties():
     assert data_source.properties["max_tokens"] == 100
 
 
+def test_extra_properties_tool_call():
+    data_source = DataSource(
+        type=DataSourceType.tool_call,
+        properties={
+            "temperature": 0.7,
+            "max_tokens": 100,
+        },
+    )
+    assert data_source.properties["temperature"] == 0.7
+    assert data_source.properties["max_tokens"] == 100
+
+
 def test_prompt_type_optional_for_synthetic():
     data_source = DataSource(
         type=DataSourceType.synthetic,

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}},