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

hud/agents/tests/test_operator.py

@@ -0,0 +1,362 @@
+"""Tests for OperatorAgent implementation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from mcp import types
+from openai import AsyncOpenAI
+from openai.types.responses.response_computer_tool_call import PendingSafetyCheck
+
+from hud.agents.operator import OperatorAgent
+from hud.environment.router import ToolRouter
+from hud.eval.context import EvalContext
+from hud.types import MCPToolCall, MCPToolResult
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+
+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 TestOperatorAgent:
+    """Test OperatorAgent class."""
+
+    @pytest.fixture
+    def mock_openai(self) -> Generator[AsyncOpenAI, None, None]:
+        """Create a mock OpenAI client."""
+        client = AsyncOpenAI(api_key="test", base_url="http://localhost")
+        client.responses.create = AsyncMock()
+        with patch("hud.agents.openai.AsyncOpenAI", return_value=client):
+            yield client
+
+    @pytest.fixture
+    def mock_eval_context_computer(self) -> MockEvalContext:
+        """Create a mock EvalContext with computer tool."""
+        return MockEvalContext(
+            tools=[
+                types.Tool(
+                    name="openai_computer",
+                    description="OpenAI computer use tool",
+                    inputSchema={},
+                )
+            ]
+        )
+
+    @pytest.mark.asyncio
+    async def test_init(self, mock_openai: AsyncOpenAI) -> None:
+        """Test agent initialization."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            model="gpt-4",
+            validate_api_key=False,
+        )
+
+        assert agent.model_name == "Operator"
+        assert agent.config.model == "gpt-4"
+        assert agent.openai_client == mock_openai
+
+    @pytest.mark.asyncio
+    async def test_format_blocks(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting content blocks."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        # Test with text blocks
+        blocks: list[types.ContentBlock] = [
+            types.TextContent(type="text", text="Hello, GPT!"),
+            types.TextContent(type="text", text="Another message"),
+        ]
+
+        messages = await agent.format_blocks(blocks)
+        assert len(messages) == 1
+        msg = cast("dict[str, Any]", messages[0])
+        assert msg["role"] == "user"
+        content = cast("list[dict[str, Any]]", msg["content"])
+        assert len(content) == 2
+        assert content[0] == {"type": "input_text", "text": "Hello, GPT!"}
+        assert content[1] == {"type": "input_text", "text": "Another message"}
+
+        # Test with mixed content
+        blocks = [
+            types.TextContent(type="text", text="Text content"),
+            types.ImageContent(type="image", data="base64data", mimeType="image/png"),
+        ]
+
+        messages = await agent.format_blocks(blocks)
+        assert len(messages) == 1
+        msg = cast("dict[str, Any]", messages[0])
+        assert msg["role"] == "user"
+        content = cast("list[dict[str, Any]]", msg["content"])
+        assert len(content) == 2
+        assert content[0] == {"type": "input_text", "text": "Text content"}
+        assert content[1] == {
+            "type": "input_image",
+            "image_url": "data:image/png;base64,base64data",
+            "detail": "auto",
+        }
+
+    @pytest.mark.asyncio
+    async def test_format_tool_results(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting tool results."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        tool_calls = [
+            MCPToolCall(name="test_tool", arguments={}, id="call_123"),
+            MCPToolCall(name="screenshot", arguments={}, id="call_456"),
+        ]
+
+        tool_results = [
+            MCPToolResult(content=[types.TextContent(type="text", text="Success")], isError=False),
+            MCPToolResult(
+                content=[types.ImageContent(type="image", data="base64data", mimeType="image/png")],
+                isError=False,
+            ),
+        ]
+
+        messages = await agent.format_tool_results(tool_calls, tool_results)
+
+        # Should return both tool results as function_call_output
+        assert len(messages) == 2
+        # First result is text
+        msg0 = cast("dict[str, Any]", messages[0])
+        assert msg0["type"] == "function_call_output"
+        assert msg0["call_id"] == "call_123"
+        output0 = cast("list[dict[str, Any]]", msg0["output"])
+        assert output0[0]["type"] == "input_text"
+        assert output0[0]["text"] == "Success"
+        # Second result is image
+        msg1 = cast("dict[str, Any]", messages[1])
+        assert msg1["type"] == "function_call_output"
+        assert msg1["call_id"] == "call_456"
+        output1 = cast("list[dict[str, Any]]", msg1["output"])
+        assert output1[0]["type"] == "input_image"
+        assert output1[0]["image_url"] == "data:image/png;base64,base64data"
+
+    @pytest.mark.asyncio
+    async def test_format_tool_results_with_error(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting tool results with errors."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        tool_calls = [
+            MCPToolCall(name="failing_tool", arguments={}, id="call_error"),
+        ]
+
+        tool_results = [
+            MCPToolResult(
+                content=[types.TextContent(type="text", text="Something went wrong")], isError=True
+            ),
+        ]
+
+        messages = await agent.format_tool_results(tool_calls, tool_results)
+
+        # Error results are returned with error flag and content
+        assert len(messages) == 1
+        msg = cast("dict[str, Any]", messages[0])
+        assert msg["type"] == "function_call_output"
+        assert msg["call_id"] == "call_error"
+        output = cast("list[dict[str, Any]]", msg["output"])
+        assert output[0]["type"] == "input_text"
+        assert output[0]["text"] == "[tool_error] true"
+        assert output[1]["type"] == "input_text"
+        assert output[1]["text"] == "Something went wrong"
+
+    @pytest.mark.asyncio
+    async def test_get_model_response(
+        self, mock_openai: AsyncOpenAI, mock_eval_context_computer: MockEvalContext
+    ) -> None:
+        """Test getting model response from OpenAI API."""
+        with patch("hud.settings.settings.telemetry_enabled", False):
+            agent = OperatorAgent.create(
+                model_client=mock_openai,
+                validate_api_key=False,
+            )
+
+            # Initialize with context
+            agent.ctx = mock_eval_context_computer
+            await agent._initialize_from_ctx(mock_eval_context_computer)
+
+            # Mock OpenAI API response for a successful computer use response
+            mock_response = MagicMock()
+            mock_response.id = "response_123"
+            mock_response.state = "completed"
+            # Mock the output message structure
+            mock_output_text = MagicMock()
+            mock_output_text.type = "output_text"
+            mock_output_text.text = "I can see the screen content."
+
+            mock_output_message = MagicMock()
+            mock_output_message.type = "message"
+            mock_output_message.content = [mock_output_text]
+
+            mock_response.output = [mock_output_message]
+
+            mock_openai.responses.create = AsyncMock(return_value=mock_response)
+
+            messages = [{"prompt": "What's on the screen?", "screenshot": None}]
+            response = await agent.get_response(messages)  # type: ignore[arg-type]
+
+            assert response.done is True
+            assert response.tool_calls == []
+
+    @pytest.mark.asyncio
+    async def test_handle_empty_response(
+        self, mock_openai: AsyncOpenAI, mock_eval_context_computer: MockEvalContext
+    ) -> None:
+        """Test handling empty response from API."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        # Initialize with context
+        agent.ctx = mock_eval_context_computer
+        await agent._initialize_from_ctx(mock_eval_context_computer)
+
+        # Mock empty response
+        mock_response = MagicMock()
+        mock_response.id = "response_empty"
+        mock_response.state = "completed"
+        mock_response.output = []  # Empty output
+
+        mock_openai.responses.create = AsyncMock(return_value=mock_response)
+
+        messages = [{"prompt": "Hi", "screenshot": None}]
+        response = await agent.get_response(messages)  # type: ignore[arg-type]
+
+        assert response.content == ""
+        assert response.tool_calls == []
+
+    @pytest.mark.asyncio
+    async def test_pending_safety_checks_initialization(self, mock_openai: AsyncOpenAI) -> None:
+        """Test that OperatorAgent initializes pending_call_id and pending_safety_checks."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        # Verify initial state
+        assert agent.pending_call_id is None
+        assert agent.pending_safety_checks == []
+
+        # Set some state
+        agent.pending_call_id = "call_id"
+        agent.pending_safety_checks = [
+            PendingSafetyCheck(id="safety_check_id", code="value", message="message")
+        ]
+
+        # Verify state was set
+        assert agent.pending_call_id == "call_id"
+        assert len(agent.pending_safety_checks) == 1
+        assert agent.pending_safety_checks[0].id == "safety_check_id"
+
+    @pytest.mark.asyncio
+    async def test_extract_tool_call_computer(self, mock_openai: AsyncOpenAI) -> None:
+        """Test that _extract_tool_call routes computer_call to openai_computer."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        # Create a mock computer_call item
+        mock_item = MagicMock()
+        mock_item.type = "computer_call"
+        mock_item.call_id = "call_123"
+        mock_item.pending_safety_checks = [
+            PendingSafetyCheck(id="check_1", code="code", message="msg")
+        ]
+        mock_item.action.to_dict.return_value = {"type": "screenshot"}
+
+        tool_call = agent._extract_tool_call(mock_item)
+
+        # Should route to openai_computer tool
+        assert tool_call is not None
+        assert tool_call.name == "openai_computer"
+        assert tool_call.id == "call_123"
+        assert tool_call.arguments == {"type": "screenshot"}
+        # Should update pending_safety_checks
+        assert agent.pending_safety_checks == mock_item.pending_safety_checks
+
+    @pytest.mark.asyncio
+    async def test_extract_tool_call_delegates_to_super(self, mock_openai: AsyncOpenAI) -> None:
+        """Test that _extract_tool_call delegates non-computer calls to parent."""
+        agent = OperatorAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        # Set up tool name map
+        agent._tool_name_map = {"test_tool": "mcp_test_tool"}
+
+        # Create a mock function_call item
+        mock_item = MagicMock()
+        mock_item.type = "function_call"
+        mock_item.call_id = "call_456"
+        mock_item.name = "test_tool"
+        mock_item.arguments = '{"arg": "value"}'
+
+        tool_call = agent._extract_tool_call(mock_item)
+
+        # Should delegate to parent and map the tool name
+        assert tool_call is not None
+        assert tool_call.name == "mcp_test_tool"
+        assert tool_call.id == "call_456"
+        assert tool_call.arguments == {"arg": "value"}

hud/agents/tests/test_run_eval.py

@@ -0,0 +1,179 @@
+"""Tests for MCPAgent.run() with EvalContext."""
+
+from __future__ import annotations
+
+from typing import Any, ClassVar
+
+import pytest
+from mcp import types
+
+from hud.agents import MCPAgent
+from hud.agents.base import BaseCreateParams
+from hud.environment.router import ToolRouter
+from hud.eval.context import EvalContext
+from hud.types import AgentResponse, BaseAgentConfig, MCPToolCall, MCPToolResult
+
+
+class MockConfig(BaseAgentConfig):
+    model_name: str = "MockAgent"
+    model: str = "mock-model"
+
+
+class MockCreateParams(BaseCreateParams, MockConfig):
+    pass
+
+
+class MockMCPAgent(MCPAgent):
+    """Mock agent for testing run()."""
+
+    metadata: ClassVar[dict[str, Any] | None] = {}
+    config_cls: ClassVar[type[BaseAgentConfig]] = MockConfig
+
+    def __init__(self, **kwargs: Any) -> None:
+        params = MockCreateParams(**kwargs)
+        super().__init__(params)
+        self._response = AgentResponse(content="Test response", tool_calls=[], done=True)
+
+    def set_response(self, response: AgentResponse) -> None:
+        self._response = response
+
+    async def get_response(self, messages: list[dict[str, Any]]) -> AgentResponse:
+        return self._response
+
+    async def format_tool_results(
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
+    ) -> list[dict[str, Any]]:
+        return [{"role": "tool", "content": str(r)} for r in tool_results]
+
+    async def get_system_messages(self) -> list[Any]:
+        return []
+
+    async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[Any]:
+        return [{"type": "text", "text": getattr(b, "text")} for b in blocks if hasattr(b, "text")]
+
+
+class MockEvalContext(EvalContext):
+    """Mock EvalContext for testing - inherits from real EvalContext."""
+
+    def __init__(self, prompt: str = "Test prompt", tools: list[types.Tool] | None = None) -> None:
+        # Core attributes
+        self.prompt = prompt
+        self._tools = tools or [types.Tool(name="test_tool", description="Test", inputSchema={})]
+        self._submitted: str | None = None
+        self.reward: float | None = None
+        self._initialized = True
+
+        # 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 True
+
+    async def list_tools(self) -> list[types.Tool]:
+        return self._tools
+
+    async def call_tool(self, call: Any, /, **kwargs: Any) -> MCPToolResult:
+        # Handle tuple format (name, args)
+        if isinstance(call, tuple):
+            name = call[0]
+        elif hasattr(call, "name"):
+            name = call.name
+        else:
+            name = str(call)
+        return MCPToolResult(
+            content=[types.TextContent(type="text", text=f"Result from {name}")],
+            isError=False,
+        )
+
+    async def submit(self, answer: str) -> None:
+        self._submitted = answer
+
+
+class TestRun:
+    """Tests for MCPAgent.run() with EvalContext."""
+
+    @pytest.mark.asyncio
+    async def test_run_basic(self) -> None:
+        """Test basic run() flow."""
+        ctx = MockEvalContext(prompt="Do the task")
+        agent = MockMCPAgent()
+
+        result = await agent.run(ctx)
+
+        assert result.done
+        assert result.content == "Test response"
+        assert ctx._submitted == "Test response"
+
+    @pytest.mark.asyncio
+    async def test_run_no_prompt_raises(self) -> None:
+        """Test run() raises when prompt is not set."""
+        ctx = MockEvalContext(prompt="")
+        agent = MockMCPAgent()
+
+        with pytest.raises(ValueError, match="prompt is not set"):
+            await agent.run(ctx)
+
+    @pytest.mark.asyncio
+    async def test_run_wrong_type_raises(self) -> None:
+        """Test run() raises TypeError for non-EvalContext."""
+        agent = MockMCPAgent()
+
+        with pytest.raises(TypeError, match="must be EvalContext"):
+            await agent.run("not an eval context")  # type: ignore[arg-type]
+
+    @pytest.mark.asyncio
+    async def test_run_clears_ctx(self) -> None:
+        """Test run() clears ctx after completion."""
+        ctx = MockEvalContext(prompt="Do the task")
+        agent = MockMCPAgent()
+
+        await agent.run(ctx)
+        assert agent.ctx is None
+
+    @pytest.mark.asyncio
+    async def test_run_no_submit_on_empty_content(self) -> None:
+        """Test run() doesn't submit when content is empty."""
+        ctx = MockEvalContext(prompt="Do the task")
+        agent = MockMCPAgent()
+        agent.set_response(AgentResponse(content="", tool_calls=[], done=True))
+
+        await agent.run(ctx)
+        assert ctx._submitted is None
+
+    @pytest.mark.asyncio
+    async def test_run_initializes_tools(self) -> None:
+        """Test run() initializes tools from context."""
+        ctx = MockEvalContext(
+            prompt="Do the task",
+            tools=[
+                types.Tool(name="tool1", description="Tool 1", inputSchema={}),
+                types.Tool(name="tool2", description="Tool 2", inputSchema={}),
+            ],
+        )
+        agent = MockMCPAgent()
+
+        await agent.run(ctx)
+
+        assert agent._initialized
+        # After cleanup, ctx is None but tools were discovered