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

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

@@ -0,0 +1,105 @@
+from enum import Enum
+from typing import Annotated
+
+from pydantic import AfterValidator
+
+ToolId = Annotated[
+    str,
+    AfterValidator(lambda v: _check_tool_id(v)),
+]
+"""
+A pydantic type that validates strings containing a valid tool ID.
+
+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>
+- More coming soon like kiln_project_tool::rag::RAG_CONFIG_ID
+"""
+
+
+class KilnBuiltInToolId(str, Enum):
+    ADD_NUMBERS = "kiln_tool::add_numbers"
+    SUBTRACT_NUMBERS = "kiln_tool::subtract_numbers"
+    MULTIPLY_NUMBERS = "kiln_tool::multiply_numbers"
+    DIVIDE_NUMBERS = "kiln_tool::divide_numbers"
+
+
+MCP_REMOTE_TOOL_ID_PREFIX = "mcp::remote::"
+RAG_TOOL_ID_PREFIX = "kiln_tool::rag::"
+MCP_LOCAL_TOOL_ID_PREFIX = "mcp::local::"
+
+
+def _check_tool_id(id: str) -> str:
+    """
+    Check that the tool ID is valid.
+    """
+    if not id or not isinstance(id, str):
+        raise ValueError(f"Invalid tool ID: {id}")
+
+    # Build in tools
+    if id in KilnBuiltInToolId.__members__.values():
+        return id
+
+    # MCP remote tools must have format: mcp::remote::<server_id>::<tool_name>
+    if id.startswith(MCP_REMOTE_TOOL_ID_PREFIX):
+        server_id, tool_name = mcp_server_and_tool_name_from_id(id)
+        if not server_id or not tool_name:
+            raise ValueError(
+                f"Invalid remote MCP tool ID: {id}. Expected format: 'mcp::remote::<server_id>::<tool_name>'."
+            )
+        return id
+
+    # MCP local tools must have format: mcp::local::<server_id>::<tool_name>
+    if id.startswith(MCP_LOCAL_TOOL_ID_PREFIX):
+        server_id, tool_name = mcp_server_and_tool_name_from_id(id)
+        if not server_id or not tool_name:
+            raise ValueError(
+                f"Invalid local MCP tool ID: {id}. Expected format: 'mcp::local::<server_id>::<tool_name>'."
+            )
+        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}")
+
+
+def mcp_server_and_tool_name_from_id(id: str) -> tuple[str, str]:
+    """
+    Get the tool server ID and tool name from the ID.
+    """
+    parts = id.split("::")
+    if len(parts) != 4:
+        # Determine if it's remote or local for the error message
+        if id.startswith(MCP_REMOTE_TOOL_ID_PREFIX):
+            raise ValueError(
+                f"Invalid remote MCP tool ID: {id}. Expected format: 'mcp::remote::<server_id>::<tool_name>'."
+            )
+        elif id.startswith(MCP_LOCAL_TOOL_ID_PREFIX):
+            raise ValueError(
+                f"Invalid local MCP tool ID: {id}. Expected format: 'mcp::local::<server_id>::<tool_name>'."
+            )
+        else:
+            raise ValueError(
+                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/__init__.py

@@ -0,0 +1,8 @@
+from kiln_ai.tools.base_tool import KilnTool, KilnToolInterface
+from kiln_ai.tools.tool_registry import tool_from_id
+
+__all__ = [
+    "KilnTool",
+    "KilnToolInterface",
+    "tool_from_id",
+]

kiln_ai/tools/base_tool.py

@@ -0,0 +1,82 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from kiln_ai.datamodel.json_schema import validate_schema_dict
+from kiln_ai.datamodel.tool_id import KilnBuiltInToolId, ToolId
+
+
+class KilnToolInterface(ABC):
+    """
+    Abstract interface defining the core API that all Kiln tools must implement.
+    This ensures consistency across all tool implementations.
+    """
+
+    @abstractmethod
+    async def run(self, **kwargs) -> Any:
+        """Execute the tool with the given parameters."""
+        pass
+
+    @abstractmethod
+    async def toolcall_definition(self) -> Dict[str, Any]:
+        """Return the OpenAI-compatible tool definition for this tool."""
+        pass
+
+    @abstractmethod
+    async def id(self) -> ToolId:
+        """Return a unique identifier for this tool."""
+        pass
+
+    @abstractmethod
+    async def name(self) -> str:
+        """Return the tool name (function name) of this tool."""
+        pass
+
+    @abstractmethod
+    async def description(self) -> str:
+        """Return a description of what this tool does."""
+        pass
+
+
+class KilnTool(KilnToolInterface):
+    """
+    Base helper class that provides common functionality for tool implementations.
+    Subclasses only need to implement run() and provide tool configuration.
+    """
+
+    def __init__(
+        self,
+        tool_id: KilnBuiltInToolId,
+        name: str,
+        description: str,
+        parameters_schema: Dict[str, Any],
+    ):
+        self._id = tool_id
+        self._name = name
+        self._description = description
+        validate_schema_dict(parameters_schema)
+        self._parameters_schema = parameters_schema
+
+    async def id(self) -> KilnBuiltInToolId:
+        return self._id
+
+    async def name(self) -> str:
+        return self._name
+
+    async def description(self) -> str:
+        return self._description
+
+    async def toolcall_definition(self) -> Dict[str, Any]:
+        """Generate OpenAI-compatible tool definition."""
+        return {
+            "type": "function",
+            "function": {
+                "name": await self.name(),
+                "description": await self.description(),
+                "parameters": self._parameters_schema,
+            },
+        }
+
+    @abstractmethod
+    async def run(self, **kwargs) -> str:
+        """Subclasses must implement the actual tool logic."""
+        pass

kiln_ai/tools/built_in_tools/__init__.py

@@ -0,0 +1,13 @@
+from kiln_ai.tools.built_in_tools.math_tools import (
+    AddTool,
+    DivideTool,
+    MultiplyTool,
+    SubtractTool,
+)
+
+__all__ = [
+    "AddTool",
+    "DivideTool",
+    "MultiplyTool",
+    "SubtractTool",
+]