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

hud/tools/grounding/grounded_tool.py

@@ -3,14 +3,15 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, 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
+
+if TYPE_CHECKING:
+    from hud.environment import Environment
 
 logger = logging.getLogger(__name__)
 
@@ -33,18 +34,18 @@ class GroundedComputerTool:
         self,
         *,
         grounder: Grounder,
-        mcp_client: AgentMCPClient,
+        ctx: Environment,
         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
+            ctx: Environment or EvalContext to call tools through
             computer_tool_name: Name of the computer tool in the environment
         """
         self._grounder = grounder
-        self._mcp_client = mcp_client
+        self._ctx = ctx
         self._computer_tool_name = computer_tool_name
 
     def get_openai_tool_schema(self) -> dict:
@@ -172,10 +173,8 @@ class GroundedComputerTool:
                 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}
-                    )
+                result = await self._ctx.call_tool(
+                    (self._computer_tool_name, {**computer_args, **kwargs})
                 )
                 return result.content
 
@@ -224,10 +223,8 @@ class GroundedComputerTool:
                 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}
-                    )
+                result = await self._ctx.call_tool(
+                    (self._computer_tool_name, {**computer_args, **kwargs})
                 )
                 return result.content
 
@@ -292,10 +289,8 @@ class GroundedComputerTool:
                 if button:
                     computer_args["button"] = button
 
-                result = await self._mcp_client.call_tool(
-                    MCPToolCall(
-                        name=self._computer_tool_name, arguments={**computer_args, **kwargs}
-                    )
+                result = await self._ctx.call_tool(
+                    (self._computer_tool_name, {**computer_args, **kwargs})
                 )
                 return result.content
 

hud/tools/grounding/grounder.py

@@ -4,15 +4,15 @@ from __future__ import annotations
 
 import base64
 import io
-import json
+import logging
 import re
 
 from openai import AsyncOpenAI
-from opentelemetry import trace
 
-from hud import instrument
 from hud.tools.grounding.config import GrounderConfig  # noqa: TC001
 
+logger = logging.getLogger(__name__)
+
 
 class Grounder:
     """Grounder that uses AsyncOpenAI to call vLLM or other model endpoints for visual grounding.
@@ -181,12 +181,6 @@ class Grounder:
 
         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:
@@ -247,12 +241,7 @@ class Grounder:
 
                 # 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)
+                logger.debug("Grounder attempt %d response: %s", attempt + 1, response_text)
 
                 # Parse coordinates from response
                 if response_text is None:
@@ -277,26 +266,16 @@ class Grounder:
                     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)
-
+                logger.debug(
+                    "Grounder success: coords=%s after %d attempts",
+                    pixel_coords,
+                    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")
-
+        logger.debug("Grounder failed after %d attempts", max_retries)
         return None

hud/tools/grounding/tests/test_grounded_tool.py

@@ -7,7 +7,7 @@ import mcp.types as types
 import pytest
 
 from hud.tools.grounding.grounded_tool import GroundedComputerTool
-from hud.types import MCPToolCall, MCPToolResult
+from hud.types import MCPToolResult
 
 
 @dataclass
@@ -17,36 +17,18 @@ class FakeResult:
     structuredContent: dict | None = None
 
 
-class FakeMCPClient:
-    """Fake MCP client that implements AgentMCPClient protocol."""
-
-    _initialized: bool
+class FakeEnvironment:
+    """Fake Environment that implements the call_tool interface."""
 
     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 {}))
+    async def call_tool(self, call: tuple[str, dict[str, Any]], /, **kwargs: Any) -> MCPToolResult:
+        """Record the tool call and return a fake result."""
+        tool_name, tool_args = call
+        self.calls.append((tool_name, tool_args))
         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."""
@@ -72,9 +54,9 @@ def _png_b64() -> str:
 
 @pytest.mark.asyncio
 async def test_click_action_grounds_and_calls_mcp() -> None:
-    client = FakeMCPClient()
+    ctx = FakeEnvironment()
     grounder = FakeGrounder(coords=(123, 456))
-    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+    tool = GroundedComputerTool(grounder=grounder, ctx=ctx)  # type: ignore
 
     blocks = await tool(
         action="click",
@@ -87,14 +69,14 @@ async def test_click_action_grounds_and_calls_mcp() -> None:
     # 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"})]
+    assert ctx.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()
+    ctx = FakeEnvironment()
     grounder = FakeGrounder(coords=(5, 6))
-    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+    tool = GroundedComputerTool(grounder=grounder, ctx=ctx)  # type: ignore
 
     # Missing element_description
     with pytest.raises(Exception) as ei:
@@ -109,9 +91,9 @@ async def test_move_and_scroll_require_element_description_and_screenshot() -> N
 
 @pytest.mark.asyncio
 async def test_drag_grounds_both_points_and_calls_mcp() -> None:
-    client = FakeMCPClient()
+    ctx = FakeEnvironment()
     grounder = FakeGrounder(coords=(10, 20))
-    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+    tool = GroundedComputerTool(grounder=grounder, ctx=ctx)  # type: ignore
 
     await tool(
         action="drag",
@@ -124,7 +106,7 @@ async def test_drag_grounds_both_points_and_calls_mcp() -> None:
     # 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]
+    name, args = ctx.calls[0]
     assert name == "computer"
     assert args["action"] == "drag"
     assert args["button"] == "left"
@@ -133,9 +115,9 @@ async def test_drag_grounds_both_points_and_calls_mcp() -> None:
 
 @pytest.mark.asyncio
 async def test_drag_requires_both_descriptions_and_screenshot() -> None:
-    client = FakeMCPClient()
+    ctx = FakeEnvironment()
     grounder = FakeGrounder()
-    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+    tool = GroundedComputerTool(grounder=grounder, ctx=ctx)  # type: ignore
 
     with pytest.raises(Exception) as ei:
         await tool(action="drag", start_element_description="a", screenshot_b64=_png_b64())
@@ -152,9 +134,9 @@ async def test_drag_requires_both_descriptions_and_screenshot() -> None:
 
 @pytest.mark.asyncio
 async def test_direct_actions_bypass_grounding_and_call_mcp() -> None:
-    client = FakeMCPClient()
+    ctx = FakeEnvironment()
     grounder = FakeGrounder()
-    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+    tool = GroundedComputerTool(grounder=grounder, ctx=ctx)  # type: ignore
 
     # Actions that bypass grounding
     for action, extra in [
@@ -166,19 +148,19 @@ async def test_direct_actions_bypass_grounding_and_call_mcp() -> None:
         ("get_dimensions", {}),
         ("get_environment", {}),
     ]:
-        client.calls.clear()
+        ctx.calls.clear()
         _ = await tool(action=action, **extra)
-        assert client.calls and client.calls[0][0] == "computer"
-        assert client.calls[0][1]["action"] == action
+        assert ctx.calls and ctx.calls[0][0] == "computer"
+        assert ctx.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()
+    ctx = FakeEnvironment()
     grounder = FakeGrounder()
-    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+    tool = GroundedComputerTool(grounder=grounder, ctx=ctx)  # type: ignore
 
     with pytest.raises(Exception) as ei:
         await tool(action="zoom")
@@ -187,9 +169,9 @@ async def test_unsupported_action_raises() -> None:
 
 @pytest.mark.asyncio
 async def test_grounding_failure_propagates_as_error() -> None:
-    client = FakeMCPClient()
+    ctx = FakeEnvironment()
     grounder = FakeGrounder(coords=None)
-    tool = GroundedComputerTool(grounder=grounder, mcp_client=client)  # type: ignore
+    tool = GroundedComputerTool(grounder=grounder, ctx=ctx)  # type: ignore
 
     with pytest.raises(Exception) as ei:
         await tool(action="click", element_description="x", screenshot_b64=_png_b64())

hud/tools/jupyter.py

@@ -0,0 +1,330 @@
+"""Jupyter execution tool.
+
+Requires the [agents] extra: pip install hud-python[agents]
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+from typing import TYPE_CHECKING, Any, ClassVar
+from uuid import uuid4
+
+from hud.tools.base import BaseTool
+from hud.tools.types import ContentResult, ToolError
+
+if TYPE_CHECKING:
+    from mcp.types import ContentBlock
+
+logger = logging.getLogger(__name__)
+
+
+def strip_ansi(output: str) -> str:
+    """Remove ANSI escape sequences from string output."""
+    pattern = re.compile(r"\x1B\[\d+(;\d+){0,2}m")
+    return pattern.sub("", output)
+
+
+class JupyterTool(BaseTool):
+    """
+    Execute Python code in a Jupyter kernel.
+    """
+
+    # Class-level kernel registry for sharing kernels
+    _kernel_registry: ClassVar[dict[str, str]] = {}
+
+    @classmethod
+    def register_shared_kernel(cls, registry_name: str, kernel_id: str) -> None:
+        """Register a kernel_id with a name for reuse.
+
+        Args:
+            registry_name: Name to register the kernel under
+            kernel_id: The kernel ID to register
+        """
+        cls._kernel_registry[registry_name] = kernel_id
+        logger.info("Registered kernel '%s': %s", registry_name, kernel_id)
+
+    @classmethod
+    def from_shared_kernel(cls, registry_name: str, **kwargs: Any) -> JupyterTool:
+        """Connect to a kernel using its registry name.
+
+        Args:
+            registry_name: Name of the registered kernel
+            **kwargs: Additional parameters for JupyterTool (url_suffix, kernel_name)
+
+        Returns:
+            JupyterTool instance connected to the registered kernel
+        """
+        kernel_id = cls._kernel_registry.get(registry_name)
+        if not kernel_id:
+            raise ValueError(f"No kernel registered with name '{registry_name}'")
+
+        logger.info("Connecting to registered kernel '%s': %s", registry_name, kernel_id)
+        return cls(kernel_id=kernel_id, **kwargs)
+
+    def __init__(
+        self,
+        url_suffix: str = "localhost:8888",
+        kernel_name: str = "python3",
+        kernel_id: str = "",
+    ) -> None:
+        """Initialize JupyterTool with connection parameters.
+
+        Args:
+            url_suffix: (Optional) Kernel gateway host:port (default: localhost:8888)
+            kernel_name: (Optional) Kernel name to use (default: python3)
+            kernel_id: (Optional) If set, connect to the existed kernel with kernel_id.
+                If empty, create new kernel
+        """
+        # Check tornado is available
+        try:
+            import tornado  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "JupyterTool requires the [agents] extra. "
+                "Install with: pip install hud-python[agents]"
+            ) from e
+
+        super().__init__(
+            env=None,
+            name="jupyter",
+            title="Jupyter Code Execution",
+            description="Execute Python code in a Jupyter kernel",
+        )
+
+        # Connection parameters
+        self._base_url = f"http://{url_suffix}"
+        self._base_ws_url = f"ws://{url_suffix}"
+        self._kernel_name = kernel_name
+
+        # Kernel state (reuse existing or create new)
+        self._kernel_id = kernel_id
+        self._ws: Any = None
+        self._initialized = False
+
+        # WebSocket heartbeat
+        self._heartbeat_interval = 10000  # 10 seconds
+        self._heartbeat_callback: Any = None
+
+    async def __call__(self, code: str, execution_timeout: int = 15) -> list[ContentBlock]:
+        """Execute Python code in the Jupyter kernel.
+
+        Args:
+            code: Python code to execute
+            execution_timeout: Execution timeout in seconds (default: 15)
+
+        Returns:
+            List of ContentBlock with execution results
+        """
+        try:
+            # Ensure kernel is ready (lazy initialization)
+            await self._ensure_kernel()
+
+            # Execute code
+            result = await self._execute(code, execution_timeout)
+
+            # Check for timeout
+            if result.startswith("[Execution timed out"):
+                return ContentResult(error=result).to_content_blocks()
+
+            # Return result
+            output = result if result.strip() else "Code executed successfully (no output)"
+            return ContentResult(output=output).to_content_blocks()
+
+        except Exception as e:
+            logger.error("Jupyter execution error: %s", e)
+            raise ToolError(f"Execution failed: {e!s}") from e
+
+    async def _ensure_kernel(self) -> None:
+        """Ensure kernel is initialized and connected."""
+        if not self._initialized:
+            logger.info("Initializing Jupyter kernel connection")
+            await self._connect()
+            self._initialized = True
+            logger.info("Jupyter kernel connected successfully")
+
+    async def _connect(self) -> None:
+        """Connect to Jupyter kernel via WebSocket."""
+        import tornado.iostream
+        from tornado.escape import json_decode, json_encode, url_escape
+        from tornado.httpclient import AsyncHTTPClient, HTTPRequest
+        from tornado.ioloop import PeriodicCallback
+        from tornado.websocket import websocket_connect
+
+        if self._ws:
+            self._ws.close()
+            self._ws = None
+
+        client = AsyncHTTPClient()
+        if not self._kernel_id:
+            # Start a new kernel
+            n_tries = 5
+            while n_tries > 0:
+                try:
+                    response = await client.fetch(
+                        f"{self._base_url}/api/kernels",
+                        method="POST",
+                        body=json_encode({"name": self._kernel_name}),
+                    )
+                    kernel = json_decode(response.body)
+                    self._kernel_id = kernel["id"]
+                    logger.info("Kernel started with ID: %s", self._kernel_id)
+                    break
+                except Exception as e:
+                    logger.warning("Kernel connection attempt failed: %s", e)
+                    n_tries -= 1
+                    await asyncio.sleep(1)
+
+            if n_tries == 0:
+                raise ConnectionRefusedError("Failed to connect to kernel gateway")
+
+        # Connect WebSocket to kernel
+        ws_req = HTTPRequest(
+            url=f"{self._base_ws_url}/api/kernels/{url_escape(self._kernel_id)}/channels"
+        )
+        self._ws = await websocket_connect(ws_req)
+        logger.info("WebSocket connected to kernel")
+
+        # Setup heartbeat to keep connection alive
+        if self._heartbeat_callback:
+            self._heartbeat_callback.stop()
+
+        async def heartbeat() -> None:
+            if not self._ws:
+                return
+            try:
+                self._ws.ping()
+            except tornado.iostream.StreamClosedError:
+                try:
+                    await self._connect()
+                except ConnectionRefusedError:
+                    logger.warning(
+                        "Failed to reconnect to kernel websocket - Is the kernel still running?"
+                    )
+
+        self._heartbeat_callback = PeriodicCallback(heartbeat, self._heartbeat_interval)
+        self._heartbeat_callback.start()
+
+    async def _execute(self, code: str, execution_timeout: int = 15) -> str:
+        """Execute code in Jupyter kernel and return output.
+
+        Args:
+            code: Python code to execute
+            execution_timeout: Execution timeout in seconds
+
+        Returns:
+            String output from the kernel
+        """
+        from tornado.escape import json_decode, json_encode
+        from tornado.httpclient import AsyncHTTPClient
+
+        if not self._ws:
+            await self._connect()
+
+        msg_id = uuid4().hex
+        self._ws.write_message(
+            json_encode(
+                {
+                    "header": {
+                        "username": "",
+                        "version": "5.0",
+                        "session": "",
+                        "msg_id": msg_id,
+                        "msg_type": "execute_request",
+                    },
+                    "parent_header": {},
+                    "channel": "shell",
+                    "content": {
+                        "code": code,
+                        "silent": False,
+                        "store_history": False,
+                        "user_expressions": {},
+                        "allow_stdin": False,
+                    },
+                    "metadata": {},
+                    "buffers": {},
+                }
+            )
+        )
+
+        outputs: list[str] = []
+
+        async def wait_for_messages() -> bool:
+            execution_done = False
+            while not execution_done:
+                msg = await self._ws.read_message()
+                msg = json_decode(msg)
+                msg_type = msg["msg_type"]
+                parent_msg_id = msg["parent_header"].get("msg_id", None)
+
+                if parent_msg_id != msg_id:
+                    continue
+
+                if msg_type == "error":
+                    traceback = "\n\n\n\n".join(msg["content"]["traceback"])
+                    outputs.append(traceback)
+                    execution_done = True
+                elif msg_type == "stream":
+                    outputs.append(msg["content"]["text"])
+                elif msg_type in ["execute_result", "display_data"]:
+                    outputs.append(msg["content"]["data"]["text/plain"])
+                    # Handle image outputs
+                    if "image/png" in msg["content"]["data"]:
+                        outputs.append(
+                            f"![image](data:image/png;base64,{msg['content']['data']['image/png']})"
+                        )
+                elif msg_type == "execute_reply":
+                    execution_done = True
+            return execution_done
+
+        async def interrupt_kernel() -> None:
+            client = AsyncHTTPClient()
+            interrupt_response = await client.fetch(
+                f"{self._base_url}/api/kernels/{self._kernel_id}/interrupt",
+                method="POST",
+                body=json_encode({"kernel_id": self._kernel_id}),
+            )
+            logger.info("Kernel interrupted: %s", interrupt_response)
+
+        try:
+            await asyncio.wait_for(wait_for_messages(), execution_timeout)
+        except TimeoutError:
+            await interrupt_kernel()
+            return f"[Execution timed out ({execution_timeout} seconds).]"
+
+        ret = "".join(outputs)
+
+        # Remove ANSI escape sequences
+        return strip_ansi(ret)
+
+    async def shutdown(self) -> None:
+        """Shutdown the kernel connection."""
+        from tornado.httpclient import AsyncHTTPClient
+
+        if self._kernel_id:
+            client = AsyncHTTPClient()
+            try:
+                await client.fetch(
+                    f"{self._base_url}/api/kernels/{self._kernel_id}",
+                    method="DELETE",
+                )
+                logger.info("Kernel %s shut down", self._kernel_id)
+            except Exception as e:
+                logger.warning("Error shutting down kernel: %s", e)
+
+            self._kernel_id = ""
+
+            if self._heartbeat_callback:
+                self._heartbeat_callback.stop()
+                self._heartbeat_callback = None
+
+            if self._ws:
+                self._ws.close()
+                self._ws = None
+
+        self._initialized = False
+
+    def get_kernel_id(self) -> str:
+        """Get the jupyter kernel id."""
+        return self._kernel_id