hud-python 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

hud/agents/tests/test_openai.py

@@ -1,237 +1,237 @@
-"""Tests for OpenAI MCP Agent implementation."""
-
-from __future__ import annotations
-
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-from mcp import types
-
-from hud.agents.openai import OperatorAgent
-from hud.types import MCPToolCall, MCPToolResult
-
-
-class TestOperatorAgent:
-    """Test OperatorAgent class."""
-
-    @pytest.fixture
-    def mock_mcp_client(self):
-        """Create a mock MCP client."""
-        mcp_client = MagicMock()
-        return mcp_client
-
-    @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
-
-    @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"
-        )
-
-        assert agent.model_name == "openai-gpt-4"
-        assert agent.model == "gpt-4"
-        assert agent.openai_client == mock_model_client
-
-    @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)
-
-        # 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) == 2
-        assert messages[0] == {"type": "input_text", "text": "Hello, GPT!"}
-        assert messages[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) == 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, model_client=mock_openai)
-
-        tool_calls = [
-            MCPToolCall(name="test_tool", arguments={}, id="call_123"),  # type: ignore
-            MCPToolCall(name="screenshot", arguments={}, id="call_456"),  # type: ignore
-        ]
-
-        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)
-
-        # 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"
-
-    @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, model_client=mock_openai)
-
-        tool_calls = [
-            MCPToolCall(name="failing_tool", arguments={}, id="call_error"),  # type: ignore
-        ]
-
-        tool_results = [
-            MCPToolResult(
-                content=[types.TextContent(type="text", text="Something went wrong")], isError=True
-            ),
-        ]
-
-        messages = await agent.format_tool_results(tool_calls, tool_results)
-
-        # Since the result has isError=True and no screenshot, returns empty list
-        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."""
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
-
-        # 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={})
-        ]
-
-        # Since OpenAI checks isinstance() on response types, we need to mock that
-        # For now, let's just test that we get the expected "No computer use tools available"
-        # when there are no matching tools
-        agent._available_tools = [
-            types.Tool(name="other_tool", description="Other tool", inputSchema={})
-        ]
-
-        messages = [{"prompt": "What's on the screen?", "screenshot": None}]
-        response = await agent.get_response(messages)
-
-        assert response.content == "No computer use tools available"
-        assert response.tool_calls == []
-        assert response.done is True
-
-    @pytest.mark.asyncio
-    async def test_get_model_response_text_only(self, mock_mcp_client, mock_openai):
-        """Test getting text-only response when no computer tools available."""
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
-
-        # Set up with no computer tools
-        agent._available_tools = []
-
-        messages = [{"prompt": "Hi", "screenshot": None}]
-        response = await agent.get_response(messages)
-
-        assert response.content == "No computer use tools available"
-        assert response.tool_calls == []
-        assert response.done is True
-
-    @pytest.mark.asyncio
-    async def test_run_with_tools(self, mock_mcp_client, mock_openai):
-        """Test running agent with tool usage."""
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
-
-        # Mock tool availability
-        agent._available_tools = [
-            types.Tool(name="search", description="Search tool", inputSchema={"type": "object"})
-        ]
-        # Base agent doesn't require server mapping for tool execution
-
-        # Mock initial response with tool use
-        initial_choice = MagicMock()
-        initial_choice.message = MagicMock(
-            content=None,
-            tool_calls=[
-                MagicMock(
-                    id="call_search",
-                    function=MagicMock(name="search", arguments='{"query": "OpenAI news"}'),
-                )
-            ],
-        )
-
-        initial_response = MagicMock()
-        initial_response.choices = [initial_choice]
-        initial_response.usage = MagicMock(prompt_tokens=10, completion_tokens=15, total_tokens=25)
-
-        # Mock follow-up response
-        final_choice = MagicMock()
-        final_choice.message = MagicMock(
-            content="Here are the latest OpenAI news...", tool_calls=None
-        )
-
-        final_response = MagicMock()
-        final_response.choices = [final_choice]
-        final_response.usage = MagicMock(prompt_tokens=20, completion_tokens=10, total_tokens=30)
-
-        mock_openai.chat.completions.create = AsyncMock(
-            side_effect=[initial_response, final_response]
-        )
-
-        # Mock tool execution
-        mock_mcp_client.call_tool = AsyncMock(
-            return_value=MCPToolResult(
-                content=[types.TextContent(type="text", text="Search results...")], isError=False
-            )
-        )
-
-        # Use a string prompt instead of a task
-        result = await agent.run("Search for OpenAI news")
-
-        # Since OpenAI integration currently returns "No computer use tools available"
-        # when the tool isn't a computer tool, we expect this
-        assert result.content == "No computer use tools available"
-        assert result.done is True
-
-    @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, model_client=mock_openai)
-
-        # Set up available tools
-        agent._available_tools = [
-            types.Tool(name="openai_computer", description="Computer tool", inputSchema={})
-        ]
-
-        # 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)
-
-        assert response.content == ""
-        assert response.tool_calls == []
+"""Tests for OpenAI MCP Agent implementation."""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from mcp import types
+
+from hud.agents.openai import OperatorAgent
+from hud.types import MCPToolCall, MCPToolResult
+
+
+class TestOperatorAgent:
+    """Test OperatorAgent class."""
+
+    @pytest.fixture
+    def mock_mcp_client(self):
+        """Create a mock MCP client."""
+        mcp_client = MagicMock()
+        return mcp_client
+
+    @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
+
+    @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"
+        )
+
+        assert agent.model_name == "openai-gpt-4"
+        assert agent.model == "gpt-4"
+        assert agent.openai_client == mock_model_client
+
+    @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)
+
+        # 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) == 2
+        assert messages[0] == {"type": "input_text", "text": "Hello, GPT!"}
+        assert messages[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) == 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, model_client=mock_openai)
+
+        tool_calls = [
+            MCPToolCall(name="test_tool", arguments={}, id="call_123"),  # type: ignore
+            MCPToolCall(name="screenshot", arguments={}, id="call_456"),  # type: ignore
+        ]
+
+        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)
+
+        # 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"
+
+    @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, model_client=mock_openai)
+
+        tool_calls = [
+            MCPToolCall(name="failing_tool", arguments={}, id="call_error"),  # type: ignore
+        ]
+
+        tool_results = [
+            MCPToolResult(
+                content=[types.TextContent(type="text", text="Something went wrong")], isError=True
+            ),
+        ]
+
+        messages = await agent.format_tool_results(tool_calls, tool_results)
+
+        # Since the result has isError=True and no screenshot, returns empty list
+        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."""
+        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
+
+        # 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={})
+        ]
+
+        # Since OpenAI checks isinstance() on response types, we need to mock that
+        # For now, let's just test that we get the expected "No computer use tools available"
+        # when there are no matching tools
+        agent._available_tools = [
+            types.Tool(name="other_tool", description="Other tool", inputSchema={})
+        ]
+
+        messages = [{"prompt": "What's on the screen?", "screenshot": None}]
+        response = await agent.get_response(messages)
+
+        assert response.content == "No computer use tools available"
+        assert response.tool_calls == []
+        assert response.done is True
+
+    @pytest.mark.asyncio
+    async def test_get_model_response_text_only(self, mock_mcp_client, mock_openai):
+        """Test getting text-only response when no computer tools available."""
+        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
+
+        # Set up with no computer tools
+        agent._available_tools = []
+
+        messages = [{"prompt": "Hi", "screenshot": None}]
+        response = await agent.get_response(messages)
+
+        assert response.content == "No computer use tools available"
+        assert response.tool_calls == []
+        assert response.done is True
+
+    @pytest.mark.asyncio
+    async def test_run_with_tools(self, mock_mcp_client, mock_openai):
+        """Test running agent with tool usage."""
+        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
+
+        # Mock tool availability
+        agent._available_tools = [
+            types.Tool(name="search", description="Search tool", inputSchema={"type": "object"})
+        ]
+        # Base agent doesn't require server mapping for tool execution
+
+        # Mock initial response with tool use
+        initial_choice = MagicMock()
+        initial_choice.message = MagicMock(
+            content=None,
+            tool_calls=[
+                MagicMock(
+                    id="call_search",
+                    function=MagicMock(name="search", arguments='{"query": "OpenAI news"}'),
+                )
+            ],
+        )
+
+        initial_response = MagicMock()
+        initial_response.choices = [initial_choice]
+        initial_response.usage = MagicMock(prompt_tokens=10, completion_tokens=15, total_tokens=25)
+
+        # Mock follow-up response
+        final_choice = MagicMock()
+        final_choice.message = MagicMock(
+            content="Here are the latest OpenAI news...", tool_calls=None
+        )
+
+        final_response = MagicMock()
+        final_response.choices = [final_choice]
+        final_response.usage = MagicMock(prompt_tokens=20, completion_tokens=10, total_tokens=30)
+
+        mock_openai.chat.completions.create = AsyncMock(
+            side_effect=[initial_response, final_response]
+        )
+
+        # Mock tool execution
+        mock_mcp_client.call_tool = AsyncMock(
+            return_value=MCPToolResult(
+                content=[types.TextContent(type="text", text="Search results...")], isError=False
+            )
+        )
+
+        # Use a string prompt instead of a task
+        result = await agent.run("Search for OpenAI news")
+
+        # Since OpenAI integration currently returns "No computer use tools available"
+        # when the tool isn't a computer tool, we expect this
+        assert result.content == "No computer use tools available"
+        assert result.done is True
+
+    @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, model_client=mock_openai)
+
+        # Set up available tools
+        agent._available_tools = [
+            types.Tool(name="openai_computer", description="Computer tool", inputSchema={})
+        ]
+
+        # 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)
+
+        assert response.content == ""
+        assert response.tool_calls == []