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

hud/agents/tests/test_openai.py

@@ -2,212 +2,448 @@
 
 from __future__ import annotations
 
-from unittest.mock import AsyncMock, MagicMock, patch
+from typing import TYPE_CHECKING, Any, cast
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from mcp import types
-
-from hud.agents.openai import OperatorAgent
+from openai import AsyncOpenAI
+from openai.types.responses import (
+    ResponseFunctionToolCall,
+    ResponseOutputMessage,
+    ResponseOutputText,
+    ResponseReasoningItem,
+)
+from openai.types.responses.response_reasoning_item import Summary
+
+from hud.agents.openai import OpenAIAgent
+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,
+        )
 
-class TestOperatorAgent:
-    """Test OperatorAgent class."""
+    async def submit(self, answer: str) -> None:
+        self._submitted = answer
 
-    @pytest.fixture
-    def mock_mcp_client(self):
-        """Create a mock MCP client."""
-        mcp_client = AsyncMock()
-        # Set up the mcp_config attribute as a regular dict, not a coroutine
-        mcp_client.mcp_config = {"test_server": {"url": "http://test"}}
-        # Mock list_tools to return the required openai_computer tool
-        mcp_client.list_tools = AsyncMock(
-            return_value=[
-                types.Tool(
-                    name="openai_computer", description="OpenAI computer use tool", inputSchema={}
-                )
-            ]
-        )
-        mcp_client.initialize = AsyncMock()
-        return mcp_client
+
+class TestOpenAIAgent:
+    """Test OpenAIAgent class."""
 
     @pytest.fixture
-    def mock_openai(self):
-        """Create a mock OpenAI client."""
-        with patch("hud.agents.openai.AsyncOpenAI") as mock:
-            client = AsyncMock()
-            mock.return_value = client
-            yield client
+    def mock_openai(self) -> Generator[AsyncOpenAI, None, None]:  # type: ignore[misc]
+        """Create a stub OpenAI client."""
+        with patch("hud.agents.openai.AsyncOpenAI") as mock_class:
+            client = AsyncOpenAI(api_key="test", base_url="http://localhost")
+            client.chat.completions.create = AsyncMock()
+            client.responses.create = AsyncMock()
+            mock_class.return_value = client
+            yield client  # type: ignore[misc]
 
     @pytest.mark.asyncio
-    async def test_init(self, mock_mcp_client):
-        """Test agent initialization."""
-        mock_model_client = MagicMock()
-        agent = OperatorAgent(
-            mcp_client=mock_mcp_client,
-            model_client=mock_model_client,
-            model="gpt-4",
-            validate_api_key=False,  # Skip validation in tests
+    async def test_init_with_client(self, mock_openai: AsyncOpenAI) -> None:
+        """Test agent initialization with provided client."""
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            model="gpt-4o",
+            validate_api_key=False,
         )
 
-        assert agent.model_name == "openai-gpt-4"
-        assert agent.model == "gpt-4"
-        assert agent.openai_client == mock_model_client
+        assert agent.model_name == "OpenAI"
+        assert agent.config.model == "gpt-4o"
+        assert agent.model == "gpt-4o"
+        assert agent.openai_client == mock_openai
+        assert agent.max_output_tokens is None
+        assert agent.temperature is None
+
+    @pytest.mark.asyncio
+    async def test_init_with_parameters(self, mock_openai: AsyncOpenAI) -> None:
+        """Test agent initialization with various parameters."""
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            model="gpt-4o",
+            max_output_tokens=2048,
+            temperature=0.7,
+            reasoning={"effort": "high"},
+            tool_choice="auto",
+            parallel_tool_calls=True,
+            validate_api_key=False,
+        )
+
+        assert agent.max_output_tokens == 2048
+        assert agent.temperature == 0.7
+        assert agent.reasoning == {"effort": "high"}
+        assert agent.tool_choice == "auto"
+        assert agent.parallel_tool_calls is True
+
+    @pytest.mark.asyncio
+    async def test_init_without_client_no_api_key(self) -> None:
+        """Test agent initialization fails without API key."""
+        with patch("hud.agents.openai.settings") as mock_settings:
+            mock_settings.openai_api_key = None
+            with pytest.raises(ValueError, match="OpenAI API key not found"):
+                OpenAIAgent.create()
 
     @pytest.mark.asyncio
-    async def test_format_blocks(self, mock_mcp_client):
-        """Test formatting content blocks."""
-        mock_model_client = MagicMock()
-        agent = OperatorAgent(
-            mcp_client=mock_mcp_client,
-            model_client=mock_model_client,
-            validate_api_key=False,  # Skip validation in tests
+    async def test_format_blocks_text_only(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting text content blocks."""
+        agent = OpenAIAgent.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"),
+            types.TextContent(type="text", text="Hello, world!"),
+            types.TextContent(type="text", text="How are you?"),
         ]
 
         messages = await agent.format_blocks(blocks)
-        assert len(messages) == 2
-        assert messages[0] == {"type": "input_text", "text": "Hello, GPT!"}
-        assert messages[1] == {"type": "input_text", "text": "Another message"}
+        assert len(messages) == 1
+        assert messages[0]["role"] == "user"
+        assert len(messages[0]["content"]) == 2
+        assert messages[0]["content"][0]["type"] == "input_text"
+        assert messages[0]["content"][0]["text"] == "Hello, world!"
 
-        # Test with mixed content
-        blocks = [
-            types.TextContent(type="text", text="Text content"),
+    @pytest.mark.asyncio
+    async def test_format_blocks_with_image(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting image content blocks."""
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+
+        blocks: list[types.ContentBlock] = [
+            types.TextContent(type="text", text="Look at this:"),
             types.ImageContent(type="image", data="base64data", mimeType="image/png"),
         ]
 
         messages = await agent.format_blocks(blocks)
-        assert len(messages) == 2
-        assert messages[0] == {"type": "input_text", "text": "Text content"}
-        assert messages[1] == {
-            "type": "input_image",
-            "image_url": "data:image/png;base64,base64data",
-        }
-
-    @pytest.mark.asyncio
-    async def test_format_tool_results(self, mock_mcp_client, mock_openai):
-        """Test formatting tool results."""
-        agent = OperatorAgent(
-            mcp_client=mock_mcp_client,
+        assert len(messages) == 1
+        assert len(messages[0]["content"]) == 2
+        assert messages[0]["content"][1]["type"] == "input_image"
+        assert messages[0]["content"][1]["image_url"] == "data:image/png;base64,base64data"  # type: ignore[typeddict-item]
+
+    @pytest.mark.asyncio
+    async def test_format_blocks_empty(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting empty content blocks."""
+        agent = OpenAIAgent.create(
             model_client=mock_openai,
-            validate_api_key=False,  # Skip validation in tests
+            validate_api_key=False,
         )
 
-        tool_calls = [
-            MCPToolCall(name="test_tool", arguments={}, id="call_123"),  # type: ignore
-            MCPToolCall(name="screenshot", arguments={}, id="call_456"),  # type: ignore
-        ]
+        messages = await agent.format_blocks([])
+        assert len(messages) == 1
+        # Empty blocks produce a single empty text item
+        assert len(messages[0]["content"]) == 1
+        assert messages[0]["content"][0]["type"] == "input_text"
+        assert messages[0]["content"][0]["text"] == ""
+
+    @pytest.mark.asyncio
+    async def test_format_tool_results_text(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting tool results with text content."""
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
 
+        tool_calls = [MCPToolCall(id="call_123", name="test_tool", arguments={})]
         tool_results = [
-            MCPToolResult(content=[types.TextContent(type="text", text="Success")], isError=False),
             MCPToolResult(
-                content=[types.ImageContent(type="image", data="base64data", mimeType="image/png")],
+                content=[types.TextContent(type="text", text="Tool output")],
                 isError=False,
-            ),
+            )
         ]
 
         messages = await agent.format_tool_results(tool_calls, tool_results)
-
-        # OpenAI's format_tool_results returns input_image with screenshot
         assert len(messages) == 1
-        assert messages[0]["type"] == "input_image"
-        assert "image_url" in messages[0]
-        assert messages[0]["image_url"] == "data:image/png;base64,base64data"
+        assert messages[0]["type"] == "function_call_output"
+        assert messages[0]["call_id"] == "call_123"
+        # Output is a list of content items
+        assert len(messages[0]["output"]) == 1
+        assert messages[0]["output"][0]["text"] == "Tool output"  # type: ignore[index]
 
     @pytest.mark.asyncio
-    async def test_format_tool_results_with_error(self, mock_mcp_client, mock_openai):
-        """Test formatting tool results with errors."""
-        agent = OperatorAgent(
-            mcp_client=mock_mcp_client,
+    async def test_format_tool_results_with_error(self, mock_openai: AsyncOpenAI) -> None:
+        """Test formatting tool results with error."""
+        agent = OpenAIAgent.create(
             model_client=mock_openai,
-            validate_api_key=False,  # Skip validation in tests
+            validate_api_key=False,
         )
 
-        tool_calls = [
-            MCPToolCall(name="failing_tool", arguments={}, id="call_error"),  # type: ignore
-        ]
-
+        tool_calls = [MCPToolCall(id="call_123", name="test_tool", arguments={})]
         tool_results = [
             MCPToolResult(
-                content=[types.TextContent(type="text", text="Something went wrong")], isError=True
-            ),
+                content=[types.TextContent(type="text", text="Error message")],
+                isError=True,
+            )
         ]
 
         messages = await agent.format_tool_results(tool_calls, tool_results)
+        assert len(messages) == 1
+        # Output is a list; first item is error indicator, second is the message
+        msg = cast("dict[str, Any]", messages[0])
+        output = cast("list[dict[str, Any]]", msg["output"])
+        assert any(item.get("text") == "[tool_error] true" for item in output)
+        assert any(item.get("text") == "Error message" for item in output)
+
+    @pytest.mark.asyncio
+    async def test_get_system_messages(self, mock_openai: AsyncOpenAI) -> None:
+        """Test getting system messages - OpenAI uses instructions field instead."""
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            system_prompt="You are a helpful assistant.",
+            validate_api_key=False,
+        )
 
-        # Since the result has isError=True and no screenshot, returns empty list
+        # OpenAI agent returns empty list - system prompt is passed via instructions
+        messages = await agent.get_system_messages()
         assert len(messages) == 0
 
     @pytest.mark.asyncio
-    async def test_get_model_response(self, mock_mcp_client, mock_openai):
-        """Test getting model response from OpenAI API."""
-        # Disable telemetry for this test to avoid backend configuration issues
-        with patch("hud.settings.settings.telemetry_enabled", False):
-            agent = OperatorAgent(
-                mcp_client=mock_mcp_client,
-                model_client=mock_openai,
-                validate_api_key=False,  # Skip validation in tests
+    async def test_convert_tools_for_openai(self, mock_openai: AsyncOpenAI) -> None:
+        """Test converting MCP tools to OpenAI format."""
+        tools = [
+            types.Tool(
+                name="my_tool",
+                description="A test tool",
+                inputSchema={"type": "object", "properties": {"x": {"type": "string"}}},
             )
+        ]
+        ctx = MockEvalContext(tools=tools)
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
 
-            # Set up available tools so agent doesn't return "No computer use tools available"
-            agent._available_tools = [
-                types.Tool(name="computer_openai", description="Computer tool", inputSchema={})
-            ]
+        # Initialize with context to trigger tool conversion
+        agent.ctx = ctx
+        await agent._initialize_from_ctx(ctx)
 
-            # 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."
+        # Check that tools were converted
+        assert len(agent._openai_tools) >= 1
+        # Find our tool
+        tool = next((t for t in agent._openai_tools if t.get("name") == "my_tool"), None)
+        assert tool is not None
+        assert tool["type"] == "function"
 
-            mock_output_message = MagicMock()
-            mock_output_message.type = "message"
-            mock_output_message.content = [mock_output_text]
+    @pytest.mark.asyncio
+    async def test_convert_tools_raises_on_incomplete(self, mock_openai: AsyncOpenAI) -> None:
+        """Test that tools without description raise error."""
+        tools = [
+            types.Tool(
+                name="incomplete_tool",
+                description=None,  # Missing description
+                inputSchema={"type": "object"},
+            )
+        ]
+        ctx = MockEvalContext(tools=tools)
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
 
-            mock_response.output = [mock_output_message]
+        agent.ctx = ctx
+        with pytest.raises(ValueError, match="requires both a description"):
+            await agent._initialize_from_ctx(ctx)
 
-            mock_openai.responses.create = AsyncMock(return_value=mock_response)
+    @pytest.mark.asyncio
+    async def test_get_response_with_text(self, mock_openai: AsyncOpenAI) -> None:
+        """Test getting response with text output."""
+        # Setup mock response
+        mock_response = AsyncMock()
+        mock_response.output = [
+            ResponseOutputMessage(
+                id="msg_123",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[ResponseOutputText(type="output_text", text="Hello!", annotations=[])],
+            )
+        ]
+        mock_openai.responses.create = AsyncMock(return_value=mock_response)
 
-            messages = [{"prompt": "What's on the screen?", "screenshot": None}]
-            response = await agent.get_response(messages)
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+        # Set empty tools to avoid needing initialization
+        agent._openai_tools = []
+        agent._initialized = True
 
-            # The test should verify that the response is processed correctly
-            # Since the isinstance checks will fail, content will be empty, but done should be True
-            assert response.done is True
-            assert response.tool_calls == []
+        response = await agent.get_response([])
+        assert response.content == "Hello!"
+        assert response.done is True
+        assert len(response.tool_calls) == 0
 
     @pytest.mark.asyncio
-    async def test_handle_empty_response(self, mock_mcp_client, mock_openai):
-        """Test handling empty response from API."""
-        agent = OperatorAgent(
-            mcp_client=mock_mcp_client,
+    async def test_get_response_with_tool_call(self, mock_openai: AsyncOpenAI) -> None:
+        """Test getting response with tool call."""
+        mock_response = AsyncMock()
+        # Tool calls come as separate output items, not inside message content
+        mock_response.output = [
+            ResponseFunctionToolCall(
+                id="call_123",
+                type="function_call",
+                call_id="call_123",
+                name="my_tool",
+                arguments='{"x": "value"}',
+            )
+        ]
+        mock_openai.responses.create = AsyncMock(return_value=mock_response)
+
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
+        agent._openai_tools = []
+        agent._tool_name_map = {"my_tool": "my_tool"}
+        agent._initialized = True
+
+        response = await agent.get_response([])
+        assert response.done is False
+        assert len(response.tool_calls) == 1
+        assert response.tool_calls[0].name == "my_tool"
+        assert response.tool_calls[0].arguments == {"x": "value"}
+
+    @pytest.mark.asyncio
+    async def test_get_response_with_reasoning(self, mock_openai: AsyncOpenAI) -> None:
+        """Test getting response with reasoning."""
+        mock_response = AsyncMock()
+        mock_response.output = [
+            ResponseReasoningItem(
+                id="reason_123",
+                type="reasoning",
+                summary=[Summary(type="summary_text", text="Thinking about it...")],
+            ),
+            ResponseOutputMessage(
+                id="msg_123",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[ResponseOutputText(type="output_text", text="Answer!", annotations=[])],
+            ),
+        ]
+        mock_openai.responses.create = AsyncMock(return_value=mock_response)
+
+        agent = OpenAIAgent.create(
             model_client=mock_openai,
-            validate_api_key=False,  # Skip validation in tests
+            validate_api_key=False,
         )
+        agent._openai_tools = []
+        agent._initialized = True
+
+        response = await agent.get_response([])
+        # Reasoning is stored separately from content
+        assert response.reasoning == "Thinking about it..."
+        assert response.content == "Answer!"
+
 
-        # Set up available tools
-        agent._available_tools = [
-            types.Tool(name="openai_computer", description="Computer tool", inputSchema={})
+class TestOpenAIToolConversion:
+    """Tests for tool conversion to OpenAI format."""
+
+    @pytest.fixture
+    def mock_openai(self) -> Generator[AsyncOpenAI, None, None]:  # type: ignore[misc]
+        """Create a stub OpenAI client."""
+        with patch("hud.agents.openai.AsyncOpenAI") as mock_class:
+            client = AsyncOpenAI(api_key="test", base_url="http://localhost")
+            client.responses.create = AsyncMock()
+            mock_class.return_value = client
+            yield client  # type: ignore[misc]
+
+    @pytest.mark.asyncio
+    async def test_shell_tool_conversion(self, mock_openai: AsyncOpenAI) -> None:
+        """Test that shell tool is converted to native format."""
+        tools = [
+            types.Tool(
+                name="shell",
+                description="Execute shell commands",
+                inputSchema={"type": "object"},
+            )
         ]
+        ctx = MockEvalContext(tools=tools)
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
 
-        # Mock empty response
-        mock_response = MagicMock()
-        mock_response.id = "response_empty"
-        mock_response.state = "completed"
-        mock_response.output = []  # Empty output
+        agent.ctx = ctx
+        await agent._initialize_from_ctx(ctx)
 
-        mock_openai.responses.create = AsyncMock(return_value=mock_response)
+        # Check for native shell tool
+        shell_tool = next((t for t in agent._openai_tools if t.get("type") == "shell"), None)
+        assert shell_tool is not None
+
+    @pytest.mark.asyncio
+    async def test_computer_tool_conversion(self, mock_openai: AsyncOpenAI) -> None:
+        """Test that computer tool is converted to function format."""
+        tools = [
+            types.Tool(
+                name="computer",
+                description="Control computer",
+                inputSchema={"type": "object"},
+            )
+        ]
+        ctx = MockEvalContext(tools=tools)
+        agent = OpenAIAgent.create(
+            model_client=mock_openai,
+            validate_api_key=False,
+        )
 
-        messages = [{"prompt": "Hi", "screenshot": None}]
-        response = await agent.get_response(messages)
+        agent.ctx = ctx
+        await agent._initialize_from_ctx(ctx)
 
-        assert response.content == ""
-        assert response.tool_calls == []
+        # Computer tool is converted to a regular function tool
+        computer_tool = next(
+            (t for t in agent._openai_tools if t.get("name") == "computer"),
+            None,
+        )
+        assert computer_tool is not None
+        assert computer_tool.get("type") == "function"