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

hud/agents/tests/test_client.py

@@ -15,7 +15,6 @@ from hud.types import MCPToolResult
 logger = logging.getLogger(__name__)
 
 
-@patch("hud.clients.base.setup_hud_telemetry")
 class TestMCPClient:
     """Test MCPClient class."""
 
@@ -34,7 +33,7 @@ class TestMCPClient:
             yield mock_instance
 
     @pytest.mark.asyncio
-    async def test_connect_single_server(self, mock_telemetry, mock_mcp_use_client):
+    async def test_connect_single_server(self, mock_mcp_use_client):
         """Test connecting to a single server."""
         config = {"test_server": {"command": "python", "args": ["-m", "test_server"]}}
 
@@ -77,7 +76,7 @@ class TestMCPClient:
         assert names == {"tool1", "tool2"}
 
     @pytest.mark.asyncio
-    async def test_connect_multiple_servers(self, mock_telemetry, mock_mcp_use_client):
+    async def test_connect_multiple_servers(self, mock_mcp_use_client):
         """Test connecting to multiple servers."""
         config = {
             "server1": {"command": "python", "args": ["-m", "server1"]},
@@ -129,7 +128,7 @@ class TestMCPClient:
         assert names == {"server1_tool1", "server2_tool2"}
 
     @pytest.mark.asyncio
-    async def test_call_tool(self, mock_telemetry, mock_mcp_use_client):
+    async def test_call_tool(self, mock_mcp_use_client):
         """Test calling a tool."""
         config = {"test": {"command": "test"}}
         client = MCPClient(mcp_config=config)
@@ -180,7 +179,7 @@ class TestMCPClient:
         )
 
     @pytest.mark.asyncio
-    async def test_call_tool_not_found(self, mock_telemetry, mock_mcp_use_client):
+    async def test_call_tool_not_found(self, mock_mcp_use_client):
         """Test calling a non-existent tool."""
         config = {"test": {"command": "test"}}
         client = MCPClient(mcp_config=config)
@@ -208,7 +207,7 @@ class TestMCPClient:
         assert "Tool 'nonexistent' not found" in text_content
 
     @pytest.mark.asyncio
-    async def test_get_telemetry_data(self, mock_telemetry, mock_mcp_use_client):
+    async def test_get_telemetry_data(self, mock_mcp_use_client):
         """Test getting telemetry data."""
         config = {"test": {"command": "test"}}
         client = MCPClient(mcp_config=config)
@@ -245,7 +244,7 @@ class TestMCPClient:
         assert isinstance(telemetry_data, dict)
 
     @pytest.mark.asyncio
-    async def test_close(self, mock_telemetry, mock_mcp_use_client):
+    async def test_close(self, mock_mcp_use_client):
         """Test closing client connections."""
         config = {"test": {"command": "test"}}
         client = MCPClient(mcp_config=config)
@@ -267,7 +266,7 @@ class TestMCPClient:
         mock_mcp_use_client.close_all_sessions.assert_called_once()
 
     @pytest.mark.asyncio
-    async def test_context_manager(self, mock_telemetry, mock_mcp_use_client):
+    async def test_context_manager(self, mock_mcp_use_client):
         """Test using client as context manager."""
         mock_session = MagicMock()
         mock_session.connector = MagicMock()
@@ -291,7 +290,7 @@ class TestMCPClient:
         mock_mcp_use_client.close_all_sessions.assert_called_once()
 
     @pytest.mark.asyncio
-    async def test_get_available_tools(self, mock_telemetry, mock_mcp_use_client):
+    async def test_get_available_tools(self, mock_mcp_use_client):
         """Test getting available tools."""
         config = {"test": {"command": "test"}}
         client = MCPClient(mcp_config=config)
@@ -319,7 +318,7 @@ class TestMCPClient:
         assert names == {"tool1", "tool2"}
 
     @pytest.mark.asyncio
-    async def test_get_tool_map(self, mock_telemetry, mock_mcp_use_client):
+    async def test_get_tool_map(self, mock_mcp_use_client):
         """Test getting tool map."""
         config = {"test": {"command": "test"}}
         client = MCPClient(mcp_config=config)

hud/agents/tests/test_gemini.py

@@ -0,0 +1,369 @@
+"""Tests for Gemini MCP Agent implementation."""
+
+from __future__ import annotations
+
+import base64
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from google import genai
+from google.genai import types as genai_types
+from mcp import types
+
+from hud.agents.gemini import GeminiAgent
+from hud.environment.router import ToolRouter
+from hud.eval.context import EvalContext
+from hud.types import MCPToolCall, MCPToolResult
+
+
+class MockEvalContext(EvalContext):
+    """Mock EvalContext for testing."""
+
+    def __init__(self, tools: list[types.Tool] | None = None) -> None:
+        # Core attributes
+        self.prompt = "Test prompt"
+        self._tools = tools or []
+        self._submitted: str | None = None
+        self.reward: float | None = None
+
+        # Environment attributes
+        self._router = ToolRouter()
+        self._agent_include: list[str] | None = None
+        self._agent_exclude: list[str] | None = None
+
+        # EvalContext attributes
+        self._task = None
+        self.trace_id = "test-trace-id"
+        self.eval_name = "test-eval"
+        self.job_id: str | None = None
+        self.group_id: str | None = None
+        self.index = 0
+        self.variants: dict[str, Any] = {}
+        self.answer: str | None = None
+        self.system_prompt: str | None = None
+        self.error: BaseException | None = None
+        self.metadata: dict[str, Any] = {}
+        self.results: list[Any] = []
+        self._is_summary = False
+
+    def as_tools(self) -> list[types.Tool]:
+        return self._tools
+
+    @property
+    def has_scenario(self) -> bool:
+        return False
+
+    async def list_tools(self) -> list[types.Tool]:
+        return self._tools
+
+    async def call_tool(self, call: Any, /, **kwargs: Any) -> MCPToolResult:
+        return MCPToolResult(
+            content=[types.TextContent(type="text", text="ok")],
+            isError=False,
+        )
+
+    async def submit(self, answer: str) -> None:
+        self._submitted = answer
+
+
+class TestGeminiAgent:
+    """Test GeminiAgent base class."""
+
+    @pytest.fixture
+    def mock_gemini_client(self) -> MagicMock:
+        """Create a stub Gemini client."""
+        client = MagicMock(spec=genai.Client)
+        client.api_key = "test_key"
+        client.models = MagicMock()
+        client.models.list = MagicMock(return_value=iter([]))
+        client.models.generate_content = MagicMock()
+        # Set up async interface (aio.models.generate_content)
+        client.aio = MagicMock()
+        client.aio.models = MagicMock()
+        client.aio.models.generate_content = AsyncMock()
+        return client
+
+    @pytest.mark.asyncio
+    async def test_init(self, mock_gemini_client: MagicMock) -> None:
+        """Test agent initialization."""
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            model="gemini-2.5-flash",
+            validate_api_key=False,
+        )
+
+        assert agent.model_name == "Gemini"
+        assert agent.config.model == "gemini-2.5-flash"
+        assert agent.gemini_client == mock_gemini_client
+
+    @pytest.mark.asyncio
+    async def test_init_without_model_client(self) -> None:
+        """Test agent initialization without model client."""
+        with (
+            patch("hud.settings.settings.gemini_api_key", "test_key"),
+            patch("hud.agents.gemini.genai.Client") as mock_client_class,
+        ):
+            mock_client = MagicMock()
+            mock_client.api_key = "test_key"
+            mock_client.models = MagicMock()
+            mock_client.models.list = MagicMock(return_value=iter([]))
+            mock_client_class.return_value = mock_client
+
+            agent = GeminiAgent.create(
+                model="gemini-2.5-flash",
+                validate_api_key=False,
+            )
+
+            assert agent.gemini_client is not None
+
+    @pytest.mark.asyncio
+    async def test_format_blocks_text_only(self, mock_gemini_client: MagicMock) -> None:
+        """Test formatting text content blocks."""
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            validate_api_key=False,
+        )
+
+        blocks: list[types.ContentBlock] = [
+            types.TextContent(type="text", text="Hello, world!"),
+            types.TextContent(type="text", text="How are you?"),
+        ]
+
+        messages = await agent.format_blocks(blocks)
+        assert len(messages) == 1
+        assert messages[0].role == "user"
+        assert messages[0].parts is not None
+        assert len(messages[0].parts) == 2
+
+    @pytest.mark.asyncio
+    async def test_format_blocks_with_image(self, mock_gemini_client: MagicMock) -> None:
+        """Test formatting image content blocks."""
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            validate_api_key=False,
+        )
+
+        # Create a tiny valid base64 PNG
+        png_data = base64.b64encode(b"\x89PNG\r\n\x1a\n").decode()
+
+        blocks: list[types.ContentBlock] = [
+            types.TextContent(type="text", text="Look at this:"),
+            types.ImageContent(type="image", data=png_data, mimeType="image/png"),
+        ]
+
+        messages = await agent.format_blocks(blocks)
+        assert len(messages) == 1
+        assert messages[0].parts is not None
+        assert len(messages[0].parts) == 2
+
+    @pytest.mark.asyncio
+    async def test_format_tool_results(self, mock_gemini_client: MagicMock) -> None:
+        """Test formatting tool results."""
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            validate_api_key=False,
+        )
+
+        tool_calls = [MCPToolCall(id="call_123", name="test_tool", arguments={})]
+        tool_results = [
+            MCPToolResult(
+                content=[types.TextContent(type="text", text="Tool output")],
+                isError=False,
+            )
+        ]
+
+        messages = await agent.format_tool_results(tool_calls, tool_results)
+        assert len(messages) == 1
+        assert messages[0].role == "user"
+
+    @pytest.mark.asyncio
+    async def test_get_system_messages(self, mock_gemini_client: MagicMock) -> None:
+        """Test that system messages return empty (Gemini uses system_instruction)."""
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            system_prompt="You are a helpful assistant.",
+            validate_api_key=False,
+        )
+
+        messages = await agent.get_system_messages()
+        # Gemini doesn't use system messages in the message list
+        assert messages == []
+
+    @pytest.mark.asyncio
+    async def test_get_response_text_only(self, mock_gemini_client: MagicMock) -> None:
+        """Test getting text-only response."""
+        # Disable telemetry for this test
+        with patch("hud.settings.settings.telemetry_enabled", False):
+            agent = GeminiAgent.create(
+                model_client=mock_gemini_client,
+                validate_api_key=False,
+            )
+            # Set up agent as initialized (no tools needed for this test)
+            agent.gemini_tools = []
+            agent._initialized = True
+
+            # Mock the API response with text only
+            mock_response = MagicMock()
+            mock_candidate = MagicMock()
+
+            text_part = MagicMock()
+            text_part.text = "Task completed successfully"
+            text_part.function_call = None
+
+            mock_candidate.content = MagicMock()
+            mock_candidate.content.parts = [text_part]
+
+            mock_response.candidates = [mock_candidate]
+
+            mock_gemini_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
+
+            messages = [
+                genai_types.Content(role="user", parts=[genai_types.Part.from_text(text="Status?")])
+            ]
+            response = await agent.get_response(messages)
+
+            assert response.content == "Task completed successfully"
+            assert response.tool_calls == []
+            assert response.done is True
+
+    @pytest.mark.asyncio
+    async def test_get_response_with_thinking(self, mock_gemini_client: MagicMock) -> None:
+        """Test getting response with thinking content."""
+        with patch("hud.settings.settings.telemetry_enabled", False):
+            agent = GeminiAgent.create(
+                model_client=mock_gemini_client,
+                validate_api_key=False,
+            )
+            # Set up agent as initialized (no tools needed for this test)
+            agent.gemini_tools = []
+            agent._initialized = True
+
+            mock_response = MagicMock()
+            mock_candidate = MagicMock()
+
+            thinking_part = MagicMock()
+            thinking_part.text = "Let me reason through this..."
+            thinking_part.function_call = None
+            thinking_part.thought = True
+
+            text_part = MagicMock()
+            text_part.text = "Here is my answer"
+            text_part.function_call = None
+            text_part.thought = False
+
+            mock_candidate.content = MagicMock()
+            mock_candidate.content.parts = [thinking_part, text_part]
+
+            mock_response.candidates = [mock_candidate]
+
+            mock_gemini_client.aio.models.generate_content = AsyncMock(return_value=mock_response)
+
+            messages = [
+                genai_types.Content(
+                    role="user", parts=[genai_types.Part.from_text(text="Hard question")]
+                )
+            ]
+            response = await agent.get_response(messages)
+
+            assert response.content == "Here is my answer"
+            assert response.reasoning == "Let me reason through this..."
+
+    @pytest.mark.asyncio
+    async def test_convert_tools_for_gemini(self, mock_gemini_client: MagicMock) -> None:
+        """Test converting MCP tools to Gemini format."""
+        tools = [
+            types.Tool(
+                name="my_tool",
+                description="A test tool",
+                inputSchema={"type": "object", "properties": {"x": {"type": "string"}}},
+            )
+        ]
+        ctx = MockEvalContext(tools=tools)
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            validate_api_key=False,
+        )
+
+        agent.ctx = ctx
+        await agent._initialize_from_ctx(ctx)
+
+        # Check that tools were converted
+        assert len(agent.gemini_tools) == 1
+        # Gemini tools have function_declarations - cast to genai Tool type
+        gemini_tool = agent.gemini_tools[0]
+        assert isinstance(gemini_tool, genai_types.Tool)
+        assert gemini_tool.function_declarations is not None
+        assert gemini_tool.function_declarations[0].name == "my_tool"
+
+
+class TestGeminiToolConversion:
+    """Tests for tool conversion to Gemini format."""
+
+    @pytest.fixture
+    def mock_gemini_client(self) -> MagicMock:
+        """Create a stub Gemini client."""
+        client = MagicMock(spec=genai.Client)
+        client.api_key = "test_key"
+        client.models = MagicMock()
+        client.models.list = MagicMock(return_value=iter([]))
+        # Set up async interface
+        client.aio = MagicMock()
+        client.aio.models = MagicMock()
+        client.aio.models.generate_content = AsyncMock()
+        return client
+
+    @pytest.mark.asyncio
+    async def test_tool_with_properties(self, mock_gemini_client: MagicMock) -> None:
+        """Test tool with input properties."""
+        tools = [
+            types.Tool(
+                name="search",
+                description="Search the web",
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string", "description": "Search query"},
+                        "limit": {"type": "integer", "description": "Max results"},
+                    },
+                    "required": ["query"],
+                },
+            )
+        ]
+        ctx = MockEvalContext(tools=tools)
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            validate_api_key=False,
+        )
+
+        agent.ctx = ctx
+        await agent._initialize_from_ctx(ctx)
+
+        assert len(agent.gemini_tools) == 1
+        gemini_tool = agent.gemini_tools[0]
+        # Gemini tools have function_declarations - cast to genai Tool type
+        assert isinstance(gemini_tool, genai_types.Tool)
+        assert gemini_tool.function_declarations is not None
+        assert gemini_tool.function_declarations[0].name == "search"
+        assert gemini_tool.function_declarations[0].parameters_json_schema is not None
+
+    @pytest.mark.asyncio
+    async def test_tool_without_schema(self, mock_gemini_client: MagicMock) -> None:
+        """Test tool without description raises error."""
+        # Create a tool with inputSchema but no description
+        tools = [
+            types.Tool(
+                name="incomplete",
+                description=None,
+                inputSchema={"type": "object"},
+            )
+        ]
+        ctx = MockEvalContext(tools=tools)
+        agent = GeminiAgent.create(
+            model_client=mock_gemini_client,
+            validate_api_key=False,
+        )
+
+        agent.ctx = ctx
+        with pytest.raises(ValueError, match="requires both a description"):
+            await agent._initialize_from_ctx(ctx)

hud/agents/tests/test_grounded_openai_agent.py

@@ -1,60 +1,16 @@
 from __future__ import annotations
 
-import json
 from typing import Any
 
 import mcp.types as types
 import pytest
+from openai import AsyncOpenAI
 
 from hud.agents.grounded_openai import GroundedOpenAIChatAgent
 from hud.tools.grounding import GrounderConfig
 from hud.types import MCPToolCall, MCPToolResult
 
 
-class DummyOpenAI:
-    class chat:  # type: ignore[no-redef]
-        class completions:
-            @staticmethod
-            async def create(**kwargs: Any) -> Any:
-                # Return a minimal object mimicking OpenAI response
-                class Msg:
-                    def __init__(self) -> None:
-                        self.content = "Thinking..."
-                        self.tool_calls = [
-                            type(
-                                "ToolCall",
-                                (),
-                                {
-                                    "id": "call_1",
-                                    "function": type(
-                                        "Fn",
-                                        (),
-                                        {
-                                            "name": "computer",
-                                            "arguments": json.dumps(
-                                                {
-                                                    "action": "click",
-                                                    "element_description": "blue button",
-                                                }
-                                            ),
-                                        },
-                                    ),
-                                },
-                            )()
-                        ]
-
-                class Choice:
-                    def __init__(self) -> None:
-                        self.message = Msg()
-                        self.finish_reason = "tool_calls"
-
-                class Resp:
-                    def __init__(self) -> None:
-                        self.choices = [Choice()]
-
-                return Resp()
-
-
 class FakeMCPClient:
     def __init__(self) -> None:
         self.tools: list[types.Tool] = [
@@ -62,6 +18,7 @@ class FakeMCPClient:
             types.Tool(name="setup", description="internal functions", inputSchema={}),
         ]
         self.called: list[MCPToolCall] = []
+        self._initialized = True
 
     async def initialize(self, mcp_config: dict[str, dict[str, Any]] | None = None) -> None:
         return None
@@ -77,6 +34,10 @@ class FakeMCPClient:
     def mcp_config(self) -> dict[str, dict[str, Any]]:
         return {"local": {"command": "echo", "args": ["ok"]}}
 
+    @property
+    def is_connected(self) -> bool:
+        return self._initialized
+
     async def shutdown(self) -> None:
         return None
 
@@ -109,19 +70,20 @@ class DummyGroundedTool:
 
 @pytest.mark.asyncio
 async def test_call_tools_injects_screenshot_and_delegates(monkeypatch: pytest.MonkeyPatch) -> None:
-    # Agent with fake OpenAI client and fake MCP client
+    # Agent with fake OpenAI client
     grounder_cfg = GrounderConfig(api_base="http://example", model="qwen")
-    agent = GroundedOpenAIChatAgent(
+    fake_openai = AsyncOpenAI(api_key="test")
+    agent = GroundedOpenAIChatAgent.create(
         grounder_config=grounder_cfg,
-        openai_client=DummyOpenAI(),
-        model_name="gpt-4o-mini",
-        mcp_client=FakeMCPClient(),
+        openai_client=fake_openai,
+        model="gpt-4o-mini",
         initial_screenshot=False,
     )
 
     # Inject a dummy grounded tool to observe args without full initialization
     dummy_tool = DummyGroundedTool()
     agent.grounded_tool = dummy_tool  # type: ignore
+    agent._initialized = True  # Mark as initialized to skip context initialization
 
     # Seed conversation history with a user image
     png_b64 = (
@@ -153,3 +115,56 @@ async def test_call_tools_injects_screenshot_and_delegates(monkeypatch: pytest.M
     assert dummy_tool.last_args["element_description"] == "blue button"
     assert "screenshot_b64" in dummy_tool.last_args
     assert isinstance(dummy_tool.last_args["screenshot_b64"], str)
+
+
+@pytest.mark.asyncio
+async def test_get_response_with_reasoning() -> None:
+    """Test that reasoning content is extracted from the response."""
+    from unittest.mock import AsyncMock, MagicMock, patch
+
+    grounder_cfg = GrounderConfig(api_base="http://example", model="qwen")
+    fake_openai = AsyncOpenAI(api_key="test")
+
+    with patch("hud.settings.settings.telemetry_enabled", False):
+        agent = GroundedOpenAIChatAgent.create(
+            grounder_config=grounder_cfg,
+            openai_client=fake_openai,
+            model="gpt-4o-mini",
+            initial_screenshot=False,
+        )
+
+        mock_response = MagicMock()
+        mock_choice = MagicMock()
+        mock_message = MagicMock()
+
+        mock_message.content = "Here is my answer"
+        mock_message.reasoning_content = "Let me think step by step..."
+        mock_message.tool_calls = None
+
+        mock_choice.message = mock_message
+        mock_choice.finish_reason = "stop"
+
+        mock_response.choices = [mock_choice]
+
+        agent.oai.chat.completions.create = AsyncMock(return_value=mock_response)
+        agent._initialized = True  # Mark as initialized to skip context initialization
+
+        # Include an image so get_response doesn't try to take a screenshot via ctx
+        png_b64 = (
+            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR4nGMAAQAABQAB"
+            "J2n0mQAAAABJRU5ErkJggg=="
+        )
+        agent.conversation_history = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{png_b64}"}},
+                    {"type": "text", "text": "Hard question"},
+                ],
+            }
+        ]
+
+        response = await agent.get_response(agent.conversation_history)
+
+        assert response.content == "Here is my answer"
+        assert response.reasoning == "Let me think step by step..."