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

kiln_ai/tools/mcp_session_manager.py

@@ -2,6 +2,7 @@ import logging
 import os
 import subprocess
 import sys
+import tempfile
 from contextlib import asynccontextmanager
 from datetime import timedelta
 from typing import AsyncGenerator
@@ -19,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:
     """
@@ -51,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)
 
@@ -164,35 +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, read_timeout_seconds=timedelta(seconds=8)
-                ) 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

@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import Any, Dict, List
+from typing import List, TypedDict
 
 from pydantic import BaseModel
 
@@ -18,7 +18,11 @@ 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
+from kiln_ai.tools.base_tool import (
+    KilnToolInterface,
+    ToolCallContext,
+    ToolCallDefinition,
+)
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
@@ -46,6 +50,10 @@ def format_search_results(search_results: List[SearchResult]) -> str:
     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.
@@ -115,7 +123,7 @@ class RagTool(KilnToolInterface):
     async def description(self) -> str:
         return self._description
 
-    async def toolcall_definition(self) -> Dict[str, Any]:
+    async def toolcall_definition(self) -> ToolCallDefinition:
         """Return the OpenAI-compatible tool definition for this tool."""
         return {
             "type": "function",
@@ -126,7 +134,10 @@ class RagTool(KilnToolInterface):
             },
         }
 
-    async def run(self, query: str) -> str:
+    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()
@@ -152,6 +163,6 @@ class RagTool(KilnToolInterface):
             store_query.query_embedding = query_embedding_result.embeddings[0].vector
 
         search_results = await vector_store_adapter.search(store_query)
-        context = format_search_results(search_results)
+        search_results_as_text = format_search_results(search_results)
 
-        return context
+        return search_results_as_text

kiln_ai/tools/test_kiln_task_tool.py

@@ -0,0 +1,527 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from kiln_ai.datamodel import Task
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName, StructuredOutputMode
+from kiln_ai.datamodel.external_tool_server import ExternalToolServer, ToolServerType
+from kiln_ai.datamodel.run_config import RunConfigProperties
+from kiln_ai.datamodel.task import TaskRunConfig
+from kiln_ai.datamodel.task_output import DataSource, DataSourceType
+from kiln_ai.tools.base_tool import ToolCallContext
+from kiln_ai.tools.kiln_task_tool import KilnTaskTool, KilnTaskToolResult
+
+
+class TestKilnTaskToolResult:
+    """Test the KilnTaskToolResult class."""
+
+    def test_init(self):
+        """Test KilnTaskToolResult initialization."""
+        output = "test output"
+        kiln_task_tool_data = "project_id:::tool_id:::task_id:::run_id"
+
+        result = KilnTaskToolResult(output, kiln_task_tool_data)
+
+        assert result.output == output
+        assert result.kiln_task_tool_data == kiln_task_tool_data
+
+    def test_init_with_empty_strings(self):
+        """Test KilnTaskToolResult initialization with empty strings."""
+        result = KilnTaskToolResult("", "")
+
+        assert result.output == ""
+        assert result.kiln_task_tool_data == ""
+
+
+class TestKilnTaskTool:
+    """Test the KilnTaskTool class."""
+
+    @pytest.fixture
+    def mock_external_tool_server(self):
+        """Create a mock ExternalToolServer for testing."""
+        return ExternalToolServer(
+            name="test_tool",
+            type=ToolServerType.kiln_task,
+            description="Test Kiln task tool",
+            properties={
+                "name": "test_task_tool",
+                "description": "A test task tool",
+                "task_id": "test_task_123",
+                "run_config_id": "test_config_456",
+                "is_archived": False,
+            },
+        )
+
+    @pytest.fixture
+    def mock_task(self):
+        """Create a mock Task for testing."""
+        task = MagicMock(spec=Task)
+        task.id = "test_task_123"
+        task.input_json_schema = None
+        task.input_schema.return_value = None
+        task.run_configs.return_value = []
+        return task
+
+    @pytest.fixture
+    def mock_run_config(self):
+        """Create a mock TaskRunConfig for testing."""
+        run_config = MagicMock(spec=TaskRunConfig)
+        run_config.id = "test_config_456"
+        run_config.run_config_properties = {
+            "model_name": "gpt-4",
+            "model_provider_name": "openai",
+            "prompt_id": "simple_prompt_builder",
+            "structured_output_mode": "default",
+        }
+        return run_config
+
+    @pytest.fixture
+    def mock_context(self):
+        """Create a mock ToolCallContext for testing."""
+        context = MagicMock(spec=ToolCallContext)
+        context.allow_saving = True
+        return context
+
+    @pytest.fixture
+    def kiln_task_tool(self, mock_external_tool_server):
+        """Create a KilnTaskTool instance for testing."""
+        return KilnTaskTool(
+            project_id="test_project",
+            tool_id="test_tool_id",
+            data_model=mock_external_tool_server,
+        )
+
+    @pytest.mark.asyncio
+    async def test_init(self, mock_external_tool_server):
+        """Test KilnTaskTool initialization."""
+        tool = KilnTaskTool(
+            project_id="test_project",
+            tool_id="test_tool_id",
+            data_model=mock_external_tool_server,
+        )
+
+        assert tool._project_id == "test_project"
+        assert tool._tool_id == "test_tool_id"
+        assert tool._tool_server_model == mock_external_tool_server
+        assert tool._name == "test_task_tool"
+        assert tool._description == "A test task tool"
+        assert tool._task_id == "test_task_123"
+        assert tool._run_config_id == "test_config_456"
+
+    @pytest.mark.asyncio
+    async def test_init_with_missing_properties(self):
+        """Test KilnTaskTool initialization with missing properties."""
+        # Create a server with minimal required properties
+        server = ExternalToolServer(
+            name="test_tool",
+            type=ToolServerType.kiln_task,
+            description="Test tool",
+            properties={
+                "name": "minimal_tool",
+                "description": "",
+                "task_id": "",
+                "run_config_id": "",
+                "is_archived": False,
+            },
+        )
+
+        tool = KilnTaskTool(
+            project_id="test_project",
+            tool_id="test_tool_id",
+            data_model=server,
+        )
+
+        assert tool._name == "minimal_tool"
+        assert tool._description == ""
+        assert tool._task_id == ""
+        assert tool._run_config_id == ""
+
+    @pytest.mark.asyncio
+    async def test_id(self, kiln_task_tool):
+        """Test the id method."""
+        result = await kiln_task_tool.id()
+        assert result == "test_tool_id"
+
+    @pytest.mark.asyncio
+    async def test_name(self, kiln_task_tool):
+        """Test the name method."""
+        result = await kiln_task_tool.name()
+        assert result == "test_task_tool"
+
+    @pytest.mark.asyncio
+    async def test_description(self, kiln_task_tool):
+        """Test the description method."""
+        result = await kiln_task_tool.description()
+        assert result == "A test task tool"
+
+    @pytest.mark.asyncio
+    async def test_toolcall_definition(self, kiln_task_tool):
+        """Test the toolcall_definition method."""
+        # Mock the parameters_schema property directly
+        kiln_task_tool.parameters_schema = {"type": "object"}
+
+        definition = await kiln_task_tool.toolcall_definition()
+
+        assert definition["type"] == "function"
+        assert definition["function"]["name"] == "test_task_tool"
+        assert definition["function"]["description"] == "A test task tool"
+        assert definition["function"]["parameters"] == {"type": "object"}
+
+    @pytest.mark.asyncio
+    async def test_run_with_plaintext_input(
+        self, kiln_task_tool, mock_context, mock_task, mock_run_config
+    ):
+        """Test the run method with plaintext input."""
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+        kiln_task_tool._run_config = mock_run_config
+
+        with (
+            patch(
+                "kiln_ai.adapters.adapter_registry.adapter_for_task"
+            ) as mock_adapter_for_task,
+            patch(
+                "kiln_ai.adapters.model_adapters.base_adapter.AdapterConfig"
+            ) as mock_adapter_config,
+        ):
+            # Mock adapter and task run
+            mock_adapter = AsyncMock()
+            mock_adapter_for_task.return_value = mock_adapter
+
+            mock_task_run = MagicMock()
+            mock_task_run.id = "run_789"
+            mock_task_run.output.output = "Task completed successfully"
+            mock_adapter.invoke.return_value = mock_task_run
+
+            # Test with plaintext input
+            result = await kiln_task_tool.run(context=mock_context, input="test input")
+
+            # Verify adapter was created correctly
+            mock_adapter_for_task.assert_called_once_with(
+                mock_task,
+                run_config_properties={
+                    "model_name": "gpt-4",
+                    "model_provider_name": "openai",
+                    "prompt_id": "simple_prompt_builder",
+                    "structured_output_mode": "default",
+                },
+                base_adapter_config=mock_adapter_config.return_value,
+            )
+
+            # Verify adapter config
+            mock_adapter_config.assert_called_once_with(
+                allow_saving=True,
+                default_tags=["tool_call"],
+            )
+
+            # Verify adapter invoke was called
+            mock_adapter.invoke.assert_called_once_with(
+                "test input",
+                input_source=DataSource(
+                    type=DataSourceType.tool_call,
+                    run_config=RunConfigProperties(
+                        model_name="gpt-4",
+                        model_provider_name=ModelProviderName.openai,
+                        prompt_id="simple_prompt_builder",
+                        structured_output_mode=StructuredOutputMode.default,
+                    ),
+                ),
+            )
+
+            # Verify result
+            assert isinstance(result, KilnTaskToolResult)
+            assert result.output == "Task completed successfully"
+            assert (
+                result.kiln_task_tool_data
+                == "test_project:::test_tool_id:::test_task_123:::run_789"
+            )
+
+    @pytest.mark.asyncio
+    async def test_run_with_structured_input(
+        self, kiln_task_tool, mock_context, mock_task, mock_run_config
+    ):
+        """Test the run method with structured input."""
+        # Setup task with JSON schema
+        mock_task.input_json_schema = {
+            "type": "object",
+            "properties": {"param1": {"type": "string"}},
+        }
+
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+        kiln_task_tool._run_config = mock_run_config
+
+        with patch(
+            "kiln_ai.adapters.adapter_registry.adapter_for_task"
+        ) as mock_adapter_for_task:
+            # Mock adapter and task run
+            mock_adapter = AsyncMock()
+            mock_adapter_for_task.return_value = mock_adapter
+
+            mock_task_run = MagicMock()
+            mock_task_run.id = "run_789"
+            mock_task_run.output.output = "Structured task completed"
+            mock_adapter.invoke.return_value = mock_task_run
+
+            # Test with structured input
+            result = await kiln_task_tool.run(
+                context=mock_context, param1="value1", param2="value2"
+            )
+
+            # Verify adapter invoke was called with kwargs
+            mock_adapter.invoke.assert_called_once_with(
+                {"param1": "value1", "param2": "value2"},
+                input_source=DataSource(
+                    type=DataSourceType.tool_call,
+                    run_config=RunConfigProperties(
+                        model_name="gpt-4",
+                        model_provider_name=ModelProviderName.openai,
+                        prompt_id="simple_prompt_builder",
+                        structured_output_mode=StructuredOutputMode.default,
+                    ),
+                ),
+            )
+
+            # Verify result
+            assert result.output == "Structured task completed"
+
+    @pytest.mark.asyncio
+    async def test_run_without_context(self, kiln_task_tool):
+        """Test the run method without context raises ValueError."""
+        with pytest.raises(
+            ValueError, match="Context is required for running a KilnTaskTool"
+        ):
+            await kiln_task_tool.run(input="test input")
+
+    @pytest.mark.asyncio
+    async def test_run_plaintext_missing_input(
+        self, kiln_task_tool, mock_context, mock_task
+    ):
+        """Test the run method with plaintext task but missing input parameter."""
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+
+        with pytest.raises(ValueError, match="Input not found in kwargs"):
+            await kiln_task_tool.run(context=mock_context, wrong_param="value")
+
+    @pytest.mark.asyncio
+    async def test_task_property_project_not_found(self, kiln_task_tool):
+        """Test _task property when project is not found."""
+        with patch("kiln_ai.tools.kiln_task_tool.project_from_id", return_value=None):
+            with pytest.raises(ValueError, match="Project not found: test_project"):
+                _ = kiln_task_tool._task
+
+    @pytest.mark.asyncio
+    async def test_task_property_task_not_found(self, kiln_task_tool):
+        """Test _task property when task is not found."""
+        mock_project = MagicMock()
+        mock_project.path = "/test/path"
+
+        with (
+            patch(
+                "kiln_ai.tools.kiln_task_tool.project_from_id",
+                return_value=mock_project,
+            ),
+            patch(
+                "kiln_ai.tools.kiln_task_tool.Task.from_id_and_parent_path",
+                return_value=None,
+            ),
+        ):
+            with pytest.raises(
+                ValueError,
+                match="Task not found: test_task_123 in project test_project",
+            ):
+                _ = kiln_task_tool._task
+
+    @pytest.mark.asyncio
+    async def test_task_property_success(self, kiln_task_tool, mock_task):
+        """Test _task property when task is found successfully."""
+        mock_project = MagicMock()
+        mock_project.path = "/test/path"
+
+        with (
+            patch(
+                "kiln_ai.tools.kiln_task_tool.project_from_id",
+                return_value=mock_project,
+            ),
+            patch(
+                "kiln_ai.tools.kiln_task_tool.Task.from_id_and_parent_path",
+                return_value=mock_task,
+            ),
+        ):
+            result = kiln_task_tool._task
+            assert result == mock_task
+
+    @pytest.mark.asyncio
+    async def test_run_config_property_not_found(self, kiln_task_tool, mock_task):
+        """Test _run_config property when run config is not found."""
+        mock_task.run_configs.return_value = []
+
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+
+        with pytest.raises(
+            ValueError,
+            match="Task run config not found: test_config_456 for task test_task_123 in project test_project",
+        ):
+            _ = kiln_task_tool._run_config
+
+    @pytest.mark.asyncio
+    async def test_run_config_property_success(
+        self, kiln_task_tool, mock_task, mock_run_config
+    ):
+        """Test _run_config property when run config is found successfully."""
+        mock_task.run_configs.return_value = [mock_run_config]
+
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+
+        result = kiln_task_tool._run_config
+        assert result == mock_run_config
+
+    @pytest.mark.asyncio
+    async def test_parameters_schema_with_json_schema(self, kiln_task_tool, mock_task):
+        """Test parameters_schema property with JSON schema."""
+        expected_schema = {
+            "type": "object",
+            "properties": {"param": {"type": "string"}},
+        }
+        mock_task.input_json_schema = expected_schema
+        mock_task.input_schema.return_value = expected_schema
+
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+
+        result = kiln_task_tool.parameters_schema
+        assert result == expected_schema
+
+    @pytest.mark.asyncio
+    async def test_parameters_schema_plaintext(self, kiln_task_tool, mock_task):
+        """Test parameters_schema property for plaintext task."""
+        mock_task.input_json_schema = None
+
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+
+        result = kiln_task_tool.parameters_schema
+
+        expected = {
+            "type": "object",
+            "properties": {
+                "input": {
+                    "type": "string",
+                    "description": "Plaintext input for the tool.",
+                }
+            },
+            "required": ["input"],
+        }
+        assert result == expected
+
+    @pytest.mark.asyncio
+    async def test_parameters_schema_none_raises_error(self, kiln_task_tool, mock_task):
+        """Test parameters_schema property when schema is None raises ValueError."""
+        # Set up a task with JSON schema but input_schema returns None
+        mock_task.input_json_schema = {
+            "type": "object",
+            "properties": {"param": {"type": "string"}},
+        }
+        mock_task.input_schema.return_value = None
+
+        # Setup mocks - directly assign the task to bypass cached property
+        kiln_task_tool._task = mock_task
+
+        with pytest.raises(
+            ValueError,
+            match="Failed to create parameters schema for tool_id test_tool_id",
+        ):
+            _ = kiln_task_tool.parameters_schema
+
+    @pytest.mark.asyncio
+    async def test_cached_properties(self, kiln_task_tool, mock_task, mock_run_config):
+        """Test that cached properties work correctly."""
+        mock_project = MagicMock()
+        mock_project.path = "/test/path"
+
+        with (
+            patch(
+                "kiln_ai.tools.kiln_task_tool.project_from_id",
+                return_value=mock_project,
+            ),
+            patch(
+                "kiln_ai.tools.kiln_task_tool.Task.from_id_and_parent_path",
+                return_value=mock_task,
+            ),
+        ):
+            # First access should call the methods
+            task1 = kiln_task_tool._task
+            task2 = kiln_task_tool._task
+
+            # Should be the same object (cached)
+            assert task1 is task2
+
+            # Verify the methods were called only once
+            assert mock_project is not None  # project_from_id was called
+            # Task.from_id_and_parent_path should have been called once
+            with patch(
+                "kiln_ai.tools.kiln_task_tool.Task.from_id_and_parent_path"
+            ) as mock_from_id:
+                mock_from_id.return_value = mock_task
+                _ = kiln_task_tool._task
+                # Should not be called again due to caching
+                mock_from_id.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_run_with_adapter_exception(
+        self, kiln_task_tool, mock_context, mock_task, mock_run_config
+    ):
+        """Test the run method when adapter raises an exception."""
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+        kiln_task_tool._run_config = mock_run_config
+
+        with patch(
+            "kiln_ai.adapters.adapter_registry.adapter_for_task"
+        ) as mock_adapter_for_task:
+            # Mock adapter to raise an exception
+            mock_adapter = AsyncMock()
+            mock_adapter.invoke.side_effect = Exception("Adapter failed")
+            mock_adapter_for_task.return_value = mock_adapter
+
+            with pytest.raises(Exception, match="Adapter failed"):
+                await kiln_task_tool.run(context=mock_context, input="test input")
+
+    @pytest.mark.asyncio
+    async def test_run_with_different_allow_saving(
+        self, kiln_task_tool, mock_task, mock_run_config
+    ):
+        """Test the run method with different allow_saving values."""
+        mock_context_false = MagicMock(spec=ToolCallContext)
+        mock_context_false.allow_saving = False
+
+        # Setup mocks
+        kiln_task_tool._task = mock_task
+        kiln_task_tool._run_config = mock_run_config
+
+        with (
+            patch(
+                "kiln_ai.adapters.adapter_registry.adapter_for_task"
+            ) as mock_adapter_for_task,
+            patch(
+                "kiln_ai.adapters.model_adapters.base_adapter.AdapterConfig"
+            ) as mock_adapter_config,
+        ):
+            mock_adapter = AsyncMock()
+            mock_adapter_for_task.return_value = mock_adapter
+
+            mock_task_run = MagicMock()
+            mock_task_run.id = "run_789"
+            mock_task_run.output.output = "Task completed"
+            mock_adapter.invoke.return_value = mock_task_run
+
+            await kiln_task_tool.run(context=mock_context_false, input="test input")
+
+            # Verify adapter config was called with allow_saving=False
+            mock_adapter_config.assert_called_once_with(
+                allow_saving=False,
+                default_tags=["tool_call"],
+            )