hud-python 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

hud/mcp_agent/base.py

@@ -1,723 +0,0 @@
-"""Base MCP Agent implementation."""
-
-from __future__ import annotations
-
-import asyncio
-import logging
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
-
-import mcp.types as types
-from mcp_use import MCPClient
-
-if TYPE_CHECKING:
-    from hud.task import Task
-
-logger = logging.getLogger(__name__)
-
-
-class BaseMCPAgent(ABC):
-    """
-    Base class for MCP-enabled agents.
-
-    This class provides the foundation for agents that interact with MCP servers,
-    handling tool discovery and filtering while leaving provider-specific
-    implementation details to subclasses.
-    """
-
-    def __init__(
-        self,
-        client: MCPClient | None = None,
-        allowed_tools: list[str] | None = None,
-        disallowed_tools: list[str] | None = None,
-        initial_screenshot: bool = False,
-        max_screenshot_history: int = 3,
-        append_tool_system_prompt: bool = True,
-        custom_system_prompt: str | None = None,
-        lifecycle_tools: dict[str, str] | None = None,
-    ) -> None:
-        """
-        Initialize the base MCP agent.
-
-        Args:
-            client: MCPClient instance for server connections
-            allowed_tools: List of tool names to allow (None = all tools)
-            disallowed_tools: List of tool names to disallow
-            initial_screenshot: Whether to capture screenshot before first prompt
-            max_screenshot_history: Maximum number of screenshots to keep in context
-            append_tool_system_prompt: Whether to append available tools to system prompt
-            custom_system_prompt: Custom system prompt to use
-            lifecycle_tools: Dict mapping lifecycle phases to tool names. Default:
-                {
-                    "setup": "setup",      # Setup phase tool
-                    "evaluate": "evaluate"  # Evaluation phase tool
-                }
-        """
-        self.client = client
-        self.allowed_tools = allowed_tools
-        self.disallowed_tools = disallowed_tools or []
-        self.initial_screenshot = initial_screenshot
-        self.max_screenshot_history = max_screenshot_history
-        self.append_tool_system_prompt = append_tool_system_prompt
-        self.custom_system_prompt = custom_system_prompt
-
-        # Default lifecycle tool mapping
-        default_lifecycle = {"setup": "setup", "evaluate": "evaluate"}
-        self.lifecycle_tools = {**default_lifecycle, **(lifecycle_tools or {})}
-
-        self._available_tools: list[types.Tool] = []
-        self._tool_map: dict[str, tuple[str, types.Tool]] = {}
-        self._sessions: dict[str, Any] = {}
-
-        if client is None:
-            self.client = MCPClient()
-
-    async def initialize(self) -> None:
-        """Initialize the agent and discover available tools."""
-        # Get existing sessions or create new ones
-        if self.client is None:
-            raise ValueError("Client is not initialized")
-
-        sessions = self.client.get_all_active_sessions()
-
-        if not sessions:
-            logger.info("No active sessions found, creating new ones...")
-            sessions = await self.client.create_all_sessions()
-
-        self._sessions = sessions
-
-        # Discover tools from all servers
-        self._available_tools = []
-        self._tool_map = {}
-
-        for server_name, session in 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:
-                    raise ValueError("Client session is not initialized")
-
-                tools_result = await session.connector.client_session.list_tools()
-
-                # Log all tools before filtering
-                logger.info(
-                    "Tools from '%s' (pre-filter): %s",
-                    server_name,
-                    [tool.name for tool in tools_result.tools],
-                )
-
-                for tool in tools_result.tools:
-                    # Always include lifecycle tools for framework use
-                    is_lifecycle_tool = tool.name in self.lifecycle_tools.values()
-
-                    # Apply filtering (but always allow lifecycle tools)
-                    if not is_lifecycle_tool:
-                        if self.allowed_tools and tool.name not in self.allowed_tools:
-                            continue
-                        if tool.name in self.disallowed_tools:
-                            continue
-
-                    self._available_tools.append(tool)
-                    # Store tool with server reference for execution
-                    self._tool_map[tool.name] = (server_name, tool)
-
-            except Exception as e:
-                logger.error("Failed to list tools from server %s: %s", server_name, e)
-
-        # Separate lifecycle tools from regular tools for clearer logging
-        lifecycle_tool_names = list(self.lifecycle_tools.values())
-        regular_tools = [
-            t.name for t in self._available_tools if t.name not in lifecycle_tool_names
-        ]
-        lifecycle_tools_found = [
-            t.name for t in self._available_tools if t.name in lifecycle_tool_names
-        ]
-
-        logger.info(
-            "Agent initialized with %s tools (%s regular, %s lifecycle)",
-            len(self._available_tools),
-            len(regular_tools),
-            len(lifecycle_tools_found),
-        )
-        if regular_tools:
-            logger.info("Regular tools: %s", regular_tools)
-        if lifecycle_tools_found:
-            logger.info("Lifecycle tools: %s", lifecycle_tools_found)
-
-    def get_available_tools(self) -> list[types.Tool]:
-        """Get list of available MCP tools for LLM use (excludes lifecycle tools)."""
-        lifecycle_tool_names = list(self.lifecycle_tools.values())
-        return [tool for tool in self._available_tools if tool.name not in lifecycle_tool_names]
-
-    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, Any]:
-        """Get active MCP sessions."""
-        return self._sessions
-
-    def get_tools_by_server(self) -> dict[str, list[types.Tool]]:
-        """Get tools grouped by server name."""
-        tools_by_server = {}
-        for server_name, tool in self._tool_map.values():
-            if server_name not in tools_by_server:
-                tools_by_server[server_name] = []
-            tools_by_server[server_name].append(tool)
-        return tools_by_server
-
-    def get_tools_by_connector(self) -> dict[Any, list[types.Tool]]:
-        """Get tools grouped by connector instance."""
-        tools_by_connector = {}
-        for server_name, tool in self._tool_map.values():
-            session = self._sessions[server_name]
-            connector = session.connector
-
-            if connector not in tools_by_connector:
-                tools_by_connector[connector] = []
-            tools_by_connector[connector].append(tool)
-        return tools_by_connector
-
-    def get_system_prompt(self) -> str:
-        """Generate system prompt with optional tool information."""
-        base_prompt = self.custom_system_prompt or "You are a helpful assistant."
-
-        if self.append_tool_system_prompt and self._available_tools:
-            tool_descriptions = []
-            for tool in self._available_tools:
-                desc = f"- {tool.name}: {tool.description}"
-                if tool.inputSchema:
-                    desc += f" (parameters: {tool.inputSchema})"
-                tool_descriptions.append(desc)
-
-            tools_prompt = "\n\nYou have access to the following tools:\n" + "\n".join(
-                tool_descriptions
-            )
-            return base_prompt + tools_prompt
-
-        return base_prompt
-
-    async def call_tool(self, tool_call: dict[str, Any]) -> types.CallToolResult:
-        """
-        Call a tool through the MCP client.
-
-        Args:
-            tool_call: Dict with 'name' and optional 'arguments' keys
-
-        Returns:
-            The raw MCP CallToolResult
-        """
-        tool_name = tool_call.get("name")
-        if not tool_name:
-            raise ValueError("Tool call must have a 'name' field")
-
-        tool_args = tool_call.get("arguments", {})
-
-        if tool_name not in self._tool_map:
-            raise ValueError(f"Tool '{tool_name}' not found or not allowed")
-
-        if self.client is None:
-            raise ValueError("Client is not initialized")
-
-        server_name, tool = self._tool_map[tool_name]
-        session = self.client.get_session(server_name)
-
-        logger.info(
-            "Calling tool '%s' on server '%s' with args: %s",
-            tool_name,
-            server_name,
-            tool_args,
-        )
-        if session.connector.client_session is None:
-            raise ValueError("Client session is not initialized")
-
-        result = await session.connector.client_session.call_tool(tool_name, tool_args)
-
-        # Log result for debugging
-        if result.isError:
-            logger.error("Tool '%s' returned error: %s", tool_name, result.content)
-        else:
-            logger.debug("Tool '%s' completed successfully", tool_name)
-
-        return result
-
-    def has_computer_tools(self) -> bool:
-        """Check if any computer control tools are available."""
-        computer_tools = {"computer", "computer_anthropic", "computer_openai", "screenshot"}
-        return any(tool.name in computer_tools for tool in self._available_tools)
-
-    def get_tool_schemas(self) -> list[dict]:
-        """Get tool schemas in a format suitable for the model."""
-        schemas = []
-        for tool in self._available_tools:
-            # Filter out lifecycle tools from LLM conversation
-            if tool.name in self.lifecycle_tools.values():
-                continue
-
-            schema = {
-                "name": tool.name,
-                "description": tool.description,
-            }
-            if tool.inputSchema:
-                schema["parameters"] = tool.inputSchema
-            schemas.append(schema)
-        return schemas
-
-    async def capture_screenshot(self) -> str | None:
-        """Capture a screenshot using available tools."""
-        if not self.has_computer_tools():
-            return None
-
-        # Try different screenshot tools
-        for tool_name in [
-            "computer",
-            "screenshot",
-            "computer_anthropic",
-            "computer_openai",
-            "anthropic_computer",
-            "openai_computer",
-        ]:
-            if tool_name in self._tool_map:
-                try:
-                    # Different tools have different APIs
-                    if tool_name == "computer_openai":
-                        tool_call = {"name": tool_name, "arguments": {"type": "screenshot"}}
-                    else:
-                        tool_call = {"name": tool_name, "arguments": {"action": "screenshot"}}
-
-                    result = await self.call_tool(tool_call)
-
-                    # Extract screenshot from result
-                    for content in result.content:
-                        if isinstance(content, types.ImageContent):
-                            logger.info("Captured screenshot")
-                            return content.data
-
-                except Exception as e:
-                    logger.warning("Failed to capture screenshot with %s: %s", tool_name, e)
-
-        return None
-
-    def process_tool_results(self, tool_results: list[dict[str, Any]]) -> dict[str, Any]:
-        """
-        Process tool results into a standardized format.
-
-        Returns a dict with:
-        - text: Combined text output from all tools
-        - screenshot: Latest screenshot if any tool returned one
-        - errors: List of any errors encountered
-        - results: List of (tool_name, content_blocks) tuples for provider-specific formatting
-        """
-        text_parts = []
-        latest_screenshot = None
-        errors = []
-        results = []
-
-        for tool_result in tool_results:
-            tool_name = tool_result["tool_name"]
-            content_blocks = []
-
-            if tool_result.get("error"):
-                error_msg = f"{tool_name}: {tool_result.get('error_message', 'Unknown error')}"
-                errors.append(error_msg)
-                text_parts.append(f"Error - {error_msg}")
-                content_blocks.append(
-                    {
-                        "type": "error",
-                        "text": tool_result.get("error_message", "Unknown error"),
-                    }
-                )
-            else:
-                result = tool_result["result"]
-                if result.isError:
-                    # Extract error from content
-                    error_text = "Tool execution failed"
-                    for content in result.content:
-                        if isinstance(content, types.TextContent):
-                            error_text = content.text
-                            break
-                    error_msg = f"{tool_name}: {error_text}"
-                    errors.append(error_msg)
-                    text_parts.append(f"Error - {error_msg}")
-                    content_blocks.append(
-                        {
-                            "type": "error",
-                            "text": error_text,
-                        }
-                    )
-                else:
-                    # Process success content
-                    tool_output = []
-                    for content in result.content:
-                        if isinstance(content, types.TextContent):
-                            tool_output.append(content.text)
-                            content_blocks.append(
-                                {
-                                    "type": "text",
-                                    "text": content.text,
-                                }
-                            )
-                        elif isinstance(content, types.ImageContent):
-                            # Keep the latest screenshot
-                            latest_screenshot = content.data
-                            content_blocks.append(
-                                {
-                                    "type": "image",
-                                    "data": content.data,
-                                }
-                            )
-
-                    if tool_output:
-                        text_parts.append(f"{tool_name}: " + " ".join(tool_output))
-
-            results.append((tool_name, content_blocks))
-
-        return {
-            "text": "\n".join(text_parts) if text_parts else "No output from tools",
-            "screenshot": latest_screenshot,
-            "errors": errors,
-            "results": results,  # List of (tool_name, content_blocks) for provider-specific use
-        }
-
-    async def run(
-        self, prompt_or_task: str | Task, max_steps: int = 10, conversation_mode: bool = False
-    ) -> dict[str, Any]:
-        """
-        Run the agent with the given prompt or task.
-
-        Args:
-            prompt_or_task: Either a string prompt for simple execution or a Task object
-            max_steps: Maximum number of steps
-            conversation_mode: If True, continue even when model returns text without tool calls
-
-        Returns:
-            For string prompts: The final response string
-            For Task objects: Evaluation result dict with 'reward', 'done', 'info' keys
-        """
-        # Import here to avoid circular imports
-        from hud.task import Task
-
-        if not self._available_tools:
-            await self.initialize()
-
-        # Handle Task objects with full lifecycle
-        if isinstance(prompt_or_task, Task):
-            return await self._run_task(prompt_or_task, max_steps)
-
-        # Handle simple string prompts (existing behavior)
-        elif isinstance(prompt_or_task, str):
-            return await self._run_prompt(prompt_or_task, max_steps, conversation_mode)
-
-        else:
-            raise TypeError(f"prompt_or_task must be str or Task, got {type(prompt_or_task)}")
-
-    async def _run_task(self, task: Task, max_steps: int = 10) -> dict[str, Any]:
-        """
-        Execute a task with setup and evaluate phases.
-
-        Args:
-            task: Task object with prompt, setup, and evaluate configs
-            max_steps: Maximum steps for task execution
-
-        Returns:
-            Evaluation result dict with 'reward', 'done', 'info' keys
-        """
-        try:
-            # Setup phase
-            if task.setup is not None:
-                setup_tool = self.lifecycle_tools.get("setup", "setup")
-                await self._call_tool_safe(setup_tool, task.setup)
-
-            # Execute the task prompt
-            await self._run_prompt(task.prompt, max_steps, conversation_mode=False)
-
-            # Evaluate phase
-            if task.evaluate is not None:
-                evaluate_tool = self.lifecycle_tools.get("evaluate", "evaluate")
-                eval_result = await self._call_tool_safe(evaluate_tool, task.evaluate)
-
-                # Return evaluation result if it's properly formatted
-                if (
-                    isinstance(eval_result, dict)
-                    and "reward" in eval_result
-                    and "done" in eval_result
-                ):
-                    return eval_result
-                elif isinstance(eval_result, dict) and "grade" in eval_result:
-                    return {
-                        "reward": eval_result.get("grade", 0.0),
-                        "done": True,
-                        "info": {
-                            "error": eval_result.get("error"),
-                            "logs": eval_result.get("logs", ""),
-                            "original_result": eval_result,
-                        },
-                    }
-                else:
-                    # Fallback for invalid evaluation format
-                    return {
-                        "reward": 0.0,
-                        "done": True,
-                        "info": {"error": "Invalid evaluation result", "eval_result": eval_result},
-                    }
-            else:
-                # No evaluation - assume success
-                return {
-                    "reward": 0.0,
-                    "done": True,
-                    "info": {"message": "Task completed (no evaluation specified)"},
-                }
-
-        except Exception as e:
-            return {"reward": 0.0, "done": True, "info": {"error": str(e)}}
-
-    async def _call_tool_safe(self, tool_name: str, arguments: Any) -> Any:
-        """
-        Safely call a tool and return its result.
-
-        Args:
-            tool_name: Name of the tool to call
-            arguments: Arguments to pass to the tool (config from task)
-
-        Returns:
-            Tool result or None if tool not available/failed
-        """
-        try:
-            if tool_name in self._tool_map:
-                tool_call = {"name": tool_name, "arguments": arguments}
-                result = await self.call_tool(tool_call)
-
-                if result.isError:
-                    logger.error("Tool %s returned error: %s", tool_name, result.content)
-                    return {"error": result.content}
-                else:
-                    # Extract content from MCP result
-                    if hasattr(result, "content") and result.content:
-                        if len(result.content) == 1:
-                            content_item = result.content[0]
-                            # Check if content_item is a text type
-                            if hasattr(content_item, "text") and hasattr(content_item, "type"):
-                                if getattr(content_item, "type", None) == "text":
-                                    # Try to parse as JSON if it looks like structured data
-                                    text = content_item.text  # type: ignore[reportAttributeAccessIssue]
-                                    if text.strip().startswith("{") and text.strip().endswith("}"):
-                                        try:
-                                            import json
-
-                                            return json.loads(text)
-                                        except json.JSONDecodeError:
-                                            return text
-                                    return text
-                            else:
-                                return content_item
-                        else:
-                            return result.content
-                    return result
-            else:
-                logger.warning("Tool %s not available", tool_name)
-                return None
-        except Exception as e:
-            logger.error("Failed to call tool %s: %s", tool_name, e)
-            return {"error": str(e)}
-
-    async def _run_prompt(
-        self,
-        prompt: str,
-        max_steps: int = 10,
-        conversation_mode: bool = False,
-    ) -> dict[str, Any]:
-        """
-        Run the agent with the given prompt.
-
-        Args:
-            prompt: The task to complete
-            max_steps: Maximum number of steps
-            conversation_mode: If True, continue even when model returns text without tool calls
-
-        Returns:
-            The final response or result
-        """
-        try:
-            latest_screenshot = None
-            if self.initial_screenshot:
-                latest_screenshot = await self.capture_screenshot()
-
-            messages = await self.create_initial_messages(prompt, latest_screenshot)
-
-            step = 0
-            while step < max_steps:
-                step += 1
-                logger.info("step %s/%s", step, max_steps)
-
-                try:
-                    response = await self.get_model_response(messages, step)
-
-                    # Log the model's response
-                    logger.info("Model response - Content: %s", response.get("content", ""))
-                    logger.info(
-                        "Model response - Tool calls: %s",
-                        [tc.get("name") for tc in response.get("tool_calls", [])],
-                    )
-                    logger.info("Model response - Done: %s", response.get("done", False))
-
-                    # Check if we should stop
-                    if response.get("done", False) and not conversation_mode:
-                        return response.get("content", "Task completed")
-
-                    tool_calls = response.get("tool_calls", [])
-                    if not tool_calls:
-                        if conversation_mode:
-                            # In conversation mode, if model responds without tools,
-                            # show the response and get user input
-                            model_response = response.get("content", "")
-                            if model_response:
-                                print(f"\n🤖 Agent: {model_response}")  # noqa: T201
-                                user_input = input("\n👤 You: ").strip()
-                                if user_input.lower() in ["exit", "quit", "bye"]:
-                                    return {
-                                        "done": True,
-                                        "reward": 0.0,
-                                        "info": {"message": "Conversation ended by user."},
-                                    }
-                                # Add user's response to the conversation
-                                # This needs to be handled by subclass-specific format
-                                user_message = await self.create_user_message(user_input)
-                                messages.append(user_message)
-                                continue
-                            else:
-                                # No content and no tools - something went wrong
-                                return {
-                                    "done": False,
-                                    "reward": 0.0,
-                                    "info": {"message": "No response generated"},
-                                }
-                        else:
-                            # In task mode, no tool calls means we're done
-                            logger.info("In task mode with no tool calls - stopping execution")
-                            logger.info(
-                                "Final message: %s",
-                                response.get("content", "No response generated"),
-                            )
-                            return {
-                                "done": True,
-                                "reward": 0.0,
-                                "info": {
-                                    "message": response.get("content", "No response generated"),
-                                },
-                            }
-
-                    # Execute tool calls
-                    tool_results = []
-                    for tool_call in tool_calls:
-                        if not tool_call.get("name"):
-                            continue
-                        try:
-                            result = await self.call_tool(tool_call)
-                            tool_results.append(
-                                {
-                                    "tool_name": tool_call["name"],
-                                    "result": result,
-                                    "error": False,
-                                }
-                            )
-                        except Exception as e:
-                            logger.error("Tool execution failed: %s", e)
-                            tool_results.append(
-                                {
-                                    "tool_name": tool_call["name"],
-                                    "error": True,
-                                    "error_message": str(e),
-                                }
-                            )
-
-                    # Process results
-                    processed_results = self.process_tool_results(tool_results)
-
-                    # Update screenshot if we got a new one
-                    if processed_results["screenshot"]:
-                        latest_screenshot = processed_results["screenshot"]
-
-                    # Format tool results for the model
-                    tool_messages = await self.format_tool_results(
-                        processed_results,
-                        response.get("tool_calls", []),
-                    )
-                    messages.extend(tool_messages)
-
-                except Exception as e:
-                    logger.error("Model call failed: %s", e)
-                    return {"done": False, "reward": 0.0, "info": {"message": f"Error: {e}"}}
-
-            return {"done": True, "reward": 0.0, "info": {"message": "Task completed"}}
-
-        except KeyboardInterrupt:
-            logger.info("Agent execution interrupted by user")
-            return {
-                "done": False,
-                "reward": 0.0,
-                "info": {"message": "Execution interrupted by user (Ctrl+C)"},
-            }
-        except asyncio.CancelledError:
-            logger.info("Agent execution cancelled")
-            return {"done": False, "reward": 0.0, "info": {"message": "Execution cancelled"}}
-
-    @abstractmethod
-    async def create_initial_messages(self, prompt: str, screenshot: str | None) -> list[Any]:
-        """
-        Create initial messages for the conversation.
-
-        Args:
-            prompt: The user's prompt
-            screenshot: Optional initial screenshot
-
-        Returns:
-            List of messages in provider-specific format
-        """
-
-    @abstractmethod
-    async def get_model_response(self, messages: list[Any], step: int) -> dict[str, Any]:
-        """
-        Get response from the model including any tool calls.
-
-        Args:
-            messages: List of messages in provider-specific format
-            step: Current step number
-
-        Returns:
-            Dict with 'content', 'tool_calls', and 'done' keys
-        """
-
-    @abstractmethod
-    async def format_tool_results(
-        self, processed_results: dict[str, Any], tool_calls: list[dict[str, Any]]
-    ) -> list[Any]:
-        """
-        Format tool results into messages for the model.
-
-        Args:
-            processed_results: Processed tool results from process_tool_results
-            tool_calls: Original tool calls from the model
-
-        Returns:
-            List of formatted messages to append to conversation
-        """
-        raise NotImplementedError
-
-    async def create_user_message(self, text: str) -> Any:
-        """
-        Create a user message in the format expected by the model.
-
-        Default implementation for text-only messages.
-        Subclasses can override for specific formats.
-
-        Args:
-            text: User's text input
-
-        Returns:
-            Formatted user message
-        """
-        return {"role": "user", "content": text}