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

kiln_ai/tools/built_in_tools/math_tools.py

@@ -0,0 +1,124 @@
+from typing import Union
+
+from kiln_ai.datamodel.tool_id import KilnBuiltInToolId
+from kiln_ai.tools.base_tool import KilnTool
+
+
+class AddTool(KilnTool):
+    """
+    A concrete tool that adds two numbers together.
+    Demonstrates how to use the KilnTool base class.
+    """
+
+    def __init__(self):
+        parameters_schema = {
+            "type": "object",
+            "properties": {
+                "a": {"type": "number", "description": "The first number to add"},
+                "b": {"type": "number", "description": "The second number to add"},
+            },
+            "required": ["a", "b"],
+        }
+
+        super().__init__(
+            tool_id=KilnBuiltInToolId.ADD_NUMBERS,
+            name="add",
+            description="Add two numbers together and return the result",
+            parameters_schema=parameters_schema,
+        )
+
+    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+        """Add two numbers and return the result."""
+        return str(a + b)
+
+
+class SubtractTool(KilnTool):
+    """
+    A concrete tool that subtracts two numbers.
+    """
+
+    def __init__(self):
+        parameters_schema = {
+            "type": "object",
+            "properties": {
+                "a": {"type": "number", "description": "The first number (minuend)"},
+                "b": {
+                    "type": "number",
+                    "description": "The second number to subtract (subtrahend)",
+                },
+            },
+            "required": ["a", "b"],
+        }
+
+        super().__init__(
+            tool_id=KilnBuiltInToolId.SUBTRACT_NUMBERS,
+            name="subtract",
+            description="Subtract the second number from the first number and return the result",
+            parameters_schema=parameters_schema,
+        )
+
+    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+        """Subtract b from a and return the result."""
+        return str(a - b)
+
+
+class MultiplyTool(KilnTool):
+    """
+    A concrete tool that multiplies two numbers together.
+    """
+
+    def __init__(self):
+        parameters_schema = {
+            "type": "object",
+            "properties": {
+                "a": {"type": "number", "description": "The first number to multiply"},
+                "b": {"type": "number", "description": "The second number to multiply"},
+            },
+            "required": ["a", "b"],
+        }
+
+        super().__init__(
+            tool_id=KilnBuiltInToolId.MULTIPLY_NUMBERS,
+            name="multiply",
+            description="Multiply two numbers together and return the result",
+            parameters_schema=parameters_schema,
+        )
+
+    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+        """Multiply two numbers and return the result."""
+        return str(a * b)
+
+
+class DivideTool(KilnTool):
+    """
+    A concrete tool that divides two numbers.
+    """
+
+    def __init__(self):
+        parameters_schema = {
+            "type": "object",
+            "properties": {
+                "a": {
+                    "type": "number",
+                    "description": "The dividend (number to be divided)",
+                },
+                "b": {
+                    "type": "number",
+                    "description": "The divisor (number to divide by)",
+                },
+            },
+            "required": ["a", "b"],
+        }
+
+        super().__init__(
+            tool_id=KilnBuiltInToolId.DIVIDE_NUMBERS,
+            name="divide",
+            description="Divide the first number by the second number and return the result",
+            parameters_schema=parameters_schema,
+        )
+
+    async def run(self, 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")
+        return str(a / b)

kiln_ai/tools/built_in_tools/test_math_tools.py

@@ -0,0 +1,204 @@
+import pytest
+
+from kiln_ai.datamodel.tool_id import KilnBuiltInToolId
+from kiln_ai.tools.built_in_tools.math_tools import (
+    AddTool,
+    DivideTool,
+    MultiplyTool,
+    SubtractTool,
+)
+
+
+class TestAddTool:
+    """Test the AddTool class."""
+
+    async def test_init(self):
+        """Test AddTool initialization."""
+        tool = AddTool()
+        assert await tool.id() == KilnBuiltInToolId.ADD_NUMBERS
+        assert await tool.name() == "add"
+        assert (
+            await tool.description() == "Add two numbers together and return the result"
+        )
+
+    async def test_toolcall_definition(self):
+        """Test AddTool toolcall definition structure."""
+        tool = AddTool()
+        definition = await tool.toolcall_definition()
+
+        assert definition["type"] == "function"
+        assert definition["function"]["name"] == "add"
+        assert (
+            definition["function"]["description"]
+            == "Add two numbers together and return the result"
+        )
+        assert "properties" in definition["function"]["parameters"]
+        assert "a" in definition["function"]["parameters"]["properties"]
+        assert "b" in definition["function"]["parameters"]["properties"]
+
+    @pytest.mark.parametrize(
+        "a, b, expected",
+        [
+            (1, 2, "3"),
+            (0, 0, "0"),
+            (-1, 1, "0"),
+            (2.5, 3.5, "6.0"),
+            (-2.5, -3.5, "-6.0"),
+            (100, 200, "300"),
+        ],
+    )
+    async def test_run_various_inputs(self, a, b, expected):
+        """Test AddTool run method with various inputs."""
+        tool = AddTool()
+        result = await tool.run(a=a, b=b)
+        assert result == expected
+
+
+class TestSubtractTool:
+    """Test the SubtractTool class."""
+
+    async def test_init(self):
+        """Test SubtractTool initialization."""
+        tool = SubtractTool()
+        assert await tool.id() == KilnBuiltInToolId.SUBTRACT_NUMBERS
+        assert await tool.name() == "subtract"
+        assert (
+            await tool.description()
+            == "Subtract the second number from the first number and return the result"
+        )
+
+    async def test_toolcall_definition(self):
+        """Test SubtractTool toolcall definition structure."""
+        tool = SubtractTool()
+        definition = await tool.toolcall_definition()
+
+        assert definition["type"] == "function"
+        assert definition["function"]["name"] == "subtract"
+        assert (
+            definition["function"]["description"]
+            == "Subtract the second number from the first number and return the result"
+        )
+        assert "properties" in definition["function"]["parameters"]
+        assert "a" in definition["function"]["parameters"]["properties"]
+        assert "b" in definition["function"]["parameters"]["properties"]
+
+    @pytest.mark.parametrize(
+        "a, b, expected",
+        [
+            (5, 3, "2"),
+            (0, 0, "0"),
+            (1, -1, "2"),
+            (5.5, 2.5, "3.0"),
+            (-2.5, -3.5, "1.0"),
+            (100, 200, "-100"),
+        ],
+    )
+    async def test_run_various_inputs(self, a, b, expected):
+        """Test SubtractTool run method with various inputs."""
+        tool = SubtractTool()
+        result = await tool.run(a=a, b=b)
+        assert result == expected
+
+
+class TestMultiplyTool:
+    """Test the MultiplyTool class."""
+
+    async def test_init(self):
+        """Test MultiplyTool initialization."""
+        tool = MultiplyTool()
+        assert await tool.id() == KilnBuiltInToolId.MULTIPLY_NUMBERS
+        assert await tool.name() == "multiply"
+        assert (
+            await tool.description()
+            == "Multiply two numbers together and return the result"
+        )
+
+    async def test_toolcall_definition(self):
+        """Test MultiplyTool toolcall definition structure."""
+        tool = MultiplyTool()
+        definition = await tool.toolcall_definition()
+
+        assert definition["type"] == "function"
+        assert definition["function"]["name"] == "multiply"
+        assert (
+            definition["function"]["description"]
+            == "Multiply two numbers together and return the result"
+        )
+        assert "properties" in definition["function"]["parameters"]
+        assert "a" in definition["function"]["parameters"]["properties"]
+        assert "b" in definition["function"]["parameters"]["properties"]
+
+    @pytest.mark.parametrize(
+        "a, b, expected",
+        [
+            (2, 3, "6"),
+            (0, 5, "0"),
+            (-2, 3, "-6"),
+            (2.5, 4, "10.0"),
+            (-2.5, -4, "10.0"),
+            (1, 1, "1"),
+        ],
+    )
+    async def test_run_various_inputs(self, a, b, expected):
+        """Test MultiplyTool run method with various inputs."""
+        tool = MultiplyTool()
+        result = await tool.run(a=a, b=b)
+        assert result == expected
+
+
+class TestDivideTool:
+    """Test the DivideTool class."""
+
+    async def test_init(self):
+        """Test DivideTool initialization."""
+        tool = DivideTool()
+        assert await tool.id() == KilnBuiltInToolId.DIVIDE_NUMBERS
+        assert await tool.name() == "divide"
+        assert (
+            await tool.description()
+            == "Divide the first number by the second number and return the result"
+        )
+
+    async def test_toolcall_definition(self):
+        """Test DivideTool toolcall definition structure."""
+        tool = DivideTool()
+        definition = await tool.toolcall_definition()
+
+        assert definition["type"] == "function"
+        assert definition["function"]["name"] == "divide"
+        assert (
+            definition["function"]["description"]
+            == "Divide the first number by the second number and return the result"
+        )
+        assert "properties" in definition["function"]["parameters"]
+        assert "a" in definition["function"]["parameters"]["properties"]
+        assert "b" in definition["function"]["parameters"]["properties"]
+
+    @pytest.mark.parametrize(
+        "a, b, expected",
+        [
+            (6, 2, "3.0"),
+            (1, 1, "1.0"),
+            (-6, 2, "-3.0"),
+            (7.5, 2.5, "3.0"),
+            (-10, -2, "5.0"),
+            (0, 5, "0.0"),
+        ],
+    )
+    async def test_run_various_inputs(self, a, b, expected):
+        """Test DivideTool run method with various inputs."""
+        tool = DivideTool()
+        result = await tool.run(a=a, b=b)
+        assert result == expected
+
+    async def test_divide_by_zero(self):
+        """Test that division by zero raises ZeroDivisionError."""
+        tool = DivideTool()
+        with pytest.raises(ZeroDivisionError, match="Cannot divide by zero"):
+            await tool.run(a=5, b=0)
+
+    async def test_divide_zero_by_zero(self):
+        """Test that zero divided by zero raises ZeroDivisionError."""
+        tool = DivideTool()
+        with pytest.raises(ZeroDivisionError, match="Cannot divide by zero"):
+            await tool.run(a=0, b=0)

kiln_ai/tools/mcp_server_tool.py

@@ -0,0 +1,95 @@
+from typing import Any, Dict
+
+from mcp.types import CallToolResult, TextContent
+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.mcp_session_manager import MCPSessionManager
+
+
+class MCPServerTool(KilnToolInterface):
+    def __init__(self, data_model: ExternalToolServer, name: str):
+        self._tool_id = f"{MCP_REMOTE_TOOL_ID_PREFIX}{data_model.id}::{name}"
+        self._tool_server_model = data_model
+        self._name = name
+        self._tool: MCPTool | None = None
+
+    async def id(self) -> ToolId:
+        return self._tool_id
+
+    async def name(self) -> str:
+        return self._name
+
+    async def description(self) -> str:
+        await self._load_tool_properties()
+        return self._description
+
+    async def toolcall_definition(self) -> Dict[str, Any]:
+        """Generate OpenAI-compatible tool definition."""
+        await self._load_tool_properties()
+        return {
+            "type": "function",
+            "function": {
+                "name": await self.name(),
+                "description": await self.description(),
+                "parameters": self._parameters_schema,
+            },
+        }
+
+    async def run(self, **kwargs) -> Any:
+        result = await self._call_tool(**kwargs)
+
+        if result.isError:
+            raise ValueError(
+                f"Tool {await self.name()} returned an error: {result.content}"
+            )
+
+        if not result.content:
+            raise ValueError("Tool returned no content")
+
+        # raise error if the first block is not a text block
+        if not isinstance(result.content[0], TextContent):
+            raise ValueError("First block must be a text block")
+
+        # raise error if there is more than one content block
+        if len(result.content) > 1:
+            raise ValueError("Tool returned multiple content blocks, expected one")
+
+        return result.content[0].text
+
+    #  Call the MCP Tool
+    async def _call_tool(self, **kwargs) -> CallToolResult:
+        async with MCPSessionManager.shared().mcp_client(
+            self._tool_server_model
+        ) as session:
+            result = await session.call_tool(
+                name=await self.name(),
+                arguments=kwargs,
+            )
+            return result
+
+    async def _load_tool_properties(self):
+        if self._tool is not None:
+            return
+
+        tool = await self._get_tool(self._name)
+        self._tool = tool
+        self._description = tool.description or "N/A"
+        self._parameters_schema = tool.inputSchema or {
+            "type": "object",
+            "properties": {},
+        }
+
+    #  Get the MCP Tool from the server
+    async def _get_tool(self, tool_name: str) -> MCPTool:
+        async with MCPSessionManager.shared().mcp_client(
+            self._tool_server_model
+        ) as session:
+            tools = await session.list_tools()
+
+        tool = next((tool for tool in tools.tools if tool.name == tool_name), None)
+        if tool is None:
+            raise ValueError(f"Tool {tool_name} not found")
+        return tool

kiln_ai/tools/mcp_session_manager.py

@@ -0,0 +1,246 @@
+import logging
+import os
+import subprocess
+import sys
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import AsyncGenerator
+
+import httpx
+from mcp import StdioServerParameters
+from mcp.client.session import ClientSession
+from mcp.client.stdio import stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.exceptions import McpError
+
+from kiln_ai.datamodel.external_tool_server import ExternalToolServer, ToolServerType
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+logger = logging.getLogger(__name__)
+
+
+class MCPSessionManager:
+    """
+    This class is a singleton that manages MCP sessions for remote MCP servers.
+    """
+
+    _shared_instance = None
+
+    def __init__(self):
+        self._shell_path = None
+
+    @classmethod
+    def shared(cls):
+        if cls._shared_instance is None:
+            cls._shared_instance = cls()
+        return cls._shared_instance
+
+    @asynccontextmanager
+    async def mcp_client(
+        self,
+        tool_server: ExternalToolServer,
+    ) -> AsyncGenerator[
+        ClientSession,
+        None,
+    ]:
+        match tool_server.type:
+            case ToolServerType.remote_mcp:
+                async with self._create_remote_mcp_session(tool_server) as session:
+                    yield session
+            case ToolServerType.local_mcp:
+                async with self._create_local_mcp_session(tool_server) as session:
+                    yield session
+            case _:
+                raise_exhaustive_enum_error(tool_server.type)
+
+    def _extract_first_exception(
+        self, exception: Exception, target_type: type | tuple[type, ...]
+    ) -> Exception | None:
+        """
+        Extract first relevant exception from ExceptionGroup or handle direct exceptions
+        """
+        # Check if the exception itself is of the target type
+        if isinstance(exception, target_type):
+            return exception
+
+        # Handle ExceptionGroup
+        if hasattr(exception, "exceptions"):
+            exceptions_attr = getattr(exception, "exceptions", None)
+            if exceptions_attr:
+                for nested_exc in exceptions_attr:
+                    result = self._extract_first_exception(nested_exc, target_type)
+                    if result:
+                        return result
+
+        return None
+
+    @asynccontextmanager
+    async def _create_remote_mcp_session(
+        self,
+        tool_server: ExternalToolServer,
+    ) -> AsyncGenerator[ClientSession, None]:
+        """
+        Create a session for a remote MCP server.
+        """
+        # Make sure the server_url is set
+        server_url = tool_server.properties.get("server_url")
+        if not server_url:
+            raise ValueError("server_url is required")
+
+        # Make a copy of the headers to avoid modifying the original object
+        headers = tool_server.properties.get("headers", {}).copy()
+
+        # Retrieve secret headers from configuration and merge with regular headers
+        secret_headers, _ = tool_server.retrieve_secrets()
+        headers.update(secret_headers)
+
+        try:
+            async with streamablehttp_client(server_url, headers=headers) as (
+                read_stream,
+                write_stream,
+                _,
+            ):
+                # Create a session using the client streams
+                async with ClientSession(read_stream, write_stream) as session:
+                    await session.initialize()
+                    yield session
+        except Exception as e:
+            # Handle HTTP errors with user-friendly messages
+
+            # Check for HTTPStatusError
+            http_error = self._extract_first_exception(e, httpx.HTTPStatusError)
+            if http_error and isinstance(http_error, httpx.HTTPStatusError):
+                raise ValueError(
+                    f"The MCP server rejected the request. "
+                    f"Status {http_error.response.status_code}. "
+                    f"Response from server:\n{http_error.response.reason_phrase}"
+                )
+
+            # Check for connection errors
+            connection_error_types = (ConnectionError, OSError, httpx.RequestError)
+            connection_error = self._extract_first_exception(e, connection_error_types)
+            if connection_error and isinstance(
+                connection_error, connection_error_types
+            ):
+                raise RuntimeError(
+                    f"Unable to connect to MCP server. Please verify the configurations are correct, the server is running, and your network connection is working. Original error: {connection_error}"
+                ) from e
+
+            # If no known error types found, re-raise the original exception
+            raise RuntimeError(
+                f"Failed to connect to the MCP Server. Check the server's docs for troubleshooting. Original error: {e}"
+            ) from e
+
+    @asynccontextmanager
+    async def _create_local_mcp_session(
+        self,
+        tool_server: ExternalToolServer,
+    ) -> AsyncGenerator[ClientSession, None]:
+        """
+        Create a session for a local MCP server.
+        """
+        command = tool_server.properties.get("command")
+        if not command:
+            raise ValueError(
+                "Attempted to start local MCP server, but no command was provided"
+            )
+
+        args = tool_server.properties.get("args", [])
+        if not isinstance(args, list):
+            raise ValueError(
+                "Attempted to start local MCP server, but args is not a list of strings"
+            )
+
+        # Make a copy of the env_vars to avoid modifying the original object
+        env_vars = tool_server.properties.get("env_vars", {}).copy()
+
+        # Retrieve secret environment variables from configuration and merge with regular env_vars
+        secret_env_vars, _ = tool_server.retrieve_secrets()
+        env_vars.update(secret_env_vars)
+
+        # Set PATH, only if not explicitly set during MCP tool setup
+        if "PATH" not in env_vars:
+            env_vars["PATH"] = self._get_path()
+
+        # Set the server parameters
+        server_params = StdioServerParameters(
+            command=command,
+            args=args,
+            env=env_vars,
+        )
+
+        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):
+        """
+        Raise a ValueError 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
+
+    def _get_path(self) -> str:
+        """
+        Builds a PATH environment variable. From environment, Kiln Config, and loading rc files.
+        """
+
+        # If the user sets a custom MCP path, use only it. This also functions as a way to disable the shell path loading.
+        custom_mcp_path = Config.shared().get_value("custom_mcp_path")
+        if custom_mcp_path is not None:
+            return custom_mcp_path
+        else:
+            return self.get_shell_path()
+
+    def get_shell_path(self) -> str:
+        # Windows has a global PATH, so we don't need to source rc files
+        if sys.platform in ("win32", "Windows"):
+            return os.environ.get("PATH", "")
+
+        # Cache
+        if self._shell_path is not None:
+            return self._shell_path
+
+        # Attempt to get shell PATH from preferred shell, which will source rc files, run scripts like `brew shellenv`, etc.
+        shell_path = None
+        try:
+            shell = os.environ.get("SHELL", "/bin/bash")
+            # Use -l (login) flag to source ~/.profile, ~/.bash_profile, ~/.zprofile, etc.
+            result = subprocess.run(
+                [shell, "-l", "-c", "echo $PATH"],
+                capture_output=True,
+                text=True,
+                timeout=3,
+            )
+            if result.returncode == 0:
+                shell_path = result.stdout.strip()
+        except (subprocess.TimeoutExpired, subprocess.SubprocessError, Exception) as e:
+            logger.error(f"Shell path exception details: {e}")
+
+        # Fallback to environment PATH
+        if shell_path is None:
+            logger.error(
+                "Error getting shell PATH. You may not be able to find MCP server commands like 'npx'. You can set a custom MCP path in the Kiln config file. See docs for details."
+            )
+            shell_path = os.environ.get("PATH", "")
+
+        self._shell_path = shell_path
+        return shell_path
+
+    def clear_shell_path_cache(self):
+        """Clear the cached shell path. Typically used when adding a new tool, which might have just been installed."""
+        self._shell_path = None