hud-python 0.4.20__py3-none-any.whl → 0.4.22__py3-none-any.whl

hud/tools/grounding/config.py

@@ -0,0 +1,54 @@
+"""Configuration for grounding models."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+SYSTEM_PROMPT = (
+    "You are a visual grounding model. Given an image and a description, "
+    "return ONLY the center pixel coordinates of the described element as a "
+    "single point in parentheses format: (x, y). Do not return bounding boxes "
+    "or multiple coordinates."
+)
+
+
+@dataclass
+class GrounderConfig:
+    """Configuration for grounding model clients.
+
+    Attributes:
+        api_base: Base URL for the grounding model API endpoint
+        model: Model identifier to use for grounding
+        api_key: API key for authentication (default: "EMPTY" for local models)
+        system_prompt: System prompt to guide the grounding model
+        output_format: Format for coordinate output ("pixels", "norm_0_1", "norm_0_999")
+        parser_regex: Regular expression to parse coordinates from model output
+        resize: Image resizing configuration dictionary
+    """
+
+    api_base: str
+    model: str
+    api_key: str = "EMPTY"
+    system_prompt: str = SYSTEM_PROMPT
+    output_format: str = "pixels"  # "pixels" | "norm_0_1" | "norm_0_999"
+    parser_regex: str = r"\((\d+),\s*(\d+)\)"
+    resize: dict[str, Any] = field(
+        default_factory=lambda: {
+            "enabled": True,
+            "min_pixels": 3136,
+            "max_pixels": 4096 * 2160,
+            "factor": 28,
+        }
+    )
+
+    def __post_init__(self) -> None:
+        """Validate configuration after initialization."""
+        if self.output_format not in ("pixels", "norm_0_1", "norm_0_999"):
+            raise ValueError(f"Invalid output_format: {self.output_format}")
+
+        if not self.api_base:
+            raise ValueError("api_base is required")
+
+        if not self.model:
+            raise ValueError("model is required")

hud/tools/grounding/grounded_tool.py

@@ -0,0 +1,314 @@
+"""Grounded computer tool that resolves element descriptions to coordinates."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from mcp import ErrorData, McpError
+from mcp.types import INVALID_PARAMS, ContentBlock
+
+from hud.clients.base import AgentMCPClient  # noqa: TC001
+from hud.tools.grounding.grounder import Grounder  # noqa: TC001
+from hud.types import MCPToolCall
+
+logger = logging.getLogger(__name__)
+
+
+class GroundedComputerTool:
+    """Computer tool wrapper that grounds element descriptions to coordinates.
+
+    This tool acts as a local wrapper that:
+    1. Accepts natural language element descriptions from the agent
+    2. Calls the environment's computer tool via MCP to take screenshots
+    3. Uses a grounding model to resolve descriptions to coordinates
+    4. Calls the environment's computer tool via MCP with resolved coordinates
+    5. Returns the result to the agent
+
+    This allows the agent to use element descriptions while ensuring all
+    computer actions happen in the correct environment.
+    """
+
+    def __init__(
+        self,
+        *,
+        grounder: Grounder,
+        mcp_client: AgentMCPClient,
+        computer_tool_name: str = "computer",
+    ) -> None:
+        """Initialize the grounded computer tool.
+
+        Args:
+            grounder: Grounder instance for visual grounding
+            mcp_client: MCP client to call the environment's computer tool
+            computer_tool_name: Name of the computer tool in the environment
+        """
+        self._grounder = grounder
+        self._mcp_client = mcp_client
+        self._computer_tool_name = computer_tool_name
+
+    def get_openai_tool_schema(self) -> dict:
+        """Get the OpenAI tool schema for the grounded computer tool.
+
+        Returns:
+            Dictionary containing the tool schema in OpenAI format
+        """
+        return {
+            "type": "function",
+            "function": {
+                "name": "computer",
+                "description": (
+                    "Control a computer by interacting with UI elements. This tool uses "
+                    "element descriptions to locate and interact with UI elements on the "
+                    "screen (e.g., 'red submit button', 'search text field', 'hamburger menu "
+                    "icon', 'close button in top right corner')."
+                ),
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "action": {
+                            "type": "string",
+                            "enum": [
+                                "click",
+                                "double_click",
+                                "move",
+                                "scroll",
+                                "drag",
+                                "type",
+                                "keypress",
+                                "wait",
+                                "screenshot",
+                                "get_current_url",
+                                "get_dimensions",
+                                "get_environment",
+                            ],
+                            "description": "The action to perform",
+                        },
+                        "element_description": {
+                            "type": "string",
+                            "description": (
+                                "Natural language description of the element for "
+                                "click/move/scroll actions"
+                            ),
+                        },
+                        "start_element_description": {
+                            "type": "string",
+                            "description": "Description of the start element for drag actions",
+                        },
+                        "end_element_description": {
+                            "type": "string",
+                            "description": "Description of the end element for drag actions",
+                        },
+                        "text": {"type": "string", "description": "Text to type"},
+                        "keys": {
+                            "type": "array",
+                            "items": {"type": "string"},
+                            "description": "Keys to press (e.g., ['ctrl', 'a'] for Ctrl+A)",
+                        },
+                        "button": {
+                            "type": "string",
+                            "enum": ["left", "right", "middle"],
+                            "description": "Mouse button to use",
+                        },
+                        "scroll_x": {"type": "integer", "description": "Horizontal scroll amount"},
+                        "scroll_y": {"type": "integer", "description": "Vertical scroll amount"},
+                    },
+                    "required": ["action"],
+                },
+            },
+        }
+
+    async def __call__(
+        self,
+        action: str,
+        # Screenshot from conversation
+        screenshot_b64: str | None = None,
+        # Grounding-specific parameters
+        element_description: str | None = None,
+        start_element_description: str | None = None,
+        end_element_description: str | None = None,
+        # Pass-through parameters
+        text: str | None = None,
+        keys: list[str] | None = None,
+        button: str | None = None,
+        scroll_x: int | None = None,
+        scroll_y: int | None = None,
+        **kwargs: Any,
+    ) -> list[ContentBlock]:
+        """Execute a computer action, grounding element descriptions to coordinates first.
+
+        This method calls the environment's computer tool through MCP to ensure
+        actions happen in the correct environment.
+
+        Args:
+            action: The action to perform
+            element_description: Description of element for click/move/scroll actions
+            start_element_description: Start element for drag actions
+            end_element_description: End element for drag actions
+            text: Text to type for type actions
+            keys: Keys to press for keypress actions
+            button: Mouse button (left, right, middle)
+            scroll_x: Horizontal scroll amount
+            scroll_y: Vertical scroll amount
+            **kwargs: Additional arguments
+
+        Returns:
+            List of ContentBlocks with action results from the environment
+        """
+        try:
+            # For actions that don't need grounding, call environment tool directly
+            if action in (
+                "screenshot",
+                "type",
+                "keypress",
+                "wait",
+                "get_current_url",
+                "get_dimensions",
+                "get_environment",
+            ):
+                computer_args: dict[str, Any] = {"action": action}
+                if text is not None:
+                    computer_args["text"] = text
+                if keys is not None:
+                    computer_args["keys"] = keys
+
+                result = await self._mcp_client.call_tool(
+                    MCPToolCall(
+                        name=self._computer_tool_name, arguments={**computer_args, **kwargs}
+                    )
+                )
+                return result.content
+
+            # For actions that need coordinates, we need to ground element descriptions
+            if action in ("click", "double_click", "move", "scroll"):
+                if not element_description:
+                    raise McpError(
+                        ErrorData(
+                            code=INVALID_PARAMS,
+                            message=f"element_description is required for {action} action",
+                        )
+                    )
+
+                if not screenshot_b64:
+                    raise McpError(
+                        ErrorData(
+                            code=INVALID_PARAMS, message="No screenshot available for grounding"
+                        )
+                    )
+
+                # Ground the element description to coordinates
+                coords = await self._grounder.predict_click(
+                    image_b64=screenshot_b64, instruction=element_description
+                )
+
+                if not coords:
+                    raise McpError(
+                        ErrorData(
+                            code=INVALID_PARAMS,
+                            message=(
+                                f"Could not locate element: '{element_description}'. "
+                                "Try a more specific description or different identifying "
+                                "features (color, position, text, etc.)"
+                            ),
+                        )
+                    )
+
+                x, y = coords
+
+                # Execute action with resolved coordinates
+                computer_args: dict[str, Any] = {"action": action, "x": x, "y": y}
+                if button:
+                    computer_args["button"] = button
+                if scroll_x is not None:
+                    computer_args["scroll_x"] = scroll_x
+                if scroll_y is not None:
+                    computer_args["scroll_y"] = scroll_y
+
+                result = await self._mcp_client.call_tool(
+                    MCPToolCall(
+                        name=self._computer_tool_name, arguments={**computer_args, **kwargs}
+                    )
+                )
+                return result.content
+
+            elif action == "drag":
+                if not start_element_description or not end_element_description:
+                    raise McpError(
+                        ErrorData(
+                            code=INVALID_PARAMS,
+                            message=(
+                                "start_element_description and end_element_description "
+                                "are required for drag action"
+                            ),
+                        )
+                    )
+
+                if not screenshot_b64:
+                    raise McpError(
+                        ErrorData(
+                            code=INVALID_PARAMS, message="No screenshot available for grounding"
+                        )
+                    )
+
+                # Ground both start and end points
+                start_coords = await self._grounder.predict_click(
+                    image_b64=screenshot_b64, instruction=start_element_description
+                )
+
+                if not start_coords:
+                    raise McpError(
+                        ErrorData(
+                            code=INVALID_PARAMS,
+                            message=(
+                                f"Could not locate start element: '{start_element_description}'. "
+                                "Try a more specific description or different identifying features."
+                            ),
+                        )
+                    )
+
+                end_coords = await self._grounder.predict_click(
+                    image_b64=screenshot_b64, instruction=end_element_description
+                )
+
+                if not end_coords:
+                    raise McpError(
+                        ErrorData(
+                            code=INVALID_PARAMS,
+                            message=(
+                                f"Could not locate end element: '{end_element_description}'. "
+                                "Try a more specific description or different identifying features."
+                            ),
+                        )
+                    )
+
+                # Execute drag with resolved coordinates
+                computer_args: dict[str, Any] = {
+                    "action": "drag",
+                    "path": [
+                        (start_coords[0], start_coords[1]),
+                        (end_coords[0], end_coords[1]),
+                    ],
+                }
+                if button:
+                    computer_args["button"] = button
+
+                result = await self._mcp_client.call_tool(
+                    MCPToolCall(
+                        name=self._computer_tool_name, arguments={**computer_args, **kwargs}
+                    )
+                )
+                return result.content
+
+            else:
+                raise McpError(
+                    ErrorData(code=INVALID_PARAMS, message=f"Unsupported action: {action}")
+                )
+
+        except McpError:
+            # Re-raise MCP errors
+            raise
+        except Exception as e:
+            logger.error("Grounded tool failed: %s", e)
+            raise McpError(
+                ErrorData(code=INVALID_PARAMS, message=f"Grounding failed: {e!s}")
+            ) from e

hud/tools/grounding/grounder.py

@@ -0,0 +1,301 @@
+"""OpenAI-based grounder for visual element detection."""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+import re
+
+from openai import AsyncOpenAI
+from opentelemetry import trace
+from PIL import Image
+
+from hud import instrument
+from hud.tools.grounding.config import GrounderConfig  # noqa: TC001
+
+
+class Grounder:
+    """Grounder that uses AsyncOpenAI to call vLLM or other model endpoints for visual grounding.
+
+    This class handles:
+    - Image resizing based on configuration
+    - API calls to grounding models via AsyncOpenAI
+    - Coordinate parsing from model outputs
+    - Coordinate format conversion (pixels, normalized)
+    """
+
+    def __init__(self, config: GrounderConfig) -> None:
+        """Initialize the grounder with configuration.
+
+        Args:
+            config: GrounderConfig with API endpoint, model, and parsing settings
+        """
+        self.config = config
+        self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.api_base)
+
+    def _resize_image(self, image_b64: str) -> tuple[str, tuple[int, int], tuple[int, int]]:
+        """Resize image according to configuration.
+
+        Args:
+            image_b64: Base64-encoded image string
+
+        Returns:
+            Tuple of (processed_base64, (original_width, original_height),
+                     (processed_width, processed_height))
+        """
+        # Decode image
+        image_bytes = base64.b64decode(image_b64)
+        img = Image.open(io.BytesIO(image_bytes))
+        original_size = (img.width, img.height)
+
+        if not self.config.resize["enabled"]:
+            return image_b64, original_size, original_size
+
+        # Calculate total pixels
+        total_pixels = img.width * img.height
+        min_pixels = self.config.resize["min_pixels"]
+        max_pixels = self.config.resize["max_pixels"]
+        factor = self.config.resize["factor"]
+
+        # Determine if resizing is needed
+        if total_pixels < min_pixels or total_pixels > max_pixels:
+            # Calculate scaling factor
+            if total_pixels < min_pixels:
+                scale = (min_pixels / total_pixels) ** 0.5
+            else:
+                scale = (max_pixels / total_pixels) ** 0.5
+
+            # Round dimensions to nearest factor
+            new_width = int((img.width * scale) // factor) * factor
+            new_height = int((img.height * scale) // factor) * factor
+
+            # Ensure minimum dimensions
+            new_width = max(new_width, factor)
+            new_height = max(new_height, factor)
+
+            # Resize image
+            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+            # Convert back to base64
+            buffer = io.BytesIO()
+            img.save(buffer, format="PNG")
+            resized_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+            return resized_b64, original_size, (new_width, new_height)
+
+        return image_b64, original_size, original_size
+
+    def _parse_coordinates(self, response_text: str) -> tuple[float, float] | None:
+        """Parse coordinates from model response.
+
+        Handles multiple formats:
+        - (x, y) format from configured regex
+        - [x1, y1, x2, y2] bounding box format (returns center point)
+        - [x, y] point format
+
+        Args:
+            response_text: Text output from the grounding model
+
+        Returns:
+            Tuple of (x, y) coordinates or None if parsing fails
+        """
+        # First try the configured regex pattern
+        match = re.search(self.config.parser_regex, response_text)
+        if match:
+            try:
+                x = float(match.group(1))
+                y = float(match.group(2))
+                return (x, y)
+            except (ValueError, IndexError):
+                # If parsing fails, continue to fallback strategies
+                pass
+
+        # Try to parse as a list/array format [x1, y1, x2, y2] or [x, y]
+        # Also handles (x1, y1, x2, y2)
+        # Updated pattern to handle both integers and floats
+        list_pattern = (
+            r"[\[\(](\d+(?:\.\d+)?)[,\s]+(\d+(?:\.\d+)?)"
+            r"(?:[,\s]+(\d+(?:\.\d+)?)[,\s]+(\d+(?:\.\d+)?))?[\]\)]"
+        )
+        list_match = re.search(list_pattern, response_text)
+        if list_match:
+            x1 = float(list_match.group(1))
+            y1 = float(list_match.group(2))
+
+            # Check if it's a bounding box (4 values) or a point (2 values)
+            if list_match.group(3) and list_match.group(4):
+                # Bounding box format - return center point
+                x2 = float(list_match.group(3))
+                y2 = float(list_match.group(4))
+                center_x = (x1 + x2) / 2
+                center_y = (y1 + y2) / 2
+                return (center_x, center_y)
+            else:
+                # Point format
+                return (x1, y1)
+
+        return None
+
+    def _convert_coordinates(
+        self,
+        coords: tuple[float, float],
+        processed_size: tuple[int, int],
+        original_size: tuple[int, int],
+    ) -> tuple[int, int]:
+        """Convert coordinates based on output format configuration and scale to original size.
+
+        Args:
+            coords: Raw coordinates from model (can be float for normalized formats)
+            processed_size: Dimensions of the processed/resized image (width, height)
+            original_size: Original image dimensions (width, height)
+
+        Returns:
+            Converted coordinates in original image pixels
+        """
+        x, y = coords
+        proc_width, proc_height = processed_size
+        orig_width, orig_height = original_size
+
+        # First convert to pixels in the processed image space
+        if self.config.output_format == "pixels":
+            # Already in pixels of processed image
+            proc_x, proc_y = x, y
+        elif self.config.output_format == "norm_0_1":
+            # Convert from 0-1 normalized to pixels
+            proc_x = x * proc_width
+            proc_y = y * proc_height
+        elif self.config.output_format == "norm_0_999":
+            # Convert from 0-999 normalized to pixels
+            proc_x = x * proc_width / 999
+            proc_y = y * proc_height / 999
+        else:
+            proc_x, proc_y = x, y
+
+        # Scale from processed image coordinates to original image coordinates
+        scale_x = orig_width / proc_width
+        scale_y = orig_height / proc_height
+
+        final_x = int(proc_x * scale_x)
+        final_y = int(proc_y * scale_y)
+
+        return (final_x, final_y)
+
+    @instrument(
+        name="Grounding.predict_click",
+        span_type="agent",
+        record_args=True,
+        record_result=True,
+    )
+    async def predict_click(
+        self, *, image_b64: str, instruction: str, max_retries: int = 3
+    ) -> tuple[int, int] | None:
+        """Predict click coordinates for the given instruction on the image.
+
+        Args:
+            image_b64: Base64-encoded screenshot
+            instruction: Natural language description of the element to click
+            max_retries: Maximum number of retry attempts (default: 3)
+
+        Returns:
+            Tuple of (x, y) pixel coordinates or None if grounding fails
+        """
+
+        # Resize image once outside the retry loop
+        processed_image, original_size, processed_size = self._resize_image(image_b64)
+
+        # Build messages once
+        messages = []
+
+        # Add system prompt if configured
+        if self.config.system_prompt:
+            messages.append(
+                {
+                    "role": "system",
+                    "content": (
+                        self.config.system_prompt
+                        + f" The image resolution is height {processed_size[1]} "
+                        + f"and width {processed_size[0]}."
+                    ),
+                }
+            )
+
+        # Add user message with image and instruction
+        messages.append(
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{processed_image}"},
+                    },
+                    {"type": "text", "text": instruction},
+                ],
+            }
+        )
+
+        # Retry loop
+        for attempt in range(max_retries):
+            try:
+                # Call the grounding model via AsyncOpenAI
+                response = await self.client.chat.completions.create(
+                    model=self.config.model,
+                    messages=messages,
+                    temperature=0.0,
+                    max_tokens=50,
+                )
+
+                # Extract response text
+                response_text = response.choices[0].message.content
+
+                # Manually record the raw response in the span
+                span = trace.get_current_span()
+                if span and span.is_recording():
+                    span.set_attribute("grounder.raw_response", json.dumps(response.model_dump()))
+                    span.set_attribute("grounder.attempt", attempt + 1)
+
+                # Parse coordinates from response
+                if response_text is None:
+                    if attempt < max_retries - 1:
+                        continue
+                    return None
+
+                coords = self._parse_coordinates(response_text)
+                if coords is None:
+                    if attempt < max_retries - 1:
+                        continue
+                    return None
+
+                # Convert coordinates to original image pixels based on output format and scaling
+                pixel_coords = self._convert_coordinates(coords, processed_size, original_size)
+
+                # Validate coordinates are within image bounds
+                x, y = pixel_coords
+                if x < 0 or y < 0 or x >= original_size[0] or y >= original_size[1]:
+                    # Clamp to image bounds
+                    x = max(0, min(x, original_size[0] - 1))
+                    y = max(0, min(y, original_size[1] - 1))
+                    pixel_coords = (x, y)
+
+                # Record successful grounding in span
+                span = trace.get_current_span()
+                if span and span.is_recording():
+                    span.set_attribute("grounder.success", True)
+                    span.set_attribute(
+                        "grounder.final_coords", f"{pixel_coords[0]},{pixel_coords[1]}"
+                    )
+                    span.set_attribute("grounder.total_attempts", attempt + 1)
+
+                return pixel_coords
+
+            except Exception:
+                if attempt < max_retries - 1:
+                    continue
+
+        # Record failure in span
+        span = trace.get_current_span()
+        if span and span.is_recording():
+            span.set_attribute("grounder.success", False)
+            span.set_attribute("grounder.total_attempts", max_retries)
+            span.set_attribute("grounder.failure_reason", "All attempts exhausted")
+
+        return None

hud/tools/grounding/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for grounding tools."""