agent/core/tools.py

"""
Tool system for the agent
Provides ToolSpec and ToolRouter for managing both built-in and MCP tools
"""

import subprocess
import warnings
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Optional

from fastmcp import Client
from lmnr import observe
from mcp.types import EmbeddedResource, ImageContent, TextContent

from agent.config import MCPServerConfig

# Suppress aiohttp deprecation warning
warnings.filterwarnings(
    "ignore", category=DeprecationWarning, module="aiohttp.connector"
)


def convert_mcp_content_to_string(content: list) -> str:
    """
    Convert MCP content blocks to a string format compatible with LLM messages.

    Based on FastMCP documentation, content can be:
    - TextContent: has .text field
    - ImageContent: has .data and .mimeType fields
    - EmbeddedResource: has .resource field with .text or .blob

    Args:
        content: List of MCP content blocks

    Returns:
        String representation of the content suitable for LLM consumption
    """
    if not content:
        return ""

    parts = []
    for item in content:
        if isinstance(item, TextContent):
            # Extract text from TextContent blocks
            parts.append(item.text)
        elif isinstance(item, ImageContent):
            # TODO: Handle images
            # For images, include a description with MIME type
            parts.append(f"[Image: {item.mimeType}]")
        elif isinstance(item, EmbeddedResource):
            # TODO: Handle embedded resources
            # For embedded resources, try to extract text
            resource = item.resource
            if hasattr(resource, "text") and resource.text:
                parts.append(resource.text)
            elif hasattr(resource, "blob") and resource.blob:
                parts.append(
                    f"[Binary data: {resource.mimeType if hasattr(resource, 'mimeType') else 'unknown'}]"
                )
            else:
                parts.append(
                    f"[Resource: {resource.uri if hasattr(resource, 'uri') else 'unknown'}]"
                )
        else:
            # Fallback: try to convert to string
            parts.append(str(item))

    return "\n".join(parts)


@dataclass
class ToolSpec:
    """Tool specification for LLM"""

    name: str
    description: str
    parameters: dict[str, Any]
    handler: Optional[Callable[[dict[str, Any]], Awaitable[tuple[str, bool]]]] = None


class ToolRouter:
    """
    Routes tool calls to appropriate handlers.
    Based on codex-rs/core/src/tools/router.rs
    """

    def __init__(self, mcp_servers: dict[str, MCPServerConfig]):
        self.tools: dict[str, ToolSpec] = {}
        self.mcp_servers: dict[str, dict[str, Any]] = {}

        for tool in create_builtin_tools():
            self.register_tool(tool)

        if mcp_servers:
            mcp_servers_payload = {}
            for name, server in mcp_servers.items():
                mcp_servers_payload[name] = server.model_dump()
            self.mcp_client = Client({"mcpServers": mcp_servers_payload})
        self._mcp_initialized = False

    def register_tool(self, tool: ToolSpec) -> None:
        self.tools[tool.name] = tool

    async def register_mcp_tools(self) -> None:
        tools = await self.mcp_client.list_tools()
        for tool in tools:
            self.register_tool(
                ToolSpec(
                    name=tool.name,
                    description=tool.description,
                    parameters=tool.inputSchema,
                    handler=None,
                )
            )

    def get_tool_specs_for_llm(self) -> list[dict[str, Any]]:
        """Get tool specifications in OpenAI format"""
        specs = []
        for tool in self.tools.values():
            specs.append(
                {
                    "type": "function",
                    "function": {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.parameters,
                    },
                }
            )
        return specs

    async def __aenter__(self) -> "ToolRouter":
        if self.mcp_client is not None:
            await self.mcp_client.__aenter__()
            await self.mcp_client.initialize()
            await self.register_mcp_tools()
            self._mcp_initialized = True
        print(f"MCP initialized: {self._mcp_initialized}")
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        if self.mcp_client is not None:
            await self.mcp_client.__aexit__(exc_type, exc, tb)
            self._mcp_initialized = False

    @observe(name="call_tool")
    async def call_tool(
        self, tool_name: str, arguments: dict[str, Any]
    ) -> tuple[str, bool]:
        """
        Call a tool and return (output_string, success_bool).

        For MCP tools, converts the CallToolResult content blocks to a string.
        For built-in tools, calls their handler directly.
        """
        # Check if this is a built-in tool with a handler
        tool = self.tools.get(tool_name)
        if tool and tool.handler:
            return await tool.handler(arguments)

        # Otherwise, use MCP client
        if self._mcp_initialized:
            result = await self.mcp_client.call_tool(tool_name, arguments)
            # Convert MCP content blocks to string
            output = convert_mcp_content_to_string(result.content)
            return output, not result.is_error

        return "MCP client not initialized", False


# ============================================================================
# BUILT-IN TOOL HANDLERS
# ============================================================================


async def bash_handler(arguments: dict[str, Any]) -> tuple[str, bool]:
    """Execute bash command"""
    try:
        command = arguments.get("command", "")
        result = subprocess.run(
            command, shell=True, capture_output=True, text=True, timeout=30
        )
        output = result.stdout if result.returncode == 0 else result.stderr
        success = result.returncode == 0
        return output, success
    except Exception as e:
        return f"Error: {str(e)}", False


async def read_file_handler(arguments: dict[str, Any]) -> tuple[str, bool]:
    """Read file contents"""
    try:
        path = arguments.get("path", "")
        with open(path, "r") as f:
            content = f.read()
        return content, True
    except Exception as e:
        return f"Error reading file: {str(e)}", False


async def write_file_handler(arguments: dict[str, Any]) -> tuple[str, bool]:
    """Write to file"""
    try:
        path = arguments.get("path", "")
        content = arguments.get("content", "")
        with open(path, "w") as f:
            f.write(content)
        return f"Successfully wrote to {path}", True
    except Exception as e:
        return f"Error writing file: {str(e)}", False


def create_builtin_tools() -> list[ToolSpec]:
    """Create built-in tool specifications"""
    return [
        ToolSpec(
            name="bash",
            description="Execute a bash command and return its output",
            parameters={
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The bash command to execute",
                    }
                },
                "required": ["command"],
            },
            handler=bash_handler,
        ),
        ToolSpec(
            name="read_file",
            description="Read the contents of a file",
            parameters={
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Path to the file to read",
                    }
                },
                "required": ["path"],
            },
            handler=read_file_handler,
        ),
        ToolSpec(
            name="write_file",
            description="Write content to a file",
            parameters={
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Path to the file to write",
                    },
                    "content": {
                        "type": "string",
                        "description": "Content to write to the file",
                    },
                },
                "required": ["path", "content"],
            },
            handler=write_file_handler,
        ),
    ]