hud-python 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

hud/{mcp → agents}/tests/test_openai.py

@@ -1,238 +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 mcp.types import CallToolRequestParams as MCPToolCall
-
-from hud.mcp.openai import OpenAIMCPAgent
-
-
-class TestOpenAIMCPAgent:
-    """Test OpenAIMCPAgent class."""
-
-    @pytest.fixture
-    def mock_mcp_client(self):
-        """Create a mock MCP client."""
-        mcp_client = MagicMock()
-        mcp_client.get_all_active_sessions = MagicMock(return_value={})
-        mcp_client.get_tool_map = MagicMock(return_value={})
-        return mcp_client
-
-    @pytest.fixture
-    def mock_openai(self):
-        """Create a mock OpenAI client."""
-        with patch("hud.mcp.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 = OpenAIMCPAgent(
-            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_create_initial_messages(self, mock_mcp_client):
-        """Test creating initial messages."""
-        mock_model_client = MagicMock()
-        agent = OpenAIMCPAgent(mcp_client=mock_mcp_client, model_client=mock_model_client)
-
-        # Test with text only
-        messages = await agent.create_initial_messages("Hello, GPT!")
-        assert len(messages) == 1
-        assert messages[0]["prompt"] == "Hello, GPT!"
-        assert messages[0]["screenshot"] is None
-
-        # Test with screenshot
-        messages = await agent.create_initial_messages("Look at this", screenshot="base64data")
-        assert len(messages) == 1
-        assert messages[0]["prompt"] == "Look at this"
-        assert messages[0]["screenshot"] == "base64data"
-
-    @pytest.mark.asyncio
-    async def test_format_tool_results(self, mock_mcp_client, mock_openai):
-        """Test formatting tool results."""
-        agent = OpenAIMCPAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
-
-        tool_calls = [
-            MCPToolCall(name="test_tool", arguments={}, call_id="call_123"),  # type: ignore
-            MCPToolCall(name="screenshot", arguments={}, call_id="call_456"),  # type: ignore
-        ]
-
-        tool_results = [
-            types.CallToolResult(
-                content=[types.TextContent(type="text", text="Success")], isError=False
-            ),
-            types.CallToolResult(
-                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 just returns a simple dict with screenshot
-        assert len(messages) == 1
-        assert messages[0]["type"] == "tool_result"
-        assert (
-            messages[0]["screenshot"] == "base64data"
-        )  # Should extract screenshot from second result
-
-    @pytest.mark.asyncio
-    async def test_format_tool_results_with_error(self, mock_mcp_client, mock_openai):
-        """Test formatting tool results with errors."""
-        agent = OpenAIMCPAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
-
-        tool_calls = [
-            MCPToolCall(name="failing_tool", arguments={}, call_id="call_error"),  # type: ignore
-        ]
-
-        tool_results = [
-            types.CallToolResult(
-                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, no screenshot should be extracted
-        assert len(messages) == 1
-        assert messages[0]["type"] == "tool_result"
-        assert messages[0]["screenshot"] is None
-
-    @pytest.mark.asyncio
-    async def test_get_model_response(self, mock_mcp_client, mock_openai):
-        """Test getting model response from OpenAI API."""
-        agent = OpenAIMCPAgent(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_model_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 = OpenAIMCPAgent(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_model_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 = OpenAIMCPAgent(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"})
-        ]
-        agent._tool_map = {
-            "search": (
-                "server1",
-                types.Tool(
-                    name="search", description="Search tool", inputSchema={"type": "object"}
-                ),
-            )
-        }
-
-        # 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
-        agent.mcp_client.call_tool = AsyncMock(
-            return_value=types.CallToolResult(
-                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 = OpenAIMCPAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
-
-        # Set up available tools
-        agent._available_tools = [
-            types.Tool(name="computer_openai", 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_model_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 == []