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

kiln_ai/datamodel/test_tool_id.py

@@ -4,10 +4,12 @@ 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,
     mcp_server_and_tool_name_from_id,
+    rag_config_id_from_id,
 )
 
 
@@ -113,6 +115,36 @@ 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::")
+
 
 class TestMcpServerAndToolNameFromId:
     """Test the mcp_server_and_tool_name_from_id function."""
@@ -197,6 +229,9 @@ 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",
         ]
 
         for tool_id in valid_ids:
@@ -212,6 +247,8 @@ class TestToolIdPydanticType:
             "mcp::remote::server",
             "mcp::local::",
             "mcp::local::server",
+            "kiln_tool::rag::",
+            "kiln_tool::rag::config::extra",
         ]
 
         for invalid_id in invalid_ids:
@@ -237,3 +274,47 @@ 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 == ""

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

@@ -26,6 +26,7 @@ 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::"
 
 
@@ -58,6 +59,15 @@ 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
+
     raise ValueError(f"Invalid tool ID: {id}")
 
 
@@ -81,3 +91,15 @@ 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]

kiln_ai/datamodel/vector_store.py

@@ -0,0 +1,141 @@
+from enum import Enum
+from typing import TYPE_CHECKING, Union
+
+from pydantic import BaseModel, Field, model_validator
+
+from kiln_ai.datamodel.basemodel import FilenameString, KilnParentedModel
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+from kiln_ai.utils.validation import (
+    validate_return_dict_prop,
+    validate_return_dict_prop_optional,
+)
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.project import Project
+
+
+class VectorStoreType(str, Enum):
+    LANCE_DB_FTS = "lancedb_fts"
+    LANCE_DB_HYBRID = "lancedb_hybrid"
+    LANCE_DB_VECTOR = "lancedb_vector"
+
+
+class LanceDBConfigBaseProperties(BaseModel):
+    similarity_top_k: int = Field(
+        description="The number of results to return from the vector store.",
+    )
+    overfetch_factor: int = Field(
+        description="The overfetch factor to use for the vector search.",
+    )
+    vector_column_name: str = Field(
+        description="The name of the vector column in the vector store.",
+    )
+    text_key: str = Field(
+        description="The name of the text column in the vector store.",
+    )
+    doc_id_key: str = Field(
+        description="The name of the document id column in the vector store.",
+    )
+    nprobes: int | None = Field(
+        description="The number of probes to use for the vector search.",
+        default=None,
+    )
+
+
+class VectorStoreConfig(KilnParentedModel):
+    name: FilenameString = Field(
+        description="A name for your own reference to identify the vector store config.",
+    )
+    description: str | None = Field(
+        description="A description for your own reference.",
+        default=None,
+    )
+    store_type: VectorStoreType = Field(
+        description="The type of vector store to use.",
+    )
+    properties: dict[str, str | int | float | None] = Field(
+        description="The properties of the vector store config, specific to the selected store_type.",
+    )
+
+    @model_validator(mode="after")
+    def validate_properties(self):
+        match self.store_type:
+            case (
+                VectorStoreType.LANCE_DB_FTS
+                | VectorStoreType.LANCE_DB_HYBRID
+                | VectorStoreType.LANCE_DB_VECTOR
+            ):
+                return self.validate_lancedb_properties(self.store_type)
+            case _:
+                raise_exhaustive_enum_error(self.store_type)
+
+    def validate_lancedb_properties(self, store_type: VectorStoreType):
+        err_msg_prefix = f"LanceDB vector store configs properties for {store_type}:"
+        validate_return_dict_prop(
+            self.properties, "similarity_top_k", int, err_msg_prefix
+        )
+        validate_return_dict_prop(
+            self.properties, "overfetch_factor", int, err_msg_prefix
+        )
+        validate_return_dict_prop(
+            self.properties, "vector_column_name", str, err_msg_prefix
+        )
+        validate_return_dict_prop(self.properties, "text_key", str, err_msg_prefix)
+        validate_return_dict_prop(self.properties, "doc_id_key", str, err_msg_prefix)
+
+        # nprobes is only used for vector and hybrid queries
+        if (
+            store_type == VectorStoreType.LANCE_DB_VECTOR
+            or store_type == VectorStoreType.LANCE_DB_HYBRID
+        ):
+            validate_return_dict_prop(self.properties, "nprobes", int, err_msg_prefix)
+
+        return self
+
+    @property
+    def lancedb_properties(self) -> LanceDBConfigBaseProperties:
+        err_msg_prefix = "LanceDB vector store configs properties:"
+        return LanceDBConfigBaseProperties(
+            similarity_top_k=validate_return_dict_prop(
+                self.properties,
+                "similarity_top_k",
+                int,
+                err_msg_prefix,
+            ),
+            overfetch_factor=validate_return_dict_prop(
+                self.properties,
+                "overfetch_factor",
+                int,
+                err_msg_prefix,
+            ),
+            vector_column_name=validate_return_dict_prop(
+                self.properties,
+                "vector_column_name",
+                str,
+                err_msg_prefix,
+            ),
+            text_key=validate_return_dict_prop(
+                self.properties,
+                "text_key",
+                str,
+                err_msg_prefix,
+            ),
+            doc_id_key=validate_return_dict_prop(
+                self.properties,
+                "doc_id_key",
+                str,
+                err_msg_prefix,
+            ),
+            nprobes=validate_return_dict_prop_optional(
+                self.properties,
+                "nprobes",
+                int,
+                err_msg_prefix,
+            ),
+        )
+
+    # Workaround to return typed parent without importing Project
+    def parent_project(self) -> Union["Project", None]:
+        if self.parent is None or self.parent.__class__.__name__ != "Project":
+            return None
+        return self.parent  # type: ignore

kiln_ai/tools/mcp_session_manager.py

@@ -3,6 +3,7 @@ import os
 import subprocess
 import sys
 from contextlib import asynccontextmanager
+from datetime import timedelta
 from typing import AsyncGenerator
 
 import httpx
@@ -171,7 +172,9 @@ class MCPSessionManager:
 
         try:
             async with stdio_client(server_params) as (read, write):
-                async with ClientSession(read, write) as session:
+                async with ClientSession(
+                    read, write, read_timeout_seconds=timedelta(seconds=8)
+                ) as session:
                     await session.initialize()
                     yield session
         except Exception as e: