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

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/base_tool.py

@@ -1,10 +1,19 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import Any, Dict
 
 from kiln_ai.datamodel.json_schema import validate_schema_dict
 from kiln_ai.datamodel.tool_id import KilnBuiltInToolId, ToolId
 
 
+@dataclass
+class ToolCallContext:
+    """Context passed to tools when they are called, containing information from the calling task."""
+
+    """Used for Kiln Tasks as Tools, to know if the tool call should save the task run it invoked to that task's Dataset."""
+    allow_saving: bool = True
+
+
 class KilnToolInterface(ABC):
     """
     Abstract interface defining the core API that all Kiln tools must implement.
@@ -12,8 +21,8 @@ class KilnToolInterface(ABC):
     """
 
     @abstractmethod
-    async def run(self, **kwargs) -> Any:
-        """Execute the tool with the given parameters."""
+    async def run(self, context: ToolCallContext | None = None, **kwargs) -> Any:
+        """Execute the tool with the given parameters and calling context if provided."""
         pass
 
     @abstractmethod
@@ -77,6 +86,6 @@ class KilnTool(KilnToolInterface):
         }
 
     @abstractmethod
-    async def run(self, **kwargs) -> str:
+    async def run(self, context: ToolCallContext | None = None, **kwargs) -> Any:
         """Subclasses must implement the actual tool logic."""
         pass

kiln_ai/tools/built_in_tools/math_tools.py

@@ -27,7 +27,9 @@ class AddTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Add two numbers and return the result."""
         return str(a + b)
 
@@ -57,7 +59,9 @@ class SubtractTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Subtract b from a and return the result."""
         return str(a - b)
 
@@ -84,7 +88,9 @@ class MultiplyTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Multiply two numbers and return the result."""
         return str(a * b)
 
@@ -117,7 +123,9 @@ class DivideTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Divide a by b and return the result."""
         if b == 0:
             raise ZeroDivisionError("Cannot divide by zero")

kiln_ai/tools/kiln_task_tool.py

@@ -0,0 +1,158 @@
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Any, Dict
+
+from kiln_ai.datamodel import Task
+from kiln_ai.datamodel.external_tool_server import ExternalToolServer
+from kiln_ai.datamodel.task import TaskRunConfig
+from kiln_ai.datamodel.task_output import DataSource, DataSourceType
+from kiln_ai.datamodel.tool_id import ToolId
+from kiln_ai.tools.base_tool import KilnToolInterface, ToolCallContext
+from kiln_ai.utils.project_utils import project_from_id
+
+
+@dataclass
+class KilnTaskToolResult:
+    output: str
+    kiln_task_tool_data: str
+
+
+class KilnTaskTool(KilnToolInterface):
+    """
+    A tool that wraps a Kiln task, allowing it to be called as a function.
+
+    This tool loads a task by ID and executes it using the specified run configuration.
+    """
+
+    def __init__(
+        self,
+        project_id: str,
+        tool_id: str,
+        data_model: ExternalToolServer,
+    ):
+        self._project_id = project_id
+        self._tool_server_model = data_model
+        self._tool_id = tool_id
+
+        self._name = data_model.properties.get("name", "")
+        self._description = data_model.properties.get("description", "")
+        self._task_id = data_model.properties.get("task_id", "")
+        self._run_config_id = data_model.properties.get("run_config_id", "")
+
+    async def id(self) -> ToolId:
+        return self._tool_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,
+            },
+        }
+
+    async def run(
+        self, context: ToolCallContext | None = None, **kwargs
+    ) -> KilnTaskToolResult:
+        """Execute the wrapped Kiln task with the given parameters and calling context."""
+        if context is None:
+            raise ValueError("Context is required for running a KilnTaskTool.")
+
+        # Determine the input format
+        if self._task.input_json_schema:
+            # Structured input - pass kwargs directly
+            input = kwargs
+        else:
+            # Plaintext input - extract from 'input' parameter
+            if "input" in kwargs:
+                input = kwargs["input"]
+            else:
+                raise ValueError(f"Input not found in kwargs: {kwargs}")
+
+        # These imports are here to avoid circular chains
+        from kiln_ai.adapters.adapter_registry import adapter_for_task
+        from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
+
+        # Create adapter and run the task using the calling task's allow_saving setting
+        adapter = adapter_for_task(
+            self._task,
+            run_config_properties=self._run_config.run_config_properties,
+            base_adapter_config=AdapterConfig(
+                allow_saving=context.allow_saving,
+                default_tags=["tool_call"],
+            ),
+        )
+        task_run = await adapter.invoke(
+            input,
+            input_source=DataSource(
+                type=DataSourceType.tool_call,
+                run_config=self._run_config.run_config_properties,
+            ),
+        )
+
+        return KilnTaskToolResult(
+            output=task_run.output.output,
+            kiln_task_tool_data=f"{self._project_id}:::{self._tool_id}:::{self._task.id}:::{task_run.id}",
+        )
+
+    @cached_property
+    def _task(self) -> Task:
+        # Load the project first
+        project = project_from_id(self._project_id)
+        if project is None:
+            raise ValueError(f"Project not found: {self._project_id}")
+
+        # Load the task from the project
+        task = Task.from_id_and_parent_path(self._task_id, project.path)
+        if task is None:
+            raise ValueError(
+                f"Task not found: {self._task_id} in project {self._project_id}"
+            )
+        return task
+
+    @cached_property
+    def _run_config(self) -> TaskRunConfig:
+        run_config = next(
+            (
+                run_config
+                for run_config in self._task.run_configs(readonly=True)
+                if run_config.id == self._run_config_id
+            ),
+            None,
+        )
+        if run_config is None:
+            raise ValueError(
+                f"Task run config not found: {self._run_config_id} for task {self._task_id} in project {self._project_id}"
+            )
+        return run_config
+
+    @cached_property
+    def parameters_schema(self) -> Dict[str, Any]:
+        if self._task.input_json_schema:
+            # Use the task's input schema directly if it exists
+            parameters_schema = self._task.input_schema()
+        else:
+            # For plaintext tasks, create a simple string input parameter
+            parameters_schema = {
+                "type": "object",
+                "properties": {
+                    "input": {
+                        "type": "string",
+                        "description": "Plaintext input for the tool.",
+                    }
+                },
+                "required": ["input"],
+            }
+        if parameters_schema is None:
+            raise ValueError(
+                f"Failed to create parameters schema for tool_id {self._tool_id}"
+            )
+        return parameters_schema

kiln_ai/tools/mcp_server_tool.py

@@ -5,7 +5,7 @@ from mcp.types import Tool as MCPTool
 
 from kiln_ai.datamodel.external_tool_server import ExternalToolServer
 from kiln_ai.datamodel.tool_id import MCP_REMOTE_TOOL_ID_PREFIX, ToolId
-from kiln_ai.tools.base_tool import KilnToolInterface
+from kiln_ai.tools.base_tool import KilnToolInterface, ToolCallContext
 from kiln_ai.tools.mcp_session_manager import MCPSessionManager
 
 
@@ -38,7 +38,7 @@ class MCPServerTool(KilnToolInterface):
             },
         }
 
-    async def run(self, **kwargs) -> Any:
+    async def run(self, context: ToolCallContext | None = None, **kwargs) -> str:
         result = await self._call_tool(**kwargs)
 
         if result.isError:

kiln_ai/tools/mcp_session_manager.py

@@ -2,7 +2,9 @@ import logging
 import os
 import subprocess
 import sys
+import tempfile
 from contextlib import asynccontextmanager
+from datetime import timedelta
 from typing import AsyncGenerator
 
 import httpx
@@ -18,6 +20,8 @@ from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 logger = logging.getLogger(__name__)
 
+LOCAL_MCP_ERROR_INSTRUCTION = "Please verify your command, arguments, and environment variables, and consult the server's documentation for the correct setup."
+
 
 class MCPSessionManager:
     """
@@ -50,6 +54,8 @@ class MCPSessionManager:
             case ToolServerType.local_mcp:
                 async with self._create_local_mcp_session(tool_server) as session:
                     yield session
+            case ToolServerType.kiln_task:
+                raise ValueError("Kiln task tools are not available from an MCP server")
             case _:
                 raise_exhaustive_enum_error(tool_server.type)
 
@@ -163,33 +169,56 @@ class MCPSessionManager:
             env_vars["PATH"] = self._get_path()
 
         # Set the server parameters
+        cwd = os.path.join(Config.settings_dir(), "cache", "mcp_cache")
+        os.makedirs(cwd, exist_ok=True)
         server_params = StdioServerParameters(
-            command=command,
-            args=args,
-            env=env_vars,
+            command=command, args=args, env=env_vars, cwd=cwd
         )
 
-        try:
-            async with stdio_client(server_params) as (read, write):
-                async with ClientSession(read, write) as session:
-                    await session.initialize()
-                    yield session
-        except Exception as e:
-            # Check for MCP errors. Things like wrong arguments would fall here.
-            mcp_error = self._extract_first_exception(e, McpError)
-            if mcp_error and isinstance(mcp_error, McpError):
-                self._raise_local_mcp_error(mcp_error)
-
-            # Re-raise the original error but with a friendlier message
-            self._raise_local_mcp_error(e)
-
-    def _raise_local_mcp_error(self, e: Exception):
+        # Create temporary file to capture MCP server stderr
+        # Use errors="replace" to handle non-UTF-8 bytes gracefully
+        with tempfile.TemporaryFile(
+            mode="w+", encoding="utf-8", errors="replace"
+        ) as err_log:
+            try:
+                async with stdio_client(server_params, errlog=err_log) as (
+                    read,
+                    write,
+                ):
+                    async with ClientSession(
+                        read, write, read_timeout_seconds=timedelta(seconds=30)
+                    ) as session:
+                        await session.initialize()
+                        yield session
+            except Exception as e:
+                # Read stderr content from temporary file for debugging
+                err_log.seek(0)  # Read from the start of the file
+                stderr_content = err_log.read()
+                if stderr_content:
+                    logger.error(
+                        f"MCP server '{tool_server.name}' stderr output: {stderr_content}"
+                    )
+
+                # Check for MCP errors. Things like wrong arguments would fall here.
+                mcp_error = self._extract_first_exception(e, McpError)
+                if mcp_error and isinstance(mcp_error, McpError):
+                    self._raise_local_mcp_error(mcp_error, stderr_content)
+
+                # Re-raise the original error but with a friendlier message
+                self._raise_local_mcp_error(e, stderr_content)
+
+    def _raise_local_mcp_error(self, e: Exception, stderr: str):
         """
-        Raise a ValueError with a friendlier message for local MCP errors.
+        Raise a RuntimeError with a friendlier message for local MCP errors.
         """
-        raise RuntimeError(
-            f"MCP server failed to start. Please verify your command, arguments, and environment variables, and consult the server's documentation for the correct setup. Original error: {e}"
-        ) from e
+        error_msg = f"'{e}'"
+
+        if stderr:
+            error_msg += f"\nMCP server error: {stderr}"
+
+        error_msg += f"\n{LOCAL_MCP_ERROR_INSTRUCTION}"
+
+        raise RuntimeError(error_msg) from e
 
     def _get_path(self) -> str:
         """

kiln_ai/tools/rag_tools.py

@@ -0,0 +1,164 @@
+from functools import cached_property
+from typing import Any, Dict, List, TypedDict
+
+from pydantic import BaseModel
+
+from kiln_ai.adapters.embedding.base_embedding_adapter import BaseEmbeddingAdapter
+from kiln_ai.adapters.embedding.embedding_registry import embedding_adapter_from_type
+from kiln_ai.adapters.vector_store.base_vector_store_adapter import (
+    BaseVectorStoreAdapter,
+    SearchResult,
+    VectorStoreQuery,
+)
+from kiln_ai.adapters.vector_store.vector_store_registry import (
+    vector_store_adapter_for_config,
+)
+from kiln_ai.datamodel.embedding import EmbeddingConfig
+from kiln_ai.datamodel.project import Project
+from kiln_ai.datamodel.rag import RagConfig
+from kiln_ai.datamodel.tool_id import ToolId
+from kiln_ai.datamodel.vector_store import VectorStoreConfig, VectorStoreType
+from kiln_ai.tools.base_tool import KilnToolInterface, ToolCallContext
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+
+class ChunkContext(BaseModel):
+    metadata: dict
+    text: str
+
+    def serialize(self) -> str:
+        metadata_str = ", ".join([f"{k}: {v}" for k, v in self.metadata.items()])
+        return f"[{metadata_str}]\n{self.text}\n\n"
+
+
+def format_search_results(search_results: List[SearchResult]) -> str:
+    results: List[ChunkContext] = []
+    for search_result in search_results:
+        results.append(
+            ChunkContext(
+                metadata={
+                    "document_id": search_result.document_id,
+                    "chunk_idx": search_result.chunk_idx,
+                },
+                text=search_result.chunk_text,
+            )
+        )
+    return "\n=========\n".join([result.serialize() for result in results])
+
+
+class RagParams(TypedDict):
+    query: str
+
+
+class RagTool(KilnToolInterface):
+    """
+    A tool that searches the vector store and returns the most relevant chunks.
+    """
+
+    def __init__(self, tool_id: str, rag_config: RagConfig):
+        self._id = tool_id
+        self._name = rag_config.tool_name
+        self._description = rag_config.tool_description
+        self._parameters_schema = {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The search query",
+                },
+            },
+            "required": ["query"],
+        }
+        self._rag_config = rag_config
+        vector_store_config = VectorStoreConfig.from_id_and_parent_path(
+            str(self._rag_config.vector_store_config_id), self.project.path
+        )
+        if vector_store_config is None:
+            raise ValueError(
+                f"Vector store config not found: {self._rag_config.vector_store_config_id}"
+            )
+        self._vector_store_config = vector_store_config
+        self._vector_store_adapter: BaseVectorStoreAdapter | None = None
+
+    @cached_property
+    def project(self) -> Project:
+        project = self._rag_config.parent_project()
+        if project is None:
+            raise ValueError(f"RAG config {self._rag_config.id} has no project")
+        return project
+
+    @cached_property
+    def embedding(
+        self,
+    ) -> tuple[EmbeddingConfig, BaseEmbeddingAdapter]:
+        embedding_config = EmbeddingConfig.from_id_and_parent_path(
+            str(self._rag_config.embedding_config_id), self.project.path
+        )
+        if embedding_config is None:
+            raise ValueError(
+                f"Embedding config not found: {self._rag_config.embedding_config_id}"
+            )
+        return embedding_config, embedding_adapter_from_type(embedding_config)
+
+    async def vector_store(
+        self,
+    ) -> BaseVectorStoreAdapter:
+        if self._vector_store_adapter is None:
+            self._vector_store_adapter = await vector_store_adapter_for_config(
+                vector_store_config=self._vector_store_config,
+                rag_config=self._rag_config,
+            )
+        return self._vector_store_adapter
+
+    async def id(self) -> ToolId:
+        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]:
+        """Return the OpenAI-compatible tool definition for this tool."""
+        return {
+            "type": "function",
+            "function": {
+                "name": await self.name(),
+                "description": await self.description(),
+                "parameters": self._parameters_schema,
+            },
+        }
+
+    async def run(self, context: ToolCallContext | None = None, **kwargs) -> str:
+        kwargs = RagParams(**kwargs)
+        query = kwargs["query"]
+
+        _, embedding_adapter = self.embedding
+
+        vector_store_adapter = await self.vector_store()
+        store_query = VectorStoreQuery(
+            query_embedding=None,
+            query_string=query,
+        )
+
+        match self._vector_store_config.store_type:
+            case VectorStoreType.LANCE_DB_HYBRID | VectorStoreType.LANCE_DB_VECTOR:
+                is_vector_query = True
+            case VectorStoreType.LANCE_DB_FTS:
+                is_vector_query = False
+            case _:
+                raise_exhaustive_enum_error(self._vector_store_config.store_type)
+
+        if is_vector_query:
+            query_embedding_result = await embedding_adapter.generate_embeddings(
+                [query]
+            )
+            if len(query_embedding_result.embeddings) == 0:
+                raise ValueError("No embeddings generated")
+            store_query.query_embedding = query_embedding_result.embeddings[0].vector
+
+        search_results = await vector_store_adapter.search(store_query)
+        search_results_as_text = format_search_results(search_results)
+
+        return search_results_as_text