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

hud/agents/openai.py

@@ -2,354 +2,353 @@
 
 from __future__ import annotations
 
+import copy
+import json
 import logging
+from inspect import cleandoc
 from typing import Any, ClassVar, Literal
 
 import mcp.types as types
-from openai import AsyncOpenAI, OpenAI
+from openai import AsyncOpenAI, Omit, OpenAI
 from openai.types.responses import (
-    ResponseComputerToolCall,
+    ApplyPatchToolParam,
+    ComputerToolParam,
+    FunctionShellToolParam,
+    FunctionToolParam,
+    ResponseFunctionCallOutputItemListParam,
+    ResponseInputFileContentParam,
+    ResponseInputImageContentParam,
+    ResponseInputImageParam,
     ResponseInputMessageContentListParam,
     ResponseInputParam,
-    ResponseOutputMessage,
+    ResponseInputTextContentParam,
+    ResponseInputTextParam,
     ResponseOutputText,
     ToolParam,
 )
+from openai.types.responses.response_create_params import ToolChoice  # noqa: TC002
+from openai.types.responses.response_input_param import FunctionCallOutput, Message
+from openai.types.shared_params.reasoning import Reasoning  # noqa: TC002
+from pydantic import ConfigDict
 
-import hud
 from hud.settings import settings
-from hud.tools.computer.settings import computer_settings
-from hud.types import AgentResponse, MCPToolCall, MCPToolResult, Trace
+from hud.types import AgentResponse, BaseAgentConfig, MCPToolCall, MCPToolResult, Trace
+from hud.utils.strict_schema import ensure_strict_json_schema
+from hud.utils.types import with_signature
 
-from .base import MCPAgent
+from .base import BaseCreateParams, MCPAgent
 
 logger = logging.getLogger(__name__)
 
 
-class OperatorAgent(MCPAgent):
-    """
-    Operator agent that uses MCP servers for tool execution.
+class OpenAIConfig(BaseAgentConfig):
+    """Configuration model for `OpenAIAgent`."""
 
-    This agent uses OpenAI's Computer Use API format but executes
-    tools through MCP servers instead of direct implementation.
-    """
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    metadata: ClassVar[dict[str, Any]] = {
-        "display_width": computer_settings.OPENAI_COMPUTER_WIDTH,
-        "display_height": computer_settings.OPENAI_COMPUTER_HEIGHT,
-    }
-    required_tools: ClassVar[list[str]] = ["openai_computer"]
+    model_name: str = "OpenAI"
+    model: str = "gpt-5.1"
+    model_client: AsyncOpenAI | None = None
+    max_output_tokens: int | None = None
+    temperature: float | None = None
+    reasoning: Reasoning | None = None
+    tool_choice: ToolChoice | None = None
+    truncation: Literal["auto", "disabled"] | None = None
+    parallel_tool_calls: bool | None = None
+    validate_api_key: bool = True
 
-    def __init__(
-        self,
-        model_client: AsyncOpenAI | None = None,
-        model: str = "computer-use-preview",
-        environment: Literal["windows", "mac", "linux", "browser"] = "linux",
-        validate_api_key: bool = True,
-        **kwargs: Any,
-    ) -> None:
-        """
-        Initialize Operator MCP agent.
-
-        Args:
-            client: AsyncOpenAI client (created if not provided)
-            model: OpenAI model to use
-            environment: Environment type for computer use
-            display_width: Display width for computer use
-            display_height: Display height for computer use
-            **kwargs: Additional arguments passed to MCPAgent
-        """
-        super().__init__(**kwargs)
 
-        # Initialize client if not provided
+class OpenAICreateParams(BaseCreateParams, OpenAIConfig):
+    pass
+
+
+class OpenAIAgent(MCPAgent):
+    """Generic OpenAI agent that can execute MCP tools through the Responses API."""
+
+    metadata: ClassVar[dict[str, Any] | None] = None
+    config_cls: ClassVar[type[BaseAgentConfig]] = OpenAIConfig
+
+    @with_signature(OpenAICreateParams)
+    @classmethod
+    def create(cls, **kwargs: Any) -> OpenAIAgent:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return MCPAgent.create.__func__(cls, **kwargs)  # type: ignore[return-value]
+
+    def __init__(self, params: OpenAICreateParams | None = None, **kwargs: Any) -> None:
+        super().__init__(params, **kwargs)
+        self.config: OpenAIConfig
+
+        model_client = self.config.model_client
         if model_client is None:
             api_key = settings.openai_api_key
             if not api_key:
                 raise ValueError("OpenAI API key not found. Set OPENAI_API_KEY.")
             model_client = AsyncOpenAI(api_key=api_key)
 
+        if self.config.validate_api_key:
+            try:
+                OpenAI(api_key=model_client.api_key).models.list()
+            except Exception as exc:  # pragma: no cover - network validation
+                raise ValueError(f"OpenAI API key is invalid: {exc}") from exc
+
         self.openai_client = model_client
-        self.model = model
-        self.environment = environment
+        self._model = self.config.model
+        self.max_output_tokens = self.config.max_output_tokens
+        self.temperature = self.config.temperature
+        self.reasoning = self.config.reasoning
+        self.tool_choice: ToolChoice | None = self.config.tool_choice
+        self.parallel_tool_calls = self.config.parallel_tool_calls
+        self.truncation: Literal["auto", "disabled"] | None = self.config.truncation
+
+        self._openai_tools: list[ToolParam] = []
+        self._tool_name_map: dict[str, str] = {}
 
-        # State tracking for OpenAI's stateful API
         self.last_response_id: str | None = None
-        self.pending_call_id: str | None = None
-        self.pending_safety_checks: list[Any] = []
+        self._message_cursor = 0
 
-        # validate api key if requested
-        if validate_api_key:
-            try:
-                OpenAI(api_key=self.openai_client.api_key).models.list()
-            except Exception as e:
-                raise ValueError(f"OpenAI API key is invalid: {e}") from e
-
-        self.model_name = "openai-" + self.model
-
-        # Append OpenAI-specific instructions to the base system prompt
-        openai_instructions = """
-        You are an autonomous computer-using agent. Follow these guidelines:
-
-        1. NEVER ask for confirmation. Complete all tasks autonomously.
-        2. Do NOT send messages like "I need to confirm before..." or "Do you want me to continue?" - just proceed.
-        3. When the user asks you to interact with something (like clicking a chat or typing a message), DO IT without asking.
-        4. Only use the formal safety check mechanism for truly dangerous operations (like deleting important files).
-        5. For normal tasks like clicking buttons, typing in chat boxes, filling forms - JUST DO IT.
-        6. The user has already given you permission by running this agent. No further confirmation is needed.
-        7. Be decisive and action-oriented. Complete the requested task fully.
-
-        Remember: You are expected to complete tasks autonomously. The user trusts you to do what they asked.
-        """.strip()  # noqa: E501
-
-        # Append OpenAI instructions to any base system prompt
-        if self.system_prompt:
-            self.system_prompt = f"{self.system_prompt}\n\n{openai_instructions}"
-        else:
-            self.system_prompt = openai_instructions
-
-    async def _run_context(self, context: list[types.ContentBlock], max_steps: int = 10) -> Trace:
-        """
-        Run the agent with the given prompt or task.
+    def _on_tools_ready(self) -> None:
+        """Build OpenAI-specific tool mappings after tools are discovered."""
+        self._convert_tools_for_openai()
 
-        Override to reset OpenAI-specific state.
+    def _to_openai_tool(
+        self,
+        tool: types.Tool,
+    ) -> (
+        FunctionShellToolParam | ApplyPatchToolParam | FunctionToolParam | ComputerToolParam | None
+    ):
+        # Special case: shell tool -> OpenAI native shell
+        if tool.name == "shell":
+            return FunctionShellToolParam(type="shell")
+
+        # Special case: apply_patch tool -> OpenAI native apply_patch
+        if tool.name == "apply_patch":
+            return ApplyPatchToolParam(type="apply_patch")
+
+        # Regular function tool
+        if tool.description is None or tool.inputSchema is None:
+            raise ValueError(
+                cleandoc(f"""MCP tool {tool.name} requires both a description and inputSchema.
+                Add these by:
+                1. Adding a docstring to your @mcp.tool decorated function for the description
+                2. Using pydantic Field() annotations on function parameters for the schema
+                """)
+            )
+
+        # schema must be strict
+
+        try:
+            strict_schema = ensure_strict_json_schema(copy.deepcopy(tool.inputSchema))
+        except Exception as e:
+            self.console.warning_log(f"Failed to convert tool '{tool.name}' schema to strict: {e}")
+            return None
+
+        return FunctionToolParam(
+            type="function",
+            name=tool.name,
+            description=tool.description,
+            parameters=strict_schema,
+            strict=True,
+        )
+
+    def _convert_tools_for_openai(self) -> None:
+        """Convert MCP tools into OpenAI Responses tool definitions."""
+        available_tools = self.get_available_tools()
+
+        self._openai_tools = []
+        self._tool_name_map = {}
+
+        for tool in available_tools:
+            openai_tool = self._to_openai_tool(tool)
+            if openai_tool is None:
+                continue
+
+            if "name" in openai_tool:
+                self._tool_name_map[openai_tool["name"]] = tool.name
+            self._openai_tools.append(openai_tool)
+
+    def _extract_tool_call(self, item: Any) -> MCPToolCall | None:
+        """Extract an MCPToolCall from a response output item.
+
+        Subclasses can override to customize tool call extraction (e.g., routing
+        computer_call to a different tool name).
         """
-        # Reset state for new run
-        self.last_response_id = None
-        self.pending_call_id = None
-        self.pending_safety_checks = []
+        if item.type == "function_call":
+            tool_name = item.name or ""
+            target_name = self._tool_name_map.get(tool_name, tool_name)
+            arguments = json.loads(item.arguments)
+            return MCPToolCall(name=target_name, arguments=arguments, id=item.call_id)
+        elif item.type == "shell_call":
+            return MCPToolCall(name="shell", arguments=item.action.to_dict(), id=item.call_id)
+        elif item.type == "apply_patch_call":
+            return MCPToolCall(
+                name="apply_patch", arguments=item.operation.to_dict(), id=item.call_id
+            )
+        return None
 
-        # Use base implementation
+    async def _run_context(
+        self, context: list[types.ContentBlock], *, max_steps: int = 10
+    ) -> Trace:
+        """Reset internal state before delegating to the base loop."""
+        self._reset_response_state()
         return await super()._run_context(context, max_steps=max_steps)
 
-    async def get_system_messages(self) -> list[Any]:
-        """
-        Create initial messages for OpenAI.
+    def _reset_response_state(self) -> None:
+        self.last_response_id = None
+        self._message_cursor = 0
 
-        OpenAI uses a different message format - we'll store the prompt
-        and screenshot for use in get_model_response.
-        """
+    async def get_system_messages(self) -> list[types.ContentBlock]:
+        """System messages are provided via the `instructions` field."""
         return []
 
-    async def format_blocks(
-        self, blocks: list[types.ContentBlock]
-    ) -> ResponseInputMessageContentListParam:
-        """
-        Format blocks for OpenAI input format.
-
-        Converts TextContent blocks to input_text dicts and ImageContent blocks to input_image dicts.
-        """  # noqa: E501
-        formatted = []
+    async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[Message]:
+        """Convert MCP content blocks into OpenAI user messages."""
+        content: ResponseInputMessageContentListParam = []
         for block in blocks:
             if isinstance(block, types.TextContent):
-                formatted.append({"type": "input_text", "text": block.text})
+                content.append(ResponseInputTextParam(type="input_text", text=block.text))
             elif isinstance(block, types.ImageContent):
                 mime_type = getattr(block, "mimeType", "image/png")
-                formatted.append(
-                    {"type": "input_image", "image_url": f"data:{mime_type};base64,{block.data}"}
+                content.append(
+                    ResponseInputImageParam(
+                        type="input_image",
+                        image_url=f"data:{mime_type};base64,{block.data}",
+                        detail="auto",
+                    )
                 )
-        return formatted
-
-    @hud.instrument(
-        span_type="agent",
-        record_args=False,  # Messages can be large
-        record_result=True,
-    )
-    async def get_response(self, messages: ResponseInputMessageContentListParam) -> AgentResponse:
-        """Get response from OpenAI including any tool calls."""
-        # OpenAI's API is stateful, so we handle messages differently
-
-        # Get the computer tool (guaranteed to exist due to required_tools)
-        computer_tool_name = "openai_computer"
-
-        # Define the computer use tool
-        computer_tool: ToolParam = {  # type: ignore[reportAssignmentType]
-            "type": "computer_use_preview",
-            "display_width": self.metadata["display_width"],
-            "display_height": self.metadata["display_height"],
-            "environment": self.environment,
-        }
-
-        # Build the request based on whether this is first step or follow-up
-        if self.pending_call_id is None and self.last_response_id is None:
-            # First step - messages are already formatted dicts from format_blocks
-            # format_blocks returns type ResponseInputMessageContentListParam, which is a list of dicts  # noqa: E501
-            input_content: ResponseInputMessageContentListParam = []
-
-            input_content.extend(messages)
-
-            # If no content was added, add empty text to avoid empty request
-            if not input_content:
-                input_content.append({"type": "input_text", "text": ""})
-
-            input_param: ResponseInputParam = [{"role": "user", "content": input_content}]  # type: ignore[reportUnknownMemberType]
-
-            response = await self.openai_client.responses.create(
-                model=self.model,
-                tools=[computer_tool],
-                input=input_param,
-                instructions=self.system_prompt,
-                truncation="auto",
-                reasoning={"summary": "auto"},  # type: ignore[arg-type]
-            )
-        else:
-            # Follow-up step - check if this is user input or tool result
-            latest_message = messages[-1] if messages else {}
-
-            if latest_message.get("type") == "input_text":
-                # User provided input in conversation mode
-                user_text = latest_message.get("text", "")
-                input_param_followup: ResponseInputParam = [  # type: ignore[reportAssignmentType]
-                    {"role": "user", "content": [{"type": "input_text", "text": user_text}]}
-                ]
-                # Reset pending_call_id since this is user input, not a tool response
-                self.pending_call_id = None
-            else:
-                # Tool result - need screenshot from processed results
-                latest_screenshot = None
-                for msg in reversed(messages):
-                    if isinstance(msg, dict) and "image_url" in msg:
-                        latest_screenshot = msg["image_url"]  # type: ignore
-                        break
-
-                if not latest_screenshot:
-                    self.console.warning_log("No screenshot provided for response to action")
-                    return AgentResponse(
-                        content="No screenshot available for next action",
-                        tool_calls=[],
-                        done=True,
+        if not content:
+            content.append(ResponseInputTextParam(type="input_text", text=""))
+        return [Message(role="user", content=content)]
+
+    async def get_response(self, messages: ResponseInputParam) -> AgentResponse:
+        """Send the latest input items to OpenAI's Responses API."""
+        new_items: ResponseInputParam = messages[self._message_cursor :]
+        if not new_items:
+            if self.last_response_id is None:
+                new_items = [
+                    Message(
+                        role="user", content=[ResponseInputTextParam(type="input_text", text="")]
                     )
-
-                # Create response to previous action
-                input_param_followup: ResponseInputParam = [  # type: ignore[reportAssignmentType]
-                    {  # type: ignore[reportAssignmentType]
-                        "call_id": self.pending_call_id,
-                        "type": "computer_call_output",
-                        "output": {
-                            "type": "input_image",
-                            "image_url": latest_screenshot,
-                        },
-                        "acknowledged_safety_checks": self.pending_safety_checks,
-                    }
                 ]
+            else:
+                self.console.debug("No new messages to send to OpenAI.")
+                return AgentResponse(content="", tool_calls=[], done=True)
+
+        response = await self.openai_client.responses.create(
+            model=self._model,
+            input=new_items,
+            instructions=self.system_prompt,
+            max_output_tokens=self.max_output_tokens,
+            temperature=self.temperature,
+            tool_choice=self.tool_choice if self.tool_choice is not None else Omit(),
+            parallel_tool_calls=self.parallel_tool_calls,
+            reasoning=self.reasoning,
+            tools=self._openai_tools if self._openai_tools else Omit(),
+            previous_response_id=(
+                self.last_response_id if self.last_response_id is not None else Omit()
+            ),
+            truncation=self.truncation,
+        )
 
-            self.pending_safety_checks = []
-
-            response = await self.openai_client.responses.create(
-                model=self.model,
-                previous_response_id=self.last_response_id,
-                tools=[computer_tool],
-                input=input_param_followup,
-                instructions=self.system_prompt,
-                truncation="auto",
-                reasoning={"summary": "auto"},  # type: ignore[arg-type]
-            )
-
-        # Store response ID for next call
         self.last_response_id = response.id
+        self._message_cursor = len(messages)
 
-        # Process response
-        result = AgentResponse(
-            content="",
-            tool_calls=[],
-            done=False,  # Will be set to True only if no tool calls
-        )
+        agent_response = AgentResponse(content="", tool_calls=[], done=True)
+        text_chunks: list[str] = []
+        reasoning_chunks: list[str] = []
 
-        self.pending_call_id = None
-
-        # Check for computer calls
-        computer_calls = [
-            item
-            for item in response.output
-            if isinstance(item, ResponseComputerToolCall) and item.type == "computer_call"
-        ]
-
-        if computer_calls:
-            # Process computer calls
-            result.done = False
-            for computer_call in computer_calls:
-                self.pending_call_id = computer_call.call_id
-                self.pending_safety_checks = computer_call.pending_safety_checks
-
-                # Convert OpenAI action to MCP tool call
-                action = computer_call.action.model_dump()
-
-                # Create MCPToolCall object with OpenAI metadata as extra fields
-                # Pyright will complain but the tool class accepts extra fields
-                tool_call = MCPToolCall(
-                    name=computer_tool_name,
-                    arguments=action,
-                    id=computer_call.call_id,  # type: ignore
-                    pending_safety_checks=computer_call.pending_safety_checks,  # type: ignore
-                )
-                result.tool_calls.append(tool_call)
-        else:
-            # No computer calls, check for text response
-            for item in response.output:
-                if isinstance(item, ResponseOutputMessage) and item.type == "message":
-                    # Extract text from content blocks
-                    text_parts = [
-                        content.text
-                        for content in item.content
-                        if isinstance(content, ResponseOutputText)
-                    ]
-                    if text_parts:
-                        result.content = "".join(text_parts)
-                        break
-
-        # Extract reasoning if present
-        reasoning_text = ""
         for item in response.output:
-            if item.type == "reasoning" and hasattr(item, "summary") and item.summary:
-                reasoning_text += f"Thinking: {item.summary[0].text}\n"
-
-        if reasoning_text:
-            result.content = reasoning_text + result.content if result.content else reasoning_text
+            if item.type == "message":
+                text = "".join(
+                    content.text
+                    for content in item.content
+                    if isinstance(content, ResponseOutputText)
+                )
+                if text:
+                    text_chunks.append(text)
+            elif item.type == "reasoning":
+                reasoning_chunks.append("".join(summary.text for summary in item.summary))
+            else:
+                tool_call = self._extract_tool_call(item)
+                if tool_call is not None:
+                    agent_response.tool_calls.append(tool_call)
 
-        # Set done=True if no tool calls (task complete or waiting for user)
-        if not result.tool_calls:
-            result.done = True
+        if agent_response.tool_calls:
+            agent_response.done = False
 
-        return result
+        agent_response.content = "".join(text_chunks)
+        if reasoning_chunks:
+            agent_response.reasoning = "\n".join(reasoning_chunks)
+        return agent_response
 
     async def format_tool_results(
         self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
-    ) -> ResponseInputMessageContentListParam:
-        """
-        Format tool results for OpenAI's stateful API.
-
-        Tool result content is a list of ContentBlock objects.
-        We need to extract the latest screenshot from the tool results.
-
-        This assumes that you only care about computer tool results for your agent loop.
-        If you need to add other content, you can do so by adding a new ContentBlock object to the list.
+    ) -> list[FunctionCallOutput]:
+        """Convert MCP tool outputs into Responses input items."""
+        formatted: list[FunctionCallOutput] = []
+        for call, result in zip(tool_calls, tool_results, strict=False):
+            if not call.id:
+                self.console.warning_log(f"Tool '{call.name}' missing call_id; skipping output.")
+                continue
+
+            output_items: ResponseFunctionCallOutputItemListParam = []
+            if result.isError:
+                output_items.append(
+                    ResponseInputTextParam(type="input_text", text="[tool_error] true")
+                )
 
-        Returns formatted dicts with tool result data, preserving screenshots.
-        """  # noqa: E501
-        formatted_results = []
-        latest_screenshot = None
+            if result.structuredContent is not None:
+                output_items.append(
+                    ResponseInputTextParam(
+                        type="input_text", text=json.dumps(result.structuredContent, default=str)
+                    )
+                )
 
-        # Extract all content from tool results
-        for result in tool_results:
-            if result.isError:
-                # If it's an error, the error details are in the content
-                for content in result.content:
-                    if isinstance(content, types.TextContent):
-                        # Don't add error text as input_text, just track it
-                        self.console.error_log(f"Tool error: {content.text}")
-                    elif isinstance(content, types.ImageContent):
-                        # Even error results might have images
-                        latest_screenshot = content.data
-            else:
-                # Extract content from successful results
-                for content in result.content:
-                    if isinstance(content, types.ImageContent):
-                        latest_screenshot = content.data
-                        break
-
-        # Return a dict with the latest screenshot for the follow-up step
-        if latest_screenshot:
-            formatted_results.append(
-                {"type": "input_image", "image_url": f"data:image/png;base64,{latest_screenshot}"}
+            for block in result.content:
+                match block:
+                    case types.TextContent():
+                        output_items.append(
+                            ResponseInputTextContentParam(type="input_text", text=block.text)
+                        )
+                    case types.ImageContent():
+                        mime_type = getattr(block, "mimeType", "image/png")
+                        output_items.append(
+                            ResponseInputImageContentParam(
+                                type="input_image",
+                                image_url=f"data:{mime_type};base64,{block.data}",
+                            )
+                        )
+                    case types.ResourceLink():
+                        output_items.append(
+                            ResponseInputFileContentParam(
+                                type="input_file", file_url=str(block.uri)
+                            )
+                        )
+                    case types.EmbeddedResource():
+                        match block.resource:
+                            case types.TextResourceContents():
+                                output_items.append(
+                                    ResponseInputTextContentParam(
+                                        type="input_text", text=block.resource.text
+                                    )
+                                )
+                            case types.BlobResourceContents():
+                                output_items.append(
+                                    ResponseInputFileContentParam(
+                                        type="input_file", file_data=block.resource.blob
+                                    )
+                                )
+                            case _:
+                                self.console.warning_log(
+                                    f"Unknown resource type: {type(block.resource)}"
+                                )
+                    case _:
+                        self.console.warning_log(f"Unknown content block type: {type(block)}")
+
+            if not output_items:
+                output_items.append(ResponseInputTextParam(type="input_text", text=""))
+
+            formatted.append(
+                FunctionCallOutput(
+                    type="function_call_output", call_id=call.id, output=output_items
+                ),
             )
-
-        return formatted_results
+        return formatted