hud-python 0.4.21__py3-none-any.whl → 0.4.23__py3-none-any.whl

hud/tools/grounding/grounder.py

@@ -0,0 +1,302 @@
+"""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 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
+        from PIL import 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."""

hud/tools/grounding/tests/test_grounded_tool.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import mcp.types as types
+import pytest
+
+from hud.tools.grounding.grounded_tool import GroundedComputerTool
+from hud.types import MCPToolCall, MCPToolResult
+
+
+@dataclass
+class FakeResult:
+    content: list[types.ContentBlock]
+    isError: bool = False
+    structuredContent: dict | None = None
+
+
+class FakeMCPClient:
+    """Fake MCP client that implements AgentMCPClient protocol."""
+
+    _initialized: bool
+
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, dict[str, Any]]] = []
+        self._initialized = False
+
+    @property
+    def mcp_config(self) -> dict[str, dict[str, Any]]:
+        return {"test": {"command": "echo", "args": ["test"]}}
+
+    @property
+    def is_connected(self) -> bool:
+        return self._initialized
+
+    async def initialize(self, mcp_config: dict[str, dict[str, Any]] | None = None) -> None:
+        self._initialized = True
+
+    async def list_tools(self) -> list[types.Tool]:
+        return [types.Tool(name="computer", description="Test tool", inputSchema={})]
+
+    async def call_tool(self, tool_call: MCPToolCall) -> MCPToolResult:
+        self.calls.append((tool_call.name, tool_call.arguments or {}))
+        return MCPToolResult(content=[types.TextContent(text="ok", type="text")], isError=False)
+
+    async def shutdown(self) -> None:
+        self._initialized = False
+
+
+class FakeGrounder:
+    """Fake grounder that implements Grounder interface."""
+
+    def __init__(self, coords: tuple[int, int] | None = (10, 20)) -> None:
+        self.coords = coords
+        self.calls: list[tuple[str, str]] = []
+
+    async def predict_click(
+        self, *, image_b64: str, instruction: str, max_retries: int = 3
+    ) -> tuple[int, int] | None:
+        self.calls.append((image_b64[:10], instruction))
+        return self.coords
+
+
+def _png_b64() -> str:
+    # 1x1 transparent PNG base64 (valid minimal image)
+    return (
+        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR4nGMAAQAABQAB"
+        "J2n0mQAAAABJRU5ErkJggg=="
+    )
+
+
+@pytest.mark.asyncio
+async def test_click_action_grounds_and_calls_mcp() -> None:
+    client = FakeMCPClient()
+    grounder = FakeGrounder(coords=(123, 456))
+    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+
+    blocks = await tool(
+        action="click",
+        element_description="red button",
+        screenshot_b64=_png_b64(),
+        button="left",
+    )
+
+    assert isinstance(blocks, list)
+    # Grounder called once
+    assert len(grounder.calls) == 1
+    # MCP called with resolved coordinates
+    assert client.calls == [("computer", {"action": "click", "x": 123, "y": 456, "button": "left"})]
+
+
+@pytest.mark.asyncio
+async def test_move_and_scroll_require_element_description_and_screenshot() -> None:
+    client = FakeMCPClient()
+    grounder = FakeGrounder(coords=(5, 6))
+    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+
+    # Missing element_description
+    with pytest.raises(Exception) as ei:
+        await tool(action="move", screenshot_b64=_png_b64())
+    assert "element_description is required" in str(ei.value)
+
+    # Missing screenshot
+    with pytest.raises(Exception) as ei2:
+        await tool(action="scroll", element_description="list", scroll_y=100)
+    assert "No screenshot available" in str(ei2.value)
+
+
+@pytest.mark.asyncio
+async def test_drag_grounds_both_points_and_calls_mcp() -> None:
+    client = FakeMCPClient()
+    grounder = FakeGrounder(coords=(10, 20))
+    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+
+    await tool(
+        action="drag",
+        start_element_description="source",
+        end_element_description="target",
+        screenshot_b64=_png_b64(),
+        button="left",
+    )
+
+    # Two grounding calls (start and end)
+    assert len(grounder.calls) == 2
+    # Drag path contains two points, same coords from fake grounder
+    name, args = client.calls[0]
+    assert name == "computer"
+    assert args["action"] == "drag"
+    assert args["button"] == "left"
+    assert args["path"] == [(10, 20), (10, 20)]
+
+
+@pytest.mark.asyncio
+async def test_drag_requires_both_descriptions_and_screenshot() -> None:
+    client = FakeMCPClient()
+    grounder = FakeGrounder()
+    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+
+    with pytest.raises(Exception) as ei:
+        await tool(action="drag", start_element_description="a", screenshot_b64=_png_b64())
+    assert "start_element_description and end_element_description" in str(ei.value)
+
+    with pytest.raises(Exception) as ei2:
+        await tool(
+            action="drag",
+            start_element_description="a",
+            end_element_description="b",
+        )
+    assert "No screenshot available" in str(ei2.value)
+
+
+@pytest.mark.asyncio
+async def test_direct_actions_bypass_grounding_and_call_mcp() -> None:
+    client = FakeMCPClient()
+    grounder = FakeGrounder()
+    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+
+    # Actions that bypass grounding
+    for action, extra in [
+        ("screenshot", {}),
+        ("type", {"text": "hello"}),
+        ("keypress", {"keys": ["ctrl", "a"]}),
+        ("wait", {}),
+        ("get_current_url", {}),
+        ("get_dimensions", {}),
+        ("get_environment", {}),
+    ]:
+        client.calls.clear()
+        _ = await tool(action=action, **extra)
+        assert client.calls and client.calls[0][0] == "computer"
+        assert client.calls[0][1]["action"] == action
+    # Grounder not invoked for these
+    assert grounder.calls == []
+
+
+@pytest.mark.asyncio
+async def test_unsupported_action_raises() -> None:
+    client = FakeMCPClient()
+    grounder = FakeGrounder()
+    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+
+    with pytest.raises(Exception) as ei:
+        await tool(action="zoom")
+    assert "Unsupported action" in str(ei.value)
+
+
+@pytest.mark.asyncio
+async def test_grounding_failure_propagates_as_error() -> None:
+    client = FakeMCPClient()
+    grounder = FakeGrounder(coords=None)
+    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+
+    with pytest.raises(Exception) as ei:
+        await tool(action="click", element_description="x", screenshot_b64=_png_b64())
+    assert "Could not locate element" in str(ei.value)

hud/tools/tests/test_playwright_tool.py

@@ -52,7 +52,7 @@ class TestPlaywrightTool:
             assert any(isinstance(b, TextContent) for b in blocks)
             # The actual call includes wait_until parameter with a Field object
             mock_page.goto.assert_called_once()
-            args, kwargs = mock_page.goto.call_args
+            args, _kwargs = mock_page.goto.call_args
             assert args[0] == "https://example.com"
             mock_ensure.assert_called_once()
 

hud/tools/tests/test_tools_init.py

@@ -33,7 +33,7 @@ class TestToolsInit:
         """Test lazy import with invalid attribute name."""
         import hud.tools as tools_module
 
-        with pytest.raises(AttributeError, match="module '.*' has no attribute 'InvalidTool'"):
+        with pytest.raises(AttributeError, match=r"module '.*' has no attribute 'InvalidTool'"):
             _ = tools_module.InvalidTool
 
     def test_direct_imports_available(self):

hud/tools/tests/test_utils.py

@@ -58,7 +58,7 @@ class TestRun:
         mock_proc.communicate = AsyncMock(return_value=(b"processed", b""))
 
         with patch("asyncio.create_subprocess_shell", return_value=mock_proc):
-            return_code, stdout, stderr = await run("cat", input="test input")
+            return_code, stdout, _stderr = await run("cat", input="test input")
 
             assert return_code == 0
             assert stdout == "processed"
@@ -91,7 +91,7 @@ class TestRun:
         ):
             mock_wait_for.return_value = (b"done", b"")
 
-            return_code, stdout, stderr = await run("sleep 1", timeout=5.0)
+            _return_code, _stdout, _stderr = await run("sleep 1", timeout=5.0)
 
             # Check that wait_for was called with the correct timeout
             mock_wait_for.assert_called_once()

hud/types.py

@@ -29,9 +29,9 @@ class MCPToolCall(CallToolRequestParams):
 
     def __rich__(self) -> str:
         """Rich representation with color formatting."""
-        from hud.utils.design import design
+        from hud.utils.hud_console import hud_console
 
-        return design.format_tool_call(self.name, self.arguments)
+        return hud_console.format_tool_call(self.name, self.arguments)
 
 
 class MCPToolResult(CallToolResult):
@@ -74,10 +74,10 @@ class MCPToolResult(CallToolResult):
 
     def __rich__(self) -> str:
         """Rich representation with color formatting."""
-        from hud.utils.design import design
+        from hud.utils.hud_console import hud_console
 
         content_summary = self._get_content_summary()
-        return design.format_tool_result(content_summary, self.isError)
+        return hud_console.format_tool_result(content_summary, self.isError)
 
 
 class AgentResponse(BaseModel):

hud/utils/__init__.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from .design import HUDDesign, design
+from .hud_console import HUDConsole, hud_console
 from .telemetry import stream
 
 __all__ = [
-    "HUDDesign",
-    "design",
+    "HUDConsole",
+    "hud_console",
     "stream",
 ]

hud/utils/agent_factories.py

@@ -0,0 +1,86 @@
+"""Factory functions for creating agents compatible with run_dataset."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from openai import AsyncOpenAI
+
+from hud.agents.grounded_openai import GroundedOpenAIChatAgent
+from hud.agents.openai_chat_generic import GenericOpenAIChatAgent
+from hud.tools.grounding import GrounderConfig
+
+
+def create_openai_agent(**kwargs: Any) -> GenericOpenAIChatAgent:
+    """Factory for GenericOpenAIChatAgent with run_dataset compatibility.
+
+    Args:
+        api_key: OpenAI API key
+        base_url: Optional custom API endpoint
+        model_name: Model to use (e.g., "gpt-4o-mini")
+        **kwargs: Additional arguments passed to GenericOpenAIChatAgent
+
+    Returns:
+        Configured GenericOpenAIChatAgent instance
+
+    Example:
+        >>> from hud.datasets import run_dataset
+        >>> from hud.utils.agent_factories import create_openai_agent
+        >>> results = await run_dataset(
+        ...     "My Eval",
+        ...     "hud-evals/SheetBench-50",
+        ...     create_openai_agent,
+        ...     {"api_key": "your-key", "model_name": "gpt-4o-mini"},
+        ... )
+    """
+    api_key = kwargs.pop("api_key", None)
+    base_url = kwargs.pop("base_url", None)
+
+    openai_client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+
+    return GenericOpenAIChatAgent(openai_client=openai_client, **kwargs)
+
+
+def create_grounded_agent(**kwargs: Any) -> GroundedOpenAIChatAgent:
+    """Factory for GroundedOpenAIChatAgent with run_dataset compatibility.
+
+    Args:
+        api_key: OpenAI API key for planning model
+        base_url: Optional custom API endpoint for planning model
+        model_name: Planning model to use (e.g., "gpt-4o-mini")
+        grounder_api_key: API key for grounding model
+        grounder_api_base: API base URL for grounding model (default: OpenRouter)
+        grounder_model: Grounding model to use (default: qwen/qwen-2.5-vl-7b-instruct)
+        **kwargs: Additional arguments passed to GroundedOpenAIChatAgent
+
+    Returns:
+        Configured GroundedOpenAIChatAgent instance
+
+    Example:
+        >>> from hud.datasets import run_dataset
+        >>> from hud.utils.agent_factories import create_grounded_agent
+        >>> results = await run_dataset(
+        ...     "Grounded Eval",
+        ...     dataset,
+        ...     create_grounded_agent,
+        ...     {
+        ...         "api_key": "openai-key",
+        ...         "grounder_api_key": "openrouter-key",
+        ...         "model_name": "gpt-4o-mini",
+        ...     },
+        ... )
+    """
+    api_key = kwargs.pop("api_key", None)
+    base_url = kwargs.pop("base_url", None)
+    grounder_api_key = kwargs.pop("grounder_api_key", None)
+    grounder_api_base = kwargs.pop("grounder_api_base", "https://openrouter.ai/api/v1")
+    grounder_model = kwargs.pop("grounder_model", "qwen/qwen-2.5-vl-7b-instruct")
+
+    openai_client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+    grounder_config = GrounderConfig(
+        api_base=grounder_api_base, model=grounder_model, api_key=grounder_api_key
+    )
+
+    return GroundedOpenAIChatAgent(
+        openai_client=openai_client, grounder_config=grounder_config, **kwargs
+    )