hud-python 0.2.10__py3-none-any.whl → 0.3.1__py3-none-any.whl

hud/mcp/claude.py

@@ -0,0 +1,321 @@
+"""Claude MCP Agent implementation."""
+
+from __future__ import annotations
+
+import copy
+import logging
+from typing import TYPE_CHECKING, Any, cast
+
+from anthropic import AsyncAnthropic, BadRequestError
+
+if TYPE_CHECKING:
+    from anthropic.types.beta import (
+        BetaCacheControlEphemeralParam,
+        BetaImageBlockParam,
+        BetaMessageParam,
+        BetaTextBlockParam,
+        BetaToolResultBlockParam,
+    )
+
+    from hud.datasets import TaskConfig
+
+import mcp.types as types
+from mcp.types import CallToolRequestParams as MCPToolCall
+from mcp.types import CallToolResult as MCPToolResult
+
+from hud.settings import settings
+
+from .base import BaseMCPAgent, ModelResponse
+
+logger = logging.getLogger(__name__)
+
+
+def base64_to_content_block(base64: str) -> BetaImageBlockParam:
+    """Convert base64 image to Claude content block."""
+    return {
+        "type": "image",
+        "source": {"type": "base64", "media_type": "image/png", "data": base64},
+    }
+
+
+def text_to_content_block(text: str) -> BetaTextBlockParam:
+    """Convert text to Claude content block."""
+    return {"type": "text", "text": text}
+
+
+def tool_use_content_block(
+    tool_use_id: str, content: list[BetaTextBlockParam | BetaImageBlockParam]
+) -> BetaToolResultBlockParam:
+    """Create tool result content block."""
+    return {"type": "tool_result", "tool_use_id": tool_use_id, "content": content}
+
+
+class ClaudeMCPAgent(BaseMCPAgent):
+    """
+    Claude agent that uses MCP servers for tool execution.
+
+    This agent uses Claude's native tool calling capabilities but executes
+    tools through MCP servers instead of direct implementation.
+    """
+
+    def __init__(
+        self,
+        model_client: AsyncAnthropic | None = None,
+        model: str = "claude-3-7-sonnet-20250219",
+        max_tokens: int = 4096,
+        display_width_px: int = 1400,
+        display_height_px: int = 850,
+        use_computer_beta: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize Claude MCP agent.
+
+        Args:
+            model_client: AsyncAnthropic client (created if not provided)
+            model: Claude model to use
+            max_tokens: Maximum tokens for response
+            display_width_px: Display width for computer use tools
+            display_height_px: Display height for computer use tools
+            use_computer_beta: Whether to use computer-use beta features
+            **kwargs: Additional arguments passed to BaseMCPAgent (including mcp_client)
+        """
+        super().__init__(**kwargs)
+
+        # Initialize client if not provided
+        if model_client is None:
+            api_key = settings.anthropic_api_key
+            if not api_key:
+                raise ValueError("Anthropic API key not found. Set ANTHROPIC_API_KEY.")
+            model_client = AsyncAnthropic(api_key=api_key)
+
+        self.anthropic_client = model_client
+        self.model = model
+        self.max_tokens = max_tokens
+        self.display_width_px = display_width_px
+        self.display_height_px = display_height_px
+        self.use_computer_beta = use_computer_beta
+
+        self.model_name = self.model
+
+        # Track mapping from Claude tool names to MCP tool names
+        self._claude_to_mcp_tool_map: dict[str, str] = {}
+
+    async def initialize(self, task: str | TaskConfig | None = None) -> None:
+        """Initialize the agent and build tool mappings."""
+        await super().initialize(task)
+        # Build tool mappings after tools are discovered
+        self._convert_tools_for_claude()
+
+    async def create_initial_messages(
+        self, prompt: str, screenshot: str | None = None
+    ) -> list[BetaMessageParam]:
+        """Create initial messages for Claude."""
+        user_content: list[BetaImageBlockParam | BetaTextBlockParam] = []
+
+        # Add prompt text
+        user_content.append(text_to_content_block(prompt))
+
+        # Add screenshot if available
+        if screenshot:
+            user_content.append(base64_to_content_block(screenshot))
+
+        # Return initial user message
+        return [
+            cast(
+                "BetaMessageParam",
+                {
+                    "role": "user",
+                    "content": user_content,
+                },
+            )
+        ]
+
+    async def get_model_response(self, messages: list[BetaMessageParam]) -> ModelResponse:
+        """Get response from Claude including any tool calls."""
+        # Get Claude tools
+        claude_tools = self._convert_tools_for_claude()
+
+        # Make API call with retry for prompt length
+        current_messages = messages.copy()
+
+        while True:
+            messages_cached = self._add_prompt_caching(current_messages)
+
+            # Build create kwargs
+            create_kwargs = {
+                "model": self.model,
+                "max_tokens": self.max_tokens,
+                "system": self.get_system_prompt(),
+                "messages": messages_cached,
+                "tools": claude_tools,
+                "tool_choice": {"type": "auto", "disable_parallel_tool_use": True},
+            }
+
+            # Add beta features if using computer tools
+            if self.use_computer_beta and any(
+                t.get("type") == "computer_20250124" for t in claude_tools
+            ):
+                create_kwargs["betas"] = ["computer-use-2025-01-24"]
+
+            try:
+                response = await self.anthropic_client.beta.messages.create(**create_kwargs)
+                break
+            except BadRequestError as e:
+                if e.message.startswith("prompt is too long"):
+                    logger.warning("Prompt too long, truncating message history")
+                    # Keep first message and last 20 messages
+                    if len(current_messages) > 21:
+                        current_messages = [current_messages[0]] + current_messages[-20:]
+                    else:
+                        raise
+                else:
+                    raise
+
+        messages.append(
+            cast(
+                "BetaMessageParam",
+                {
+                    "role": "assistant",
+                    "content": response.content,
+                },
+            )
+        )
+
+        # Process response
+        result = ModelResponse(content="", tool_calls=[], done=True)
+
+        # Extract text content and reasoning
+        text_content = ""
+        thinking_content = ""
+
+        for block in response.content:
+            if block.type == "tool_use":
+                # Map Claude tool name back to MCP tool name
+                mcp_tool_name = self._claude_to_mcp_tool_map.get(block.name, block.name)
+
+                # Create MCPToolCall object with Claude metadata as extra fields
+                # Pyright will complain but the tool class accepts extra fields
+                tool_call = MCPToolCall(
+                    name=mcp_tool_name,
+                    arguments=block.input,
+                    tool_use_id=block.id,  # type: ignore
+                    claude_name=block.name,  # type: ignore
+                )
+                result.tool_calls.append(tool_call)
+                result.done = False
+            elif block.type == "text":
+                text_content += block.text
+            elif hasattr(block, "type") and block.type == "thinking":
+                thinking_content += f"Thinking: {block.thinking}\n"
+
+        # Combine text and thinking for final content
+        if thinking_content:
+            result.content = thinking_content + text_content
+        else:
+            result.content = text_content
+
+        return result
+
+    async def format_tool_results(
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
+    ) -> list[BetaMessageParam]:
+        """Format tool results into Claude messages."""
+        # Process each tool result
+        user_content = []
+
+        for tool_call, result in zip(tool_calls, tool_results, strict=True):
+            # Extract Claude-specific metadata from extra fields
+            tool_use_id = getattr(tool_call, "tool_use_id", None)
+            if not tool_use_id:
+                logger.warning("No tool_use_id found for %s", tool_call.name)
+                continue
+
+            # Convert MCP tool results to Claude format
+            claude_blocks = []
+
+            if result.isError:
+                # Extract error message from content
+                error_msg = "Tool execution failed"
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        error_msg = content.text
+                        break
+                claude_blocks.append(text_to_content_block(f"Error: {error_msg}"))
+            else:
+                # Process success content
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        claude_blocks.append(text_to_content_block(content.text))
+                    elif isinstance(content, types.ImageContent):
+                        claude_blocks.append(base64_to_content_block(content.data))
+
+            # Add tool result
+            user_content.append(tool_use_content_block(tool_use_id, claude_blocks))
+
+        # Return as a user message containing all tool results
+        return [
+            cast(
+                "BetaMessageParam",
+                {
+                    "role": "user",
+                    "content": user_content,
+                },
+            )
+        ]
+
+    async def create_user_message(self, text: str) -> BetaMessageParam:
+        """Create a user message in Claude's format."""
+        return cast("BetaMessageParam", {"role": "user", "content": text})
+
+    def _convert_tools_for_claude(self) -> list[dict]:
+        """Convert MCP tools to Claude tool format."""
+        claude_tools = []
+        self._claude_to_mcp_tool_map = {}  # Reset mapping
+
+        for tool in self._available_tools:
+            # Special handling for computer use tools
+            if tool.name in ["computer", "computer_anthropic", "anthropic_computer"]:
+                # Use Claude's native computer use format with configurable dimensions
+                claude_tool = {
+                    "type": "computer_20250124",
+                    "name": "computer",
+                    "display_width_px": self.display_width_px,
+                    "display_height_px": self.display_height_px,
+                }
+                # Map Claude's "computer" back to the actual MCP tool name
+                self._claude_to_mcp_tool_map["computer"] = tool.name
+            elif tool.name not in self.lifecycle_tools:
+                # Convert regular tools
+                claude_tool = {
+                    "name": tool.name,
+                    "description": tool.description or f"Execute {tool.name}",
+                    "input_schema": tool.inputSchema
+                    or {
+                        "type": "object",
+                        "properties": {},
+                    },
+                }
+                # Direct mapping for non-computer tools
+                self._claude_to_mcp_tool_map[tool.name] = tool.name
+            else:
+                continue
+
+            claude_tools.append(claude_tool)
+
+        return claude_tools
+
+    def _add_prompt_caching(self, messages: list[BetaMessageParam]) -> list[BetaMessageParam]:
+        """Add prompt caching to messages."""
+        messages_cached = copy.deepcopy(messages)
+
+        # Mark last user message with cache control
+        if messages_cached and messages_cached[-1].get("role") == "user":
+            last_content = messages_cached[-1]["content"]
+            if isinstance(last_content, list):
+                for block in last_content:
+                    if block.get("type") not in ["thinking", "redacted_thinking"]:
+                        cache_control: BetaCacheControlEphemeralParam = {"type": "ephemeral"}
+                        block["cache_control"] = cache_control  # type: ignore[reportGeneralTypeIssues]
+
+        return messages_cached

hud/mcp/client.py

@@ -0,0 +1,312 @@
+"""MCP Client wrapper with automatic initialization and debugging capabilities."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any
+
+from mcp_use.client import MCPClient as MCPUseClient
+from pydantic import AnyUrl
+
+if TYPE_CHECKING:
+    from typing import Self
+
+    from mcp import types
+    from mcp_use.session import MCPSession as MCPUseSession
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClient:
+    """
+    High-level MCP client wrapper that handles initialization, tool discovery,
+    and provides debugging capabilities.
+    """
+
+    def __init__(
+        self,
+        mcp_config: dict[str, dict[str, Any]],
+        verbose: bool = False,
+    ) -> None:
+        """
+        Initialize the MCP client.
+
+        Args:
+            mcp_config: MCP server configuration dict (required)
+            verbose: Enable verbose logging of server communications
+            auto_initialize: Whether to automatically initialize on construction
+        """
+        self.verbose = verbose
+
+        # Initialize mcp_use client with proper config
+        # Use from_dict to properly initialize with config
+        config = {"mcpServers": mcp_config}
+        self._mcp_client = MCPUseClient.from_dict(config)
+
+        self._sessions: dict[str, MCPUseSession] = {}
+        self._available_tools: list[types.Tool] = []
+        self._tool_map: dict[str, tuple[str, types.Tool]] = {}
+        self._telemetry_data: dict[str, Any] = {}
+
+        # Set up verbose logging if requested
+        if self.verbose:
+            self._setup_verbose_logging()
+
+    def _setup_verbose_logging(self) -> None:
+        """Configure verbose logging for debugging."""
+        # Set MCP-related loggers to DEBUG
+        logging.getLogger("mcp").setLevel(logging.DEBUG)
+        logging.getLogger("mcp_use").setLevel(logging.DEBUG)
+        logging.getLogger("mcp.client.stdio").setLevel(logging.DEBUG)
+
+        # Add handler for server communications
+        if not any(isinstance(h, logging.StreamHandler) for h in logger.handlers):
+            handler = logging.StreamHandler()
+            handler.setFormatter(
+                logging.Formatter("[%(levelname)s] %(asctime)s - %(name)s - %(message)s")
+            )
+            logger.addHandler(handler)
+            logger.setLevel(logging.DEBUG)
+
+    async def initialize(self) -> None:
+        """Perform async initialization tasks."""
+        await self.create_sessions()
+        await self.discover_tools()
+        await self.fetch_telemetry()
+
+    async def create_sessions(self) -> dict[str, MCPUseSession]:
+        # Create all sessions at once
+        try:
+            self._sessions = await self._mcp_client.create_all_sessions()
+        except Exception as e:
+            # If session creation fails, try to get Docker logs
+            logger.error("Failed to create sessions: %s", e)
+            if self.verbose:
+                logger.info("Attempting to check Docker container status...")
+                # await self._check_docker_containers()
+            raise
+
+        # Log session details in verbose mode
+        if self.verbose and self._sessions:
+            for name, session in self._sessions.items():
+                logger.debug("  - %s: %s", name, type(session).__name__)
+
+        return self._sessions
+
+    async def discover_tools(self) -> list[types.Tool]:
+        """Discover all available tools from connected servers."""
+        logger.info("Discovering available tools...")
+
+        self._available_tools = []
+        self._tool_map = {}
+
+        for server_name, session in self._sessions.items():
+            try:
+                # Ensure session is initialized
+                if not hasattr(session, "connector") or not hasattr(
+                    session.connector, "client_session"
+                ):
+                    await session.initialize()
+
+                if session.connector.client_session is None:
+                    logger.warning("Client session not initialized for %s", server_name)
+                    continue
+
+                # List tools
+                tools_result = await session.connector.client_session.list_tools()
+
+                logger.info(
+                    "Discovered %d tools from '%s': %s",
+                    len(tools_result.tools),
+                    server_name,
+                    [tool.name for tool in tools_result.tools],
+                )
+
+                # Add to collections
+                for tool in tools_result.tools:
+                    self._available_tools.append(tool)
+                    self._tool_map[tool.name] = (server_name, tool)
+
+                # Log detailed tool info in verbose mode
+                if self.verbose:
+                    for tool in tools_result.tools:
+                        description = tool.description or ""
+                        logger.debug(
+                            "  Tool '%s': %s",
+                            tool.name,
+                            description[:100] + "..." if len(description) > 100 else description,
+                        )
+
+            except Exception as e:
+                logger.error("Error discovering tools from '%s': %s", server_name, e)
+                if self.verbose:
+                    logger.exception("Full error details:")
+
+        logger.info("Total tools discovered: %d", len(self._available_tools))
+        return self._available_tools
+
+    async def fetch_telemetry(self) -> dict[str, Any]:
+        """Fetch telemetry resource from all servers that provide it."""
+        logger.info("Fetching telemetry resources...")
+
+        for server_name, session in self._sessions.items():
+            try:
+                if not hasattr(session, "connector") or not hasattr(
+                    session.connector, "client_session"
+                ):
+                    continue
+
+                if session.connector.client_session is None:
+                    continue
+
+                # Try to read telemetry resource
+                try:
+                    result = await session.connector.client_session.read_resource(
+                        AnyUrl("telemetry://live")
+                    )
+                    if result and result.contents and len(result.contents) > 0:
+                        telemetry_data = json.loads(result.contents[0].text)  # type: ignore
+                        self._telemetry_data[server_name] = telemetry_data
+
+                        logger.info("📡 Telemetry data from server '%s':", server_name)
+                        if "live_url" in telemetry_data:
+                            logger.info("   🖥️  Live URL: %s", telemetry_data["live_url"])
+                        if "status" in telemetry_data:
+                            logger.info("   📊 Status: %s", telemetry_data["status"])
+                        if "services" in telemetry_data:
+                            logger.info("   📋 Services:")
+                            for service, status in telemetry_data["services"].items():
+                                status_icon = "✅" if status == "running" else "❌"
+                                logger.info("      %s %s: %s", status_icon, service, status)
+
+                        if self.verbose:
+                            logger.debug(
+                                "Full telemetry data:\n%s", json.dumps(telemetry_data, indent=2)
+                            )
+
+                except Exception as e:
+                    # Resource might not exist, which is fine
+                    if self.verbose:
+                        logger.debug("No telemetry resource from '%s': %s", server_name, e)
+
+            except Exception as e:
+                logger.error("Error fetching telemetry from '%s': %s", server_name, e)
+
+        return self._telemetry_data
+
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any] | None = None
+    ) -> types.CallToolResult:
+        """
+        Call a tool by name with the given arguments.
+
+        Args:
+            tool_name: Name of the tool to call
+            arguments: Tool arguments
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            ValueError: If tool not found
+        """
+        if tool_name not in self._tool_map:
+            raise ValueError(f"Tool '{tool_name}' not found")
+
+        server_name, tool = self._tool_map[tool_name]
+        session = self._sessions[server_name]
+
+        if self.verbose:
+            logger.debug(
+                "Calling tool '%s' on server '%s' with arguments: %s",
+                tool_name,
+                server_name,
+                json.dumps(arguments, indent=2) if arguments else "None",
+            )
+
+        if session.connector.client_session is None:
+            raise ValueError(f"Client session not initialized for {server_name}")
+
+        result = await session.connector.client_session.call_tool(
+            name=tool_name, arguments=arguments or {}
+        )
+
+        if self.verbose:
+            logger.debug("Tool '%s' result: %s", tool_name, result)
+
+        return result
+
+    async def read_resource(self, uri: AnyUrl) -> types.ReadResourceResult | None:
+        """
+        Read a resource by URI from any server that provides it.
+
+        Args:
+            uri: Resource URI (e.g., "telemetry://live")
+
+        Returns:
+            Resource contents or None if not found
+        """
+        for server_name, session in self._sessions.items():
+            try:
+                if not hasattr(session, "connector") or not hasattr(
+                    session.connector, "client_session"
+                ):
+                    continue
+
+                if session.connector.client_session is None:
+                    continue
+
+                result = await session.connector.client_session.read_resource(uri)
+
+                if self.verbose:
+                    logger.debug(
+                        "Successfully read resource '%s' from server '%s'", uri, server_name
+                    )
+
+                return result
+
+            except Exception as e:
+                if self.verbose:
+                    logger.debug(
+                        "Could not read resource '%s' from server '%s': %s", uri, server_name, e
+                    )
+                continue
+
+        return None
+
+    def get_available_tools(self) -> list[types.Tool]:
+        """Get list of all available tools."""
+        return self._available_tools
+
+    def get_tool_map(self) -> dict[str, tuple[str, types.Tool]]:
+        """Get mapping of tool names to (server_name, tool) tuples."""
+        return self._tool_map
+
+    def get_sessions(self) -> dict[str, MCPUseSession]:
+        """Get active MCP sessions."""
+        return self._sessions
+
+    def get_telemetry_data(self) -> dict[str, Any]:
+        """Get collected telemetry data from all servers."""
+        return self._telemetry_data
+
+    def get_all_active_sessions(self) -> dict[str, MCPUseSession]:
+        """Get all active sessions (compatibility method)."""
+        return self._sessions
+
+    async def close(self) -> None:
+        """Close all active sessions."""
+        await self._mcp_client.close_all_sessions()
+
+        self._sessions = {}
+        self._available_tools = []
+        self._tool_map = {}
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        """Async context manager exit."""
+        await self.close()