hud-python 0.4.45__py3-none-any.whl → 0.5.13__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_resolver.py

@@ -0,0 +1,192 @@
+"""Tests for model resolution and create_agent."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from hud.agents import create_agent
+from hud.agents.resolver import resolve_cls
+
+
+class TestResolveCls:
+    """Tests for resolve_cls function."""
+
+    def test_resolves_known_agent_type(self) -> None:
+        """Known AgentType strings resolve to their class."""
+        from hud.agents.claude import ClaudeAgent
+
+        cls, gateway_info = resolve_cls("claude")
+        assert cls == ClaudeAgent
+        assert gateway_info is None
+
+    def test_resolves_openai(self) -> None:
+        """Resolves 'openai' to OpenAIAgent."""
+        from hud.agents import OpenAIAgent
+
+        cls, _gateway_info = resolve_cls("openai")
+        assert cls == OpenAIAgent
+
+    def test_resolves_gemini(self) -> None:
+        """Resolves 'gemini' to GeminiAgent."""
+        from hud.agents.gemini import GeminiAgent
+
+        cls, _gateway_info = resolve_cls("gemini")
+        assert cls == GeminiAgent
+
+    def test_unknown_model_without_gateway_raises(self) -> None:
+        """Unknown model with no gateway models raises ValueError."""
+        with (
+            patch("hud.agents.resolver._fetch_gateway_models", return_value=[]),
+            pytest.raises(ValueError, match="not found"),
+        ):
+            resolve_cls("unknown-model-xyz")
+
+    def test_resolves_gateway_model(self) -> None:
+        """Resolves model found in gateway."""
+        from hud.agents import OpenAIAgent
+
+        mock_models = [
+            {"id": "gpt-4o", "model": "gpt-4o", "provider": "openai"},
+        ]
+
+        with patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models):
+            cls, info = resolve_cls("gpt-4o")
+            assert cls == OpenAIAgent
+            assert info is not None
+            assert info["id"] == "gpt-4o"
+
+    def test_resolves_anthropic_provider_to_claude(self) -> None:
+        """Provider 'anthropic' maps to ClaudeAgent."""
+        from hud.agents.claude import ClaudeAgent
+
+        mock_models = [
+            {"id": "claude-sonnet", "model": "claude-3-sonnet", "provider": "anthropic"},
+        ]
+
+        with patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models):
+            cls, _info = resolve_cls("claude-sonnet")
+            assert cls == ClaudeAgent
+
+    def test_resolves_unknown_provider_to_openai_compatible(self) -> None:
+        """Unknown provider maps to OpenAIChatAgent."""
+        from hud.agents.openai_chat import OpenAIChatAgent
+
+        mock_models = [
+            {"id": "custom-model", "model": "custom", "provider": "custom-provider"},
+        ]
+
+        with patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models):
+            cls, _info = resolve_cls("custom-model")
+            assert cls == OpenAIChatAgent
+
+
+class TestCreateAgent:
+    """Tests for create_agent function - gateway-only."""
+
+    def test_creates_with_gateway_client(self) -> None:
+        """create_agent always uses gateway routing."""
+        from hud.agents import OpenAIAgent
+
+        mock_models = [
+            {"id": "gpt-4o", "model": "gpt-4o", "provider": "openai"},
+        ]
+
+        with (
+            patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models),
+            patch.object(OpenAIAgent, "create") as mock_create,
+            patch("hud.agents.gateway.build_gateway_client") as mock_build_client,
+        ):
+            mock_client = MagicMock()
+            mock_build_client.return_value = mock_client
+            mock_agent = MagicMock()
+            mock_create.return_value = mock_agent
+
+            agent = create_agent("gpt-4o")
+
+            # Should have set model and model_client
+            call_kwargs = mock_create.call_args.kwargs
+            assert call_kwargs["model"] == "gpt-4o"
+            assert "model_client" in call_kwargs
+            assert agent == mock_agent
+
+    def test_passes_kwargs_to_create(self) -> None:
+        """Extra kwargs are passed to agent.create()."""
+        from hud.agents import OpenAIAgent
+
+        mock_models = [
+            {"id": "gpt-4o", "model": "gpt-4o", "provider": "openai"},
+        ]
+
+        with (
+            patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models),
+            patch.object(OpenAIAgent, "create") as mock_create,
+            patch("hud.agents.gateway.build_gateway_client"),
+        ):
+            mock_create.return_value = MagicMock()
+
+            create_agent("gpt-4o", temperature=0.5, max_tokens=1000)
+
+            call_kwargs = mock_create.call_args.kwargs
+            assert call_kwargs["temperature"] == 0.5
+            assert call_kwargs["max_tokens"] == 1000
+
+    def test_known_agent_type_also_uses_gateway(self) -> None:
+        """Even 'claude' string uses gateway (it's a gateway shortcut)."""
+        from hud.agents.claude import ClaudeAgent
+
+        with (
+            patch.object(ClaudeAgent, "create") as mock_create,
+            patch("hud.agents.gateway.build_gateway_client") as mock_build_client,
+        ):
+            mock_client = MagicMock()
+            mock_build_client.return_value = mock_client
+            mock_create.return_value = MagicMock()
+
+            create_agent("claude")
+
+            # Should still build gateway client
+            mock_build_client.assert_called_once()
+            call_kwargs = mock_create.call_args.kwargs
+            assert "model_client" in call_kwargs
+
+
+class TestBuildGatewayClient:
+    """Tests for build_gateway_client function."""
+
+    def test_builds_anthropic_client(self) -> None:
+        """Builds AsyncAnthropic for anthropic provider."""
+        from hud.agents.gateway import build_gateway_client
+
+        with patch("hud.settings.settings") as mock_settings:
+            mock_settings.api_key = "test-key"
+            mock_settings.hud_gateway_url = "https://gateway.hud.ai"
+
+            with patch("anthropic.AsyncAnthropic") as mock_client_cls:
+                build_gateway_client("anthropic")
+                mock_client_cls.assert_called_once()
+
+    def test_builds_openai_client_for_openai(self) -> None:
+        """Builds AsyncOpenAI for openai provider."""
+        from hud.agents.gateway import build_gateway_client
+
+        with patch("hud.settings.settings") as mock_settings:
+            mock_settings.api_key = "test-key"
+            mock_settings.hud_gateway_url = "https://gateway.hud.ai"
+
+            with patch("openai.AsyncOpenAI") as mock_client_cls:
+                build_gateway_client("openai")
+                mock_client_cls.assert_called_once()
+
+    def test_builds_openai_client_for_unknown(self) -> None:
+        """Builds AsyncOpenAI for unknown providers (openai-compatible)."""
+        from hud.agents.gateway import build_gateway_client
+
+        with patch("hud.settings.settings") as mock_settings:
+            mock_settings.api_key = "test-key"
+            mock_settings.hud_gateway_url = "https://gateway.hud.ai"
+
+            with patch("openai.AsyncOpenAI") as mock_client_cls:
+                build_gateway_client("together")
+                mock_client_cls.assert_called_once()