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

kiln_ai/datamodel/test_tool_id.py

@@ -4,10 +4,13 @@ from pydantic import BaseModel, ValidationError
 from kiln_ai.datamodel.tool_id import (
     MCP_LOCAL_TOOL_ID_PREFIX,
     MCP_REMOTE_TOOL_ID_PREFIX,
+    RAG_TOOL_ID_PREFIX,
     KilnBuiltInToolId,
     ToolId,
     _check_tool_id,
+    kiln_task_server_id_from_tool_id,
     mcp_server_and_tool_name_from_id,
+    rag_config_id_from_id,
 )
 
 
@@ -113,6 +116,69 @@ class TestCheckToolId:
         with pytest.raises(ValueError, match="Invalid tool ID"):
             _check_tool_id("mcp::wrong::server::tool")
 
+    def test_valid_rag_tools(self):
+        """Test validation of valid RAG tools."""
+        valid_ids = [
+            "kiln_tool::rag::config1",
+            "kiln_tool::rag::my_rag_config",
+            "kiln_tool::rag::test_config_123",
+        ]
+        for tool_id in valid_ids:
+            result = _check_tool_id(tool_id)
+            assert result == tool_id
+
+    def test_invalid_rag_format(self):
+        """Test validation fails for invalid RAG tool formats."""
+        # These IDs start with the RAG prefix but have invalid formats
+        rag_invalid_ids = [
+            "kiln_tool::rag::",  # Missing config ID
+            "kiln_tool::rag::config::extra",  # Too many parts
+        ]
+
+        for invalid_id in rag_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid RAG tool ID"):
+                _check_tool_id(invalid_id)
+
+    def test_rag_tool_empty_config_id(self):
+        """Test that RAG tool with empty config ID is handled properly."""
+        # This tests the case where rag_config_id_from_id returns empty string
+        # which should trigger line 66 in the source
+        with pytest.raises(ValueError, match="Invalid RAG tool ID"):
+            _check_tool_id("kiln_tool::rag::")
+
+    def test_valid_kiln_task_tools(self):
+        """Test validation of valid Kiln task tools."""
+        valid_ids = [
+            "kiln_task::server1",
+            "kiln_task::my_server",
+            "kiln_task::test_server_123",
+            "kiln_task::server_with_underscores",
+            "kiln_task::server-with-dashes",
+            "kiln_task::server.with.dots",
+        ]
+        for tool_id in valid_ids:
+            result = _check_tool_id(tool_id)
+            assert result == tool_id
+
+    def test_invalid_kiln_task_format(self):
+        """Test validation fails for invalid Kiln task tool formats."""
+        # These IDs start with the Kiln task prefix but have invalid formats
+        kiln_task_invalid_ids = [
+            "kiln_task::",  # Missing server ID
+            "kiln_task::server::extra",  # Too many parts
+            "kiln_task::server::tool::extra",  # Too many parts
+        ]
+
+        for invalid_id in kiln_task_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid Kiln task tool ID"):
+                _check_tool_id(invalid_id)
+
+    def test_kiln_task_tool_empty_server_id(self):
+        """Test that Kiln task tool with empty server ID is handled properly."""
+        # This tests the case where kiln_task_server_id_from_tool_id returns empty string which should raise an error
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID"):
+            _check_tool_id("kiln_task::")
+
 
 class TestMcpServerAndToolNameFromId:
     """Test the mcp_server_and_tool_name_from_id function."""
@@ -188,7 +254,7 @@ class TestToolIdPydanticType:
             model = self._ModelWithToolId(tool_id=tool_id.value)
             assert model.tool_id == tool_id.value
 
-    def test_valid_mcp_tools(self):
+    def test_valid_tool_ids(self):
         """Test ToolId validates MCP remote and local tools."""
         valid_ids = [
             # Remote MCP tools
@@ -197,6 +263,12 @@ class TestToolIdPydanticType:
             # Local MCP tools
             "mcp::local::server1::tool1",
             "mcp::local::my_server::my_tool",
+            # RAG tools
+            "kiln_tool::rag::config1",
+            "kiln_tool::rag::my_rag_config",
+            # Kiln task tools
+            "kiln_task::server1",
+            "kiln_task::my_server",
         ]
 
         for tool_id in valid_ids:
@@ -212,6 +284,10 @@ class TestToolIdPydanticType:
             "mcp::remote::server",
             "mcp::local::",
             "mcp::local::server",
+            "kiln_tool::rag::",
+            "kiln_tool::rag::config::extra",
+            "kiln_task::",
+            "kiln_task::server::extra",
         ]
 
         for invalid_id in invalid_ids:
@@ -237,3 +313,113 @@ class TestConstants:
     def test_mcp_local_tool_id_prefix(self):
         """Test the MCP local tool ID prefix constant."""
         assert MCP_LOCAL_TOOL_ID_PREFIX == "mcp::local::"
+
+    def test_rag_tool_id_prefix(self):
+        """Test the RAG tool ID prefix constant."""
+        assert RAG_TOOL_ID_PREFIX == "kiln_tool::rag::"
+
+
+class TestRagConfigIdFromId:
+    """Test the rag_config_id_from_id function."""
+
+    def test_valid_rag_ids(self):
+        """Test parsing valid RAG tool IDs."""
+        test_cases = [
+            ("kiln_tool::rag::config1", "config1"),
+            ("kiln_tool::rag::my_rag_config", "my_rag_config"),
+            ("kiln_tool::rag::test_config_123", "test_config_123"),
+            ("kiln_tool::rag::a", "a"),  # Minimal valid case
+        ]
+
+        for tool_id, expected in test_cases:
+            result = rag_config_id_from_id(tool_id)
+            assert result == expected
+
+    def test_invalid_rag_ids(self):
+        """Test parsing fails for invalid RAG tool IDs."""
+        # Test various invalid formats that should trigger line 104
+        invalid_ids = [
+            "kiln_tool::rag::config::extra",  # Too many parts (4 parts)
+            "wrong::rag::config",  # Wrong prefix
+            "kiln_tool::wrong::config",  # Wrong middle part
+            "rag::config",  # Too few parts (2 parts)
+            "",  # Empty string
+            "single_part",  # Only 1 part
+        ]
+
+        for invalid_id in invalid_ids:
+            with pytest.raises(ValueError, match="Invalid RAG tool ID"):
+                rag_config_id_from_id(invalid_id)
+
+    def test_rag_id_with_empty_config_id(self):
+        """Test that RAG tool ID with empty config ID returns empty string."""
+        # This is actually valid according to the parser - it returns empty string
+        # The validation for empty config ID happens in _check_tool_id
+        result = rag_config_id_from_id("kiln_tool::rag::")
+        assert result == ""
+
+
+class TestKilnTaskServerIdFromToolId:
+    """Test the kiln_task_server_id_from_tool_id function."""
+
+    def test_valid_kiln_task_ids(self):
+        """Test parsing valid Kiln task tool IDs."""
+        test_cases = [
+            ("kiln_task::server1", "server1"),
+            ("kiln_task::my_server", "my_server"),
+            ("kiln_task::test_server_123", "test_server_123"),
+            ("kiln_task::a", "a"),  # Minimal valid case
+            ("kiln_task::server_with_underscores", "server_with_underscores"),
+            ("kiln_task::server-with-dashes", "server-with-dashes"),
+            ("kiln_task::server.with.dots", "server.with.dots"),
+        ]
+
+        for tool_id, expected in test_cases:
+            result = kiln_task_server_id_from_tool_id(tool_id)
+            assert result == expected
+
+    def test_invalid_kiln_task_ids(self):
+        """Test parsing fails for invalid Kiln task tool IDs."""
+        # Test various invalid formats
+        invalid_ids = [
+            "kiln_task::",  # Empty server ID
+            "kiln_task::server::extra",  # Too many parts (3 parts)
+            "kiln_task::server::tool::extra",  # Too many parts (4 parts)
+            "wrong::server",  # Wrong prefix
+            "kiln_wrong::server",  # Wrong prefix
+            "task::server",  # Too few parts (2 parts)
+            "",  # Empty string
+            "single_part",  # Only 1 part
+            "kiln_task",  # Missing server ID
+        ]
+
+        for invalid_id in invalid_ids:
+            with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+                kiln_task_server_id_from_tool_id(invalid_id)
+
+    def test_kiln_task_id_with_empty_server_id(self):
+        """Test that Kiln task tool ID with empty server ID raises error."""
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+            kiln_task_server_id_from_tool_id("kiln_task::")
+
+    def test_kiln_task_id_with_whitespace_server_id(self):
+        """Test that Kiln task tool ID with whitespace-only server ID raises error."""
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+            kiln_task_server_id_from_tool_id("kiln_task::")
+
+    def test_kiln_task_id_with_multiple_colons(self):
+        """Test that Kiln task tool ID with multiple colons raises error."""
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+            kiln_task_server_id_from_tool_id("kiln_task::server::extra")
+
+    def test_kiln_task_id_case_sensitivity(self):
+        """Test that Kiln task tool IDs are case sensitive."""
+        # These should work
+        result1 = kiln_task_server_id_from_tool_id("kiln_task::Server")
+        assert result1 == "Server"
+
+        result2 = kiln_task_server_id_from_tool_id("kiln_task::SERVER")
+        assert result2 == "SERVER"
+
+        result3 = kiln_task_server_id_from_tool_id("kiln_task::server")
+        assert result3 == "server"

kiln_ai/datamodel/test_vector_store.py

@@ -0,0 +1,320 @@
+import pytest
+from pydantic import ValidationError
+
+from kiln_ai.datamodel.project import Project
+from kiln_ai.datamodel.vector_store import (
+    LanceDBConfigBaseProperties,
+    VectorStoreConfig,
+    VectorStoreType,
+)
+
+
+@pytest.fixture
+def mock_project(tmp_path):
+    project_path = tmp_path / "test_project" / "project.kiln"
+    project_path.parent.mkdir()
+
+    project = Project(name="Test Project", path=project_path)
+    project.save_to_file()
+
+    return project
+
+
+@pytest.fixture
+def mock_vector_store_fts_config_properties():
+    return {
+        "similarity_top_k": 10,
+        "overfetch_factor": 2,
+        "vector_column_name": "vector",
+        "text_key": "text",
+        "doc_id_key": "doc_id",
+    }
+
+
+@pytest.fixture
+def mock_vector_store_vector_config_properties():
+    return {
+        "similarity_top_k": 10,
+        "overfetch_factor": 2,
+        "vector_column_name": "vector",
+        "text_key": "text",
+        "doc_id_key": "doc_id",
+        "nprobes": 1,
+    }
+
+
+class TestVectorStoreType:
+    def test_vector_store_type_values(self):
+        """Test that VectorStoreType enum has expected values."""
+        assert VectorStoreType.LANCE_DB_FTS == "lancedb_fts"
+        assert VectorStoreType.LANCE_DB_HYBRID == "lancedb_hybrid"
+        assert VectorStoreType.LANCE_DB_VECTOR == "lancedb_vector"
+
+
+class TestLanceDBConfigBaseProperties:
+    def test_valid_lance_db_config_base_properties(self):
+        """Test creating valid LanceDBConfigBaseProperties."""
+        config = LanceDBConfigBaseProperties(
+            similarity_top_k=10,
+            overfetch_factor=2,
+            vector_column_name="vector",
+            text_key="text",
+            doc_id_key="doc_id",
+            nprobes=1,
+        )
+
+        assert config.similarity_top_k == 10
+        assert config.overfetch_factor == 2
+        assert config.vector_column_name == "vector"
+        assert config.text_key == "text"
+        assert config.doc_id_key == "doc_id"
+        assert config.nprobes == 1
+
+    def test_lance_db_config_base_properties_without_nprobes(self):
+        """Test creating LanceDBConfigBaseProperties without nprobes."""
+        config = LanceDBConfigBaseProperties(
+            similarity_top_k=10,
+            overfetch_factor=2,
+            vector_column_name="vector",
+            text_key="text",
+            doc_id_key="doc_id",
+        )
+
+        assert config.similarity_top_k == 10
+        assert config.nprobes is None
+
+
+class TestVectorStoreConfig:
+    def test_invalid_store_type(self, mock_vector_store_fts_config_properties):
+        """Test creating VectorStoreConfig with invalid store type."""
+        with pytest.raises(ValidationError, match="Input should be"):
+            VectorStoreConfig(
+                name="test_store",
+                store_type="invalid_type",  # type: ignore
+                properties=mock_vector_store_fts_config_properties,
+            )
+
+    def test_invalid_store_type_after_creation(
+        self, mock_vector_store_fts_config_properties
+    ):
+        """Test creating VectorStoreConfig with invalid store type after creation."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+        )
+        with pytest.raises(ValidationError, match="Input should be"):
+            config.store_type = "invalid_type"  # type: ignore
+
+    def test_valid_lance_db_fts_vector_store_config(
+        self, mock_vector_store_fts_config_properties
+    ):
+        """Test creating valid VectorStoreConfig with LanceDB FTS."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+        )
+
+        assert config.name == "test_store"
+        assert config.store_type == VectorStoreType.LANCE_DB_FTS
+        assert config.properties["similarity_top_k"] == 10
+        assert config.properties["overfetch_factor"] == 2
+        assert config.properties["vector_column_name"] == "vector"
+        assert config.properties["text_key"] == "text"
+        assert config.properties["doc_id_key"] == "doc_id"
+
+    def test_valid_lance_db_vector_store_config(
+        self, mock_vector_store_vector_config_properties
+    ):
+        """Test creating valid VectorStoreConfig with LanceDB Vector."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_VECTOR,
+            properties=mock_vector_store_vector_config_properties,
+        )
+
+        assert config.name == "test_store"
+        assert config.store_type == VectorStoreType.LANCE_DB_VECTOR
+        assert config.properties["similarity_top_k"] == 10
+        assert config.properties["nprobes"] == 1
+
+    def test_valid_lance_db_hybrid_store_config(
+        self, mock_vector_store_vector_config_properties
+    ):
+        """Test creating valid VectorStoreConfig with LanceDB Hybrid."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_HYBRID,
+            properties=mock_vector_store_vector_config_properties,
+        )
+
+        assert config.name == "test_store"
+        assert config.store_type == VectorStoreType.LANCE_DB_HYBRID
+        assert config.properties["nprobes"] == 1
+
+    def test_vector_store_config_missing_required_property(
+        self, mock_vector_store_fts_config_properties
+    ):
+        """Test VectorStoreConfig validation fails when required property is missing."""
+        mock_vector_store_fts_config_properties.pop("similarity_top_k")
+        with pytest.raises(
+            ValidationError,
+            match=r".*similarity_top_k is a required property",
+        ):
+            VectorStoreConfig(
+                name="test_store",
+                store_type=VectorStoreType.LANCE_DB_FTS,
+                properties=mock_vector_store_fts_config_properties,
+            )
+
+    def test_vector_store_config_invalid_property_type(
+        self, mock_vector_store_fts_config_properties
+    ):
+        """Test VectorStoreConfig validation fails when property has wrong type."""
+        mock_vector_store_fts_config_properties["similarity_top_k"] = "not_an_int"
+        with pytest.raises(
+            ValidationError,
+            match=r".*similarity_top_k must be of type",
+        ):
+            VectorStoreConfig(
+                name="test_store",
+                store_type=VectorStoreType.LANCE_DB_FTS,
+                properties=mock_vector_store_fts_config_properties,
+            )
+
+    def test_vector_store_config_fts_missing_nprobes_is_valid(
+        self, mock_vector_store_fts_config_properties
+    ):
+        """Test VectorStoreConfig with FTS type doesn't require nprobes."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+        )
+        assert config.store_type == VectorStoreType.LANCE_DB_FTS
+
+    def test_vector_store_config_vector_missing_nprobes_fails(
+        self, mock_vector_store_vector_config_properties
+    ):
+        """Test VectorStoreConfig with VECTOR type requires nprobes."""
+        mock_vector_store_vector_config_properties.pop("nprobes")
+        with pytest.raises(
+            ValidationError,
+            match=r".*nprobes is a required property",
+        ):
+            VectorStoreConfig(
+                name="test_store",
+                store_type=VectorStoreType.LANCE_DB_VECTOR,
+                properties=mock_vector_store_vector_config_properties,
+            )
+
+    def test_lancedb_properties(self, mock_vector_store_vector_config_properties):
+        """Test lancedb_properties method returns correct LanceDBConfigBaseProperties."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_VECTOR,
+            properties=mock_vector_store_vector_config_properties,
+        )
+
+        props = config.lancedb_properties
+
+        assert isinstance(props, LanceDBConfigBaseProperties)
+        assert props.similarity_top_k == 10
+        assert props.overfetch_factor == 2
+        assert props.vector_column_name == "vector"
+        assert props.text_key == "text"
+        assert props.doc_id_key == "doc_id"
+        assert props.nprobes == 1
+
+    def test_vector_store_config_inherits_from_kiln_parented_model(
+        self, mock_vector_store_fts_config_properties
+    ):
+        """Test that VectorStoreConfig inherits from KilnParentedModel."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+        )
+
+        # Check that it has the expected base fields
+        assert hasattr(config, "id")
+        assert hasattr(config, "v")
+        assert hasattr(config, "created_at")
+        assert hasattr(config, "created_by")
+        assert hasattr(config, "parent")
+
+    @pytest.mark.parametrize(
+        "name",
+        ["valid_name", "valid name", "valid-name", "valid_name_123", "VALID_NAME"],
+    )
+    def test_vector_store_config_valid_names(
+        self, name, mock_vector_store_fts_config_properties
+    ):
+        """Test VectorStoreConfig accepts valid names."""
+        config = VectorStoreConfig(
+            name=name,
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+        )
+        assert config.name == name
+
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "",
+            "a" * 121,  # Too long
+        ],
+    )
+    def test_vector_store_config_invalid_names(
+        self, name, mock_vector_store_fts_config_properties
+    ):
+        """Test VectorStoreConfig rejects invalid names."""
+        with pytest.raises(ValidationError):
+            VectorStoreConfig(
+                name=name,
+                store_type=VectorStoreType.LANCE_DB_FTS,
+                properties=mock_vector_store_fts_config_properties,
+            )
+
+    def test_parent_project(
+        self, mock_project, mock_vector_store_fts_config_properties
+    ):
+        """Test that parent project is returned correctly."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+            parent=mock_project,
+        )
+
+        assert config.parent_project() is mock_project
+
+    def test_vector_store_config_parent_project_none(
+        self, mock_vector_store_fts_config_properties
+    ):
+        """Test that parent project is None if not set."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+        )
+
+        assert config.parent_project() is None
+
+    def test_project_has_vector_store_configs(
+        self, mock_project, mock_vector_store_fts_config_properties
+    ):
+        """Test that project has vector store configs."""
+        config = VectorStoreConfig(
+            name="test_store",
+            store_type=VectorStoreType.LANCE_DB_FTS,
+            properties=mock_vector_store_fts_config_properties,
+            parent=mock_project,
+        )
+        config.save_to_file()
+
+        assert len(mock_project.vector_store_configs(readonly=True)) == 1
+        assert config.id in [
+            vc.id for vc in mock_project.vector_store_configs(readonly=True)
+        ]

kiln_ai/datamodel/tool_id.py

@@ -14,6 +14,7 @@ Tool IDs can be one of:
 - A kiln built-in tool name: kiln_tool::add_numbers
 - A remote MCP tool: mcp::remote::<server_id>::<tool_name>
 - A local MCP tool: mcp::local::<server_id>::<tool_name>
+- A Kiln task tool: kiln_task::<server_id>
 - More coming soon like kiln_project_tool::rag::RAG_CONFIG_ID
 """
 
@@ -26,7 +27,9 @@ class KilnBuiltInToolId(str, Enum):
 
 
 MCP_REMOTE_TOOL_ID_PREFIX = "mcp::remote::"
+RAG_TOOL_ID_PREFIX = "kiln_tool::rag::"
 MCP_LOCAL_TOOL_ID_PREFIX = "mcp::local::"
+KILN_TASK_TOOL_ID_PREFIX = "kiln_task::"
 
 
 def _check_tool_id(id: str) -> str:
@@ -58,6 +61,24 @@ def _check_tool_id(id: str) -> str:
             )
         return id
 
+    # RAG tools must have format: kiln_tool::rag::<rag_config_id>
+    if id.startswith(RAG_TOOL_ID_PREFIX):
+        rag_config_id = rag_config_id_from_id(id)
+        if not rag_config_id:
+            raise ValueError(
+                f"Invalid RAG tool ID: {id}. Expected format: 'kiln_tool::rag::<rag_config_id>'."
+            )
+        return id
+
+    # Kiln task tools must have format: kiln_task::<server_id>
+    if id.startswith(KILN_TASK_TOOL_ID_PREFIX):
+        server_id = kiln_task_server_id_from_tool_id(id)
+        if not server_id:
+            raise ValueError(
+                f"Invalid Kiln task tool ID: {id}. Expected format: 'kiln_task::<server_id>'."
+            )
+        return id
+
     raise ValueError(f"Invalid tool ID: {id}")
 
 
@@ -81,3 +102,40 @@ def mcp_server_and_tool_name_from_id(id: str) -> tuple[str, str]:
                 f"Invalid MCP tool ID: {id}. Expected format: 'mcp::(remote|local)::<server_id>::<tool_name>'."
             )
     return parts[2], parts[3]  # server_id, tool_name
+
+
+def rag_config_id_from_id(id: str) -> str:
+    """
+    Get the RAG config ID from the ID.
+    """
+    parts = id.split("::")
+    if not id.startswith(RAG_TOOL_ID_PREFIX) or len(parts) != 3:
+        raise ValueError(
+            f"Invalid RAG tool ID: {id}. Expected format: 'kiln_tool::rag::<rag_config_id>'."
+        )
+    return parts[2]
+
+
+def kiln_task_server_id_from_tool_id(tool_id: str) -> str:
+    """
+    Get the server ID from the tool ID.
+    """
+    if not tool_id.startswith(KILN_TASK_TOOL_ID_PREFIX):
+        raise ValueError(
+            f"Invalid Kiln task tool ID format: {tool_id}. Expected format: 'kiln_task::<server_id>'."
+        )
+
+    # Remove prefix and split on ::
+    remaining = tool_id[len(KILN_TASK_TOOL_ID_PREFIX) :]
+    if not remaining:
+        raise ValueError(
+            f"Invalid Kiln task tool ID format: {tool_id}. Expected format: 'kiln_task::<server_id>'."
+        )
+    parts = remaining.split("::")
+
+    if len(parts) != 1 or not parts[0].strip():
+        raise ValueError(
+            f"Invalid Kiln task tool ID format: {tool_id}. Expected format: 'kiln_task::<server_id>'."
+        )
+
+    return parts[0]  # server_id