hud-python 0.4.45__py3-none-any.whl → 0.5.1__py3-none-any.whl

hud/agents/gemini.py

@@ -0,0 +1,275 @@
+"""Gemini MCP Agent implementation."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, ClassVar, cast
+
+import mcp.types as types
+from google import genai
+from google.genai import types as genai_types
+from pydantic import ConfigDict
+
+from hud.settings import settings
+from hud.types import AgentResponse, BaseAgentConfig, MCPToolCall, MCPToolResult
+from hud.utils.hud_console import HUDConsole
+from hud.utils.types import with_signature
+
+from .base import BaseCreateParams, MCPAgent
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiConfig(BaseAgentConfig):
+    """Configuration for `GeminiAgent`."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    model_name: str = "Gemini"
+    model: str = "gemini-3-pro-preview"
+    model_client: genai.Client | None = None
+    temperature: float = 1.0
+    top_p: float = 0.95
+    top_k: int = 40
+    max_output_tokens: int = 8192
+    validate_api_key: bool = True
+
+
+class GeminiCreateParams(BaseCreateParams, GeminiConfig):
+    pass
+
+
+class GeminiAgent(MCPAgent):
+    """
+    Gemini agent that uses MCP servers for tool execution.
+
+    This agent uses Gemini's native tool calling capabilities but executes
+    tools through MCP servers instead of direct implementation.
+    """
+
+    metadata: ClassVar[dict[str, Any] | None] = None
+    config_cls: ClassVar[type[BaseAgentConfig]] = GeminiConfig
+
+    @with_signature(GeminiCreateParams)
+    @classmethod
+    def create(cls, **kwargs: Any) -> GeminiAgent:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return MCPAgent.create.__func__(cls, **kwargs)  # type: ignore[return-value]
+
+    def __init__(self, params: GeminiCreateParams | None = None, **kwargs: Any) -> None:
+        super().__init__(params, **kwargs)
+        self.config: GeminiConfig
+
+        model_client = self.config.model_client
+        if model_client is None:
+            api_key = settings.gemini_api_key
+            if not api_key:
+                raise ValueError("Gemini API key not found. Set GEMINI_API_KEY.")
+            model_client = genai.Client(api_key=api_key)
+
+        if self.config.validate_api_key:
+            try:
+                list(model_client.models.list(config=genai_types.ListModelsConfig(page_size=1)))
+            except Exception as e:
+                raise ValueError(f"Gemini API key is invalid: {e}") from e
+
+        self.gemini_client = model_client
+        self.temperature = self.config.temperature
+        self.top_p = self.config.top_p
+        self.top_k = self.config.top_k
+        self.max_output_tokens = self.config.max_output_tokens
+        self.hud_console = HUDConsole(logger=logger)
+
+        # Track mapping from Gemini tool names to MCP tool names
+        self._gemini_to_mcp_tool_map: dict[str, str] = {}
+        self.gemini_tools: genai_types.ToolListUnion = []
+
+    def _on_tools_ready(self) -> None:
+        """Build Gemini-specific tool mappings after tools are discovered."""
+        self._convert_tools_for_gemini()
+
+    async def get_system_messages(self) -> list[genai_types.Content]:
+        """No system messages for Gemini because applied in get_response"""
+        return []
+
+    async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[genai_types.Content]:
+        """Format messages for Gemini."""
+        # Convert MCP content types to Gemini content types
+        gemini_parts: list[genai_types.Part] = []
+
+        for block in blocks:
+            if isinstance(block, types.TextContent):
+                gemini_parts.append(genai_types.Part(text=block.text))
+            elif isinstance(block, types.ImageContent):
+                # Convert MCP ImageContent to Gemini format
+                # Need to decode base64 string to bytes
+                import base64
+
+                image_bytes = base64.b64decode(block.data)
+                gemini_parts.append(
+                    genai_types.Part.from_bytes(data=image_bytes, mime_type=block.mimeType)
+                )
+            else:
+                # For other types, try to handle but log a warning
+                self.hud_console.log(f"Unknown content block type: {type(block)}", level="warning")
+
+        return [genai_types.Content(role="user", parts=gemini_parts)]
+
+    async def get_response(self, messages: list[genai_types.Content]) -> AgentResponse:
+        """Get response from Gemini including any tool calls."""
+        # Build generate content config
+        generate_config = genai_types.GenerateContentConfig(
+            temperature=self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            max_output_tokens=self.max_output_tokens,
+            tools=self.gemini_tools,
+            system_instruction=self.system_prompt,
+        )
+
+        # Use async API to avoid blocking the event loop
+        response = await self.gemini_client.aio.models.generate_content(
+            model=self.config.model,
+            contents=cast("Any", messages),
+            config=generate_config,
+        )
+
+        # Append assistant response (including any function_call) so that
+        # subsequent FunctionResponse messages correspond to a prior FunctionCall
+        if response.candidates and len(response.candidates) > 0 and response.candidates[0].content:
+            messages.append(response.candidates[0].content)
+
+        # Process response
+        result = AgentResponse(content="", tool_calls=[], done=True)
+        collected_tool_calls: list[MCPToolCall] = []
+
+        if not response.candidates:
+            self.hud_console.warning("Response has no candidates")
+            return result
+
+        candidate = response.candidates[0]
+
+        # Extract text content and function calls
+        text_content = ""
+        thinking_content = ""
+
+        if candidate.content and candidate.content.parts:
+            for part in candidate.content.parts:
+                if part.function_call:
+                    tool_call = self._extract_tool_call(part)
+                    if tool_call is not None:
+                        collected_tool_calls.append(tool_call)
+                elif part.thought is True and part.text:
+                    if thinking_content:
+                        thinking_content += "\n"
+                    thinking_content += part.text
+                elif part.text:
+                    text_content += part.text
+
+        # Assign collected tool calls and mark done status
+        if collected_tool_calls:
+            result.tool_calls = collected_tool_calls
+            result.done = False
+
+        result.content = text_content
+        if thinking_content:
+            result.reasoning = thinking_content
+
+        return result
+
+    def _extract_tool_call(self, part: genai_types.Part) -> MCPToolCall | None:
+        """Extract an MCPToolCall from a function call part.
+
+        Subclasses can override to customize tool call extraction (e.g., normalizing
+        computer use calls to a different schema).
+        """
+        if not part.function_call:
+            return None
+
+        func_name = part.function_call.name or ""
+        mcp_tool_name = self._gemini_to_mcp_tool_map.get(func_name, func_name)
+        raw_args = dict(part.function_call.args) if part.function_call.args else {}
+
+        return MCPToolCall(
+            name=mcp_tool_name,
+            arguments=raw_args,
+        )
+
+    async def format_tool_results(
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
+    ) -> list[genai_types.Content]:
+        """Format tool results into Gemini messages."""
+        # Process each tool result
+        function_responses = []
+
+        for tool_call, result in zip(tool_calls, tool_results, strict=True):
+            # Get the Gemini function name from metadata
+            gemini_name = getattr(tool_call, "gemini_name", tool_call.name)
+
+            # Convert MCP tool results to Gemini format
+            response_dict: dict[str, Any] = {}
+
+            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
+                response_dict["error"] = error_msg
+            else:
+                # Process success content
+                response_dict["success"] = True
+                # Add text content to response
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        response_dict["output"] = content.text
+                        break
+
+            # Create function response
+            function_response = genai_types.FunctionResponse(
+                name=gemini_name,
+                response=response_dict,
+            )
+            function_responses.append(function_response)
+
+        # Return as a user message containing all function responses
+        return [
+            genai_types.Content(
+                role="user",
+                parts=[genai_types.Part(function_response=fr) for fr in function_responses],
+            )
+        ]
+
+    async def create_user_message(self, text: str) -> genai_types.Content:
+        """Create a user message in Gemini's format."""
+        return genai_types.Content(role="user", parts=[genai_types.Part(text=text)])
+
+    def _convert_tools_for_gemini(self) -> genai_types.ToolListUnion:
+        """Convert MCP tools to Gemini tool format."""
+        self._gemini_to_mcp_tool_map = {}  # Reset mapping
+        self.gemini_tools = []
+
+        for tool in self.get_available_tools():
+            gemini_tool = self._to_gemini_tool(tool)
+            if gemini_tool is None:
+                continue
+
+            self._gemini_to_mcp_tool_map[tool.name] = tool.name
+            self.gemini_tools.append(gemini_tool)
+
+        return self.gemini_tools
+
+    def _to_gemini_tool(self, tool: types.Tool) -> genai_types.Tool | None:
+        """Convert a single MCP tool to Gemini tool format.
+
+        Subclasses can override to customize tool conversion (e.g., for computer use).
+        """
+        # Ensure parameters have proper Schema format
+        if tool.description is None or tool.inputSchema is None:
+            raise ValueError(f"MCP tool {tool.name} requires both a description and inputSchema.")
+        function_decl = genai_types.FunctionDeclaration(
+            name=tool.name,
+            description=tool.description,
+            parameters_json_schema=tool.inputSchema,
+        )
+        return genai_types.Tool(function_declarations=[function_decl])

hud/agents/gemini_cua.py

@@ -0,0 +1,335 @@
+"""Gemini Computer Use Agent implementation."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, ClassVar
+
+import mcp.types as types
+from google.genai import types as genai_types
+from pydantic import ConfigDict, Field
+
+from hud.tools.computer.settings import computer_settings
+from hud.types import AgentResponse, BaseAgentConfig, MCPToolCall, MCPToolResult
+from hud.utils.types import with_signature
+
+from .base import BaseCreateParams, MCPAgent
+from .gemini import GeminiAgent, GeminiConfig
+
+logger = logging.getLogger(__name__)
+
+# Predefined Gemini computer use functions
+PREDEFINED_COMPUTER_USE_FUNCTIONS = [
+    "open_web_browser",
+    "click_at",
+    "hover_at",
+    "type_text_at",
+    "scroll_document",
+    "scroll_at",
+    "wait_5_seconds",
+    "go_back",
+    "go_forward",
+    "search",
+    "navigate",
+    "key_combination",
+    "drag_and_drop",
+]
+
+GEMINI_CUA_INSTRUCTIONS = """
+You are an autonomous computer-using agent. Follow these guidelines:
+
+1. NEVER ask for confirmation. Complete all tasks autonomously.
+2. Do NOT send messages like "I need to confirm before..." or "Do you want me to
+   continue?" - just proceed.
+3. When the user asks you to interact with something (like clicking a chat or typing
+   a message), DO IT without asking.
+4. Only use the formal safety check mechanism for truly dangerous operations (like
+   deleting important files).
+5. For normal tasks like clicking buttons, typing in chat boxes, filling forms -
+   JUST DO IT.
+6. The user has already given you permission by running this agent. No further
+   confirmation is needed.
+7. Be decisive and action-oriented. Complete the requested task fully.
+
+Remember: You are expected to complete tasks autonomously. The user trusts you to do
+what they asked.
+""".strip()
+
+
+class GeminiCUAConfig(GeminiConfig):
+    """Configuration for `GeminiCUAAgent`."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    model_name: str = "GeminiCUA"
+    model: str = "gemini-2.5-computer-use-preview-10-2025"
+    excluded_predefined_functions: list[str] = Field(default_factory=list)
+
+
+class GeminiCUACreateParams(BaseCreateParams, GeminiCUAConfig):
+    pass
+
+
+class GeminiCUAAgent(GeminiAgent):
+    """
+    Gemini Computer Use Agent that extends GeminiAgent with computer use capabilities.
+
+    This agent uses Gemini's native computer use capabilities but executes
+    tools through MCP servers instead of direct implementation.
+    """
+
+    metadata: ClassVar[dict[str, Any] | None] = {
+        "display_width": computer_settings.GEMINI_COMPUTER_WIDTH,
+        "display_height": computer_settings.GEMINI_COMPUTER_HEIGHT,
+    }
+    required_tools: ClassVar[list[str]] = ["gemini_computer"]
+    config_cls: ClassVar[type[BaseAgentConfig]] = GeminiCUAConfig
+
+    @with_signature(GeminiCUACreateParams)
+    @classmethod
+    def create(cls, **kwargs: Any) -> GeminiCUAAgent:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return MCPAgent.create.__func__(cls, **kwargs)  # type: ignore[return-value]
+
+    def __init__(self, params: GeminiCUACreateParams | None = None, **kwargs: Any) -> None:
+        super().__init__(params, **kwargs)  # type: ignore[arg-type]
+        self.config: GeminiCUAConfig  # type: ignore[assignment]
+
+        self._computer_tool_name = "gemini_computer"
+        self.excluded_predefined_functions = list(self.config.excluded_predefined_functions)
+
+        # Context management: Maximum number of recent turns to keep screenshots for
+        # Configurable via GEMINI_MAX_RECENT_TURN_WITH_SCREENSHOTS environment variable
+        self.max_recent_turn_with_screenshots = (
+            computer_settings.GEMINI_MAX_RECENT_TURN_WITH_SCREENSHOTS
+        )
+
+        # Add computer use instructions
+        if self.system_prompt:
+            self.system_prompt = f"{self.system_prompt}\n\n{GEMINI_CUA_INSTRUCTIONS}"
+        else:
+            self.system_prompt = GEMINI_CUA_INSTRUCTIONS
+
+    def _to_gemini_tool(self, tool: types.Tool) -> genai_types.Tool | None:
+        """Convert a single MCP tool to Gemini tool format.
+
+        Handles gemini_computer tool specially by using Gemini's native ComputerUse.
+        """
+        if tool.name == self._computer_tool_name:
+            # Use Gemini's native computer use capability
+            return genai_types.Tool(
+                computer_use=genai_types.ComputerUse(
+                    environment=genai_types.Environment.ENVIRONMENT_BROWSER,
+                    excluded_predefined_functions=self.excluded_predefined_functions,
+                )
+            )
+
+        # For non-computer tools, use the parent implementation
+        return super()._to_gemini_tool(tool)
+
+    async def get_response(self, messages: list[genai_types.Content]) -> AgentResponse:
+        """Get response from Gemini including any tool calls.
+
+        Extends parent to trim old screenshots before making API call.
+        """
+        # Trim screenshots from older turns to manage context growth
+        self._remove_old_screenshots(messages)
+
+        return await super().get_response(messages)
+
+    async def format_tool_results(
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
+    ) -> list[genai_types.Content]:
+        """Format tool results into Gemini messages.
+
+        Handles computer tool results specially with screenshots and URLs.
+        """
+        # Process each tool result
+        function_responses = []
+
+        for tool_call, result in zip(tool_calls, tool_results, strict=True):
+            # Get the Gemini function name from metadata
+            gemini_name = getattr(tool_call, "gemini_name", tool_call.name)
+
+            # Check if this is a computer use tool call
+            is_computer_call = tool_call.name == self._computer_tool_name
+
+            # Convert MCP tool results to Gemini format
+            response_dict: dict[str, Any] = {}
+            url = None
+
+            if result.isError:
+                # Extract error message from content
+                error_msg = "Tool execution failed"
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        # Check if this is a URL metadata block
+                        if content.text.startswith("__URL__:"):
+                            url = content.text.replace("__URL__:", "")
+                        else:
+                            error_msg = content.text
+                            break
+                response_dict["error"] = error_msg
+                # for gemini cua agent, if a nonexistend computer tool is called, it won't
+                # #technically count as a computer tool call, but we still need to return a url
+                response_dict["url"] = url if url else "about:blank"
+            else:
+                # Process success content
+                response_dict["success"] = True
+
+            # Extract URL and screenshot from content (for computer use)
+            screenshot_parts = []
+            if is_computer_call:
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        # Check if this is a URL metadata block
+                        if content.text.startswith("__URL__:"):
+                            url = content.text.replace("__URL__:", "")
+                    elif isinstance(content, types.ImageContent):
+                        # Decode base64 string to bytes for FunctionResponseBlob
+                        import base64
+
+                        image_bytes = base64.b64decode(content.data)
+                        screenshot_parts.append(
+                            genai_types.FunctionResponsePart(
+                                inline_data=genai_types.FunctionResponseBlob(
+                                    mime_type=content.mimeType or "image/png",
+                                    data=image_bytes,
+                                )
+                            )
+                        )
+
+                # Add URL to response dict (required by Gemini Computer Use model)
+                # URL must ALWAYS be present per Gemini API requirements
+                response_dict["url"] = url if url else "about:blank"
+
+                # For Gemini Computer Use actions, always acknowledge safety decisions
+                requires_ack = False
+                if tool_call.arguments:
+                    requires_ack = bool(tool_call.arguments.get("safety_decision"))
+                if requires_ack:
+                    response_dict["safety_acknowledgement"] = True
+            else:
+                # For non-computer tools, add text content to response
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        response_dict["output"] = content.text
+                        break
+
+            # Create function response
+            function_response = genai_types.FunctionResponse(
+                name=gemini_name,
+                response=response_dict,
+                parts=screenshot_parts if screenshot_parts else None,
+            )
+            function_responses.append(function_response)
+
+        # Return as a user message containing all function responses
+        return [
+            genai_types.Content(
+                role="user",
+                parts=[genai_types.Part(function_response=fr) for fr in function_responses],
+            )
+        ]
+
+    def _extract_tool_call(self, part: genai_types.Part) -> MCPToolCall | None:
+        """Extract an MCPToolCall from a function call part.
+
+        Routes predefined Gemini Computer Use functions to the gemini_computer tool
+        and normalizes the arguments to MCP tool schema.
+        """
+        if not part.function_call:
+            return None
+
+        func_name = part.function_call.name or ""
+        raw_args = dict(part.function_call.args) if part.function_call.args else {}
+
+        # Route predefined computer use functions to the computer tool
+        if func_name in PREDEFINED_COMPUTER_USE_FUNCTIONS:
+            # Normalize Gemini Computer Use calls to MCP tool schema
+            # Ensure 'action' is present and equals the Gemini function name
+            normalized_args: dict[str, Any] = {"action": func_name}
+
+            # Map common argument shapes used by Gemini Computer Use
+            # 1) Coordinate arrays → x/y
+            coord = raw_args.get("coordinate") or raw_args.get("coordinates")
+            if isinstance(coord, list | tuple) and len(coord) >= 2:
+                try:
+                    normalized_args["x"] = int(coord[0])
+                    normalized_args["y"] = int(coord[1])
+                except (TypeError, ValueError):
+                    # Fall back to raw if casting fails
+                    pass
+
+            # Destination coordinate arrays → destination_x/destination_y
+            dest = (
+                raw_args.get("destination")
+                or raw_args.get("destination_coordinate")
+                or raw_args.get("destinationCoordinate")
+            )
+            if isinstance(dest, list | tuple) and len(dest) >= 2:
+                try:
+                    normalized_args["destination_x"] = int(dest[0])
+                    normalized_args["destination_y"] = int(dest[1])
+                except (TypeError, ValueError):
+                    pass
+
+            # Pass through supported fields if present (including direct coords)
+            for key in (
+                "text",
+                "press_enter",
+                "clear_before_typing",
+                "safety_decision",
+                "direction",
+                "magnitude",
+                "url",
+                "keys",
+                "x",
+                "y",
+                "destination_x",
+                "destination_y",
+            ):
+                if key in raw_args:
+                    normalized_args[key] = raw_args[key]
+
+            return MCPToolCall(
+                name=self._computer_tool_name,
+                arguments=normalized_args,
+                gemini_name=func_name,  # type: ignore[arg-type]
+            )
+
+        # Non-computer tools: use parent implementation
+        return super()._extract_tool_call(part)
+
+    def _remove_old_screenshots(self, messages: list[genai_types.Content]) -> None:
+        """
+        Remove screenshots from old turns to manage context length.
+        Keeps only the last N turns with screenshots (configured via
+        self.max_recent_turn_with_screenshots).
+        """
+        turn_with_screenshots_found = 0
+
+        for content in reversed(messages):
+            if content.role == "user" and content.parts:
+                # Check if content has screenshots (function responses with images)
+                has_screenshot = False
+                for part in content.parts:
+                    if (
+                        part.function_response
+                        and part.function_response.parts
+                        and part.function_response.name in PREDEFINED_COMPUTER_USE_FUNCTIONS
+                    ):
+                        has_screenshot = True
+                        break
+
+                if has_screenshot:
+                    turn_with_screenshots_found += 1
+                    # Remove the screenshot image if the number of screenshots exceeds the limit
+                    if turn_with_screenshots_found > self.max_recent_turn_with_screenshots:
+                        for part in content.parts:
+                            if (
+                                part.function_response
+                                and part.function_response.parts
+                                and part.function_response.name in PREDEFINED_COMPUTER_USE_FUNCTIONS
+                            ):
+                                # Clear the parts (screenshots)
+                                part.function_response.parts = None