cua-agent 0.4.34__py3-none-any.whl → 0.4.36__py3-none-any.whl

agent/integrations/hud/__init__.py

@@ -8,21 +8,22 @@ Exports:
 - run_full_dataset(dataset, ...)
 - MCPComputerAgent
 """
+
 import time
 from typing import Any, Optional
 
 from agent.computers import is_agent_computer
-from datasets import load_dataset, Dataset
-from hud.datasets import Task, run_dataset
+from datasets import Dataset, load_dataset
 from hud import trace
+from hud.datasets import Task, run_dataset
 
 from .agent import MCPComputerAgent
 
-
 # ---------------------------------------------------------------------------
 # Single-task runner
 # ---------------------------------------------------------------------------
 
+
 async def run_single_task(
     dataset: str | Dataset | list[dict[str, Any]],
     *,
@@ -47,24 +48,20 @@ async def run_single_task(
 
     # Load dataset and pick a sample
     if isinstance(dataset, str):
-        dataset = load_dataset(dataset, split="train") # type: ignore[arg-type]
+        dataset = load_dataset(dataset, split="train")  # type: ignore[arg-type]
     elif isinstance(dataset, list):
         dataset = dataset
     else:
         dataset = dataset["train"]
-    
+
     sample_task = dataset[task_id]  # type: ignore[index]
     task_prompt = sample_task.get("prompt", f"Task {sample_task.get('id', 0)}")  # type: ignore[attr-defined]
 
     # Filter any existing Computer tools
     # The eval framework will add its own Computer tool per task
     if tools:
-        tools = [
-            tool 
-            for tool in tools 
-            if not is_agent_computer(tool)
-        ]
-    
+        tools = [tool for tool in tools if not is_agent_computer(tool)]
+
     with trace(name=task_prompt):
         task = Task(**sample_task)  # type: ignore[arg-type]
 
@@ -87,13 +84,14 @@ async def run_single_task(
         )
         print(f"Running: {task_prompt}")
         result = await agent.run(task, max_steps=10)
-        print(f"✅ Reward: {getattr(result, 'reward')}")
+        print(f"✅ Reward: {result.reward}")
 
 
 # ---------------------------------------------------------------------------
 # Full-dataset runner
 # ---------------------------------------------------------------------------
 
+
 async def run_full_dataset(
     dataset: str | Dataset | list[dict[str, Any]],
     *,
@@ -121,9 +119,9 @@ async def run_full_dataset(
 
     # Run with our MCP-based agent class.
     if isinstance(dataset, str):
-        dataset_name = dataset.split('/')[-1]
+        dataset_name = dataset.split("/")[-1]
         job_name = job_name or f"Evaluation {dataset_name}"
-        dataset = load_dataset(dataset, split=split) # type: ignore[arg-type]
+        dataset = load_dataset(dataset, split=split)  # type: ignore[arg-type]
     else:
         dataset_name = "custom"
         job_name = job_name or f"Evaluation {time.strftime('%H:%M %Y-%m-%d')}"
@@ -131,12 +129,8 @@ async def run_full_dataset(
     # Filter any existing Computer tools
     # The eval framework will add its own Computer tool per task
     if tools:
-        tools = [
-            tool 
-            for tool in tools 
-            if not is_agent_computer(tool)
-        ]
-    
+        tools = [tool for tool in tools if not is_agent_computer(tool)]
+
     # Execute evaluation
     return await run_dataset(
         name=job_name,
@@ -170,4 +164,4 @@ __all__ = [
     "run_single_task",
     "run_full_dataset",
     "MCPComputerAgent",
-]
+]

agent/integrations/hud/agent.py

@@ -9,26 +9,26 @@ Key differences from the OpenAI OperatorAgent variant:
 - Planning is executed via `ComputerAgent.run(messages)`.
 - The first yielded result per step is returned as the agent response.
 """
+
 from __future__ import annotations
 
+import base64
 import io
+import uuid
+from pathlib import Path
 from typing import Any, ClassVar, Optional
 
+import hud
+import mcp.types as types
 from agent.agent import ComputerAgent as BaseComputerAgent
 from agent.callbacks import PromptInstructionsCallback
 from agent.callbacks.trajectory_saver import TrajectorySaverCallback
+from agent.computers import is_agent_computer
+from agent.responses import make_failed_tool_call_items
 from hud.agents import MCPAgent
 from hud.tools.computer.settings import computer_settings
 from hud.types import AgentResponse, MCPToolCall, MCPToolResult, Trace
-
-from agent.responses import make_failed_tool_call_items
-from agent.computers import is_agent_computer
 from PIL import Image
-import mcp.types as types
-import hud
-import uuid
-import base64
-from pathlib import Path
 
 
 class MCPComputerAgent(MCPAgent):
@@ -114,8 +114,10 @@ class MCPComputerAgent(MCPAgent):
         self.last_screenshot_b64 = None
 
         buffer = io.BytesIO()
-        Image.new('RGB', (self.metadata["display_width"], self.metadata["display_height"])).save(buffer, format='PNG')
-        self.last_screenshot_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+        Image.new("RGB", (self.metadata["display_width"], self.metadata["display_height"])).save(
+            buffer, format="PNG"
+        )
+        self.last_screenshot_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
 
         # Ensure a computer shim is present so width/height/environment are known
         computer_shim = {
@@ -128,12 +130,8 @@ class MCPComputerAgent(MCPAgent):
         }
         agent_tools: list[Any] = [computer_shim]
         if tools:
-            agent_tools.extend([
-                tool 
-                for tool in tools 
-                if not is_agent_computer(tool)
-            ])
-        
+            agent_tools.extend([tool for tool in tools if not is_agent_computer(tool)])
+
         agent_kwargs = {
             "model": self.model,
             "trajectory_dir": trajectory_dir,
@@ -150,9 +148,7 @@ class MCPComputerAgent(MCPAgent):
             "telemetry_enabled": telemetry_enabled,
         }
 
-        self.computer_agent = BaseComputerAgent(
-            **agent_kwargs
-        )
+        self.computer_agent = BaseComputerAgent(**agent_kwargs)
 
     async def get_system_messages(self) -> list[Any]:
         """Create initial messages.
@@ -161,9 +157,7 @@ class MCPComputerAgent(MCPAgent):
         """
         return []
 
-    async def format_blocks(
-        self, blocks: list[types.ContentBlock]
-    ) -> list[dict[str, Any]]:
+    async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[dict[str, Any]]:
         """
         Format blocks for OpenAI input format.
 
@@ -200,42 +194,49 @@ class MCPComputerAgent(MCPAgent):
 
         # Call the ComputerAgent LLM API
         async for result in self.computer_agent.run(messages):  # type: ignore[arg-type]
-            items = result['output']
+            items = result["output"]
             if not items or tool_calls:
                 break
 
             for item in items:
-                if item['type'] in ['reasoning', 'message', 'computer_call', 'function_call', 'function_call_output']:
+                if item["type"] in [
+                    "reasoning",
+                    "message",
+                    "computer_call",
+                    "function_call",
+                    "function_call_output",
+                ]:
                     agent_result.append(item)
-                
+
                 # Add messages to output text
-                if item['type'] == 'reasoning':
+                if item["type"] == "reasoning":
                     output_text.extend(
-                        f"Reasoning: {summary['text']}"
-                        for summary in item['summary']
+                        f"Reasoning: {summary['text']}" for summary in item["summary"]
                     )
-                elif item['type'] == 'message':
-                    if isinstance(item['content'], list):
+                elif item["type"] == "message":
+                    if isinstance(item["content"], list):
                         output_text.extend(
-                            item['text'] 
-                            for item in item['content']
-                            if item['type'] == 'output_text'
+                            item["text"]
+                            for item in item["content"]
+                            if item["type"] == "output_text"
                         )
-                    elif isinstance(item['content'], str):
-                        output_text.append(item['content'])
-                
+                    elif isinstance(item["content"], str):
+                        output_text.append(item["content"])
+
                 # If we get a tool call, we're not done
-                if item['type'] == 'computer_call':
+                if item["type"] == "computer_call":
                     id = item["call_id"]
-                    tool_calls.append(MCPToolCall(
-                        name="openai_computer",
-                        arguments=item["action"],
-                        id=id,
-                    ))
+                    tool_calls.append(
+                        MCPToolCall(
+                            name="openai_computer",
+                            arguments=item["action"],
+                            id=id,
+                        )
+                    )
                     is_done = False
                     self.tool_call_inputs[id] = agent_result
                     break
-            
+
             # if we have tool calls, we should exit the loop
             if tool_calls:
                 break
@@ -247,7 +248,7 @@ class MCPComputerAgent(MCPAgent):
             tool_calls=tool_calls,
             done=is_done,
         )
-    
+
     def _log_image(self, image_b64: str):
         callbacks = self.computer_agent.callbacks
         for callback in callbacks:
@@ -257,9 +258,7 @@ class MCPComputerAgent(MCPAgent):
                 callback._save_artifact("screenshot_after", image_bytes)
 
     async def format_tool_results(
-        self,
-        tool_calls: list[MCPToolCall],
-        tool_results: list[MCPToolResult]
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
     ) -> list[dict[str, Any]]:
         """Extract latest screenshot from tool results in dict form.
 
@@ -274,45 +273,60 @@ class MCPComputerAgent(MCPAgent):
                 previous_output = self.previous_output.copy() or []
 
                 # First we need to remove any pending computer_calls from the end of previous_output
-                while previous_output and previous_output[-1]['type'] == 'computer_call':
+                while previous_output and previous_output[-1]["type"] == "computer_call":
                     previous_output.pop()
                 messages.extend(previous_output)
 
                 # If the call is a 'response', don't add the result
-                if call.name == 'response':
+                if call.name == "response":
                     continue
                 # Otherwise, if we have a result, we should add it to the messages
                 content = [
-                    { "type": "input_text", "text": content.text } if isinstance(content, types.TextContent)
-                    else { "type": "input_image", "image_url": f"data:image/png;base64,{content.data}" } if isinstance(content, types.ImageContent)
-                    else { "type": "input_text", "text": "" }
+                    (
+                        {"type": "input_text", "text": content.text}
+                        if isinstance(content, types.TextContent)
+                        else (
+                            {
+                                "type": "input_image",
+                                "image_url": f"data:image/png;base64,{content.data}",
+                            }
+                            if isinstance(content, types.ImageContent)
+                            else {"type": "input_text", "text": ""}
+                        )
+                    )
                     for content in result.content
                 ]
-                messages.append({
-                    "role": "user",
-                    "content": content,
-                })
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": content,
+                    }
+                )
 
                 continue
-                
+
             # Add the assistant's computer call
             messages.extend(self.tool_call_inputs[call.id])
-            
+
             if result.isError:
-                error_text = "".join([
-                    content.text
-                    for content in result.content
-                    if isinstance(content, types.TextContent)
-                ])
+                error_text = "".join(
+                    [
+                        content.text
+                        for content in result.content
+                        if isinstance(content, types.TextContent)
+                    ]
+                )
 
                 # Replace computer call with failed tool call
                 messages.pop()
-                messages.extend(make_failed_tool_call_items(
-                    tool_name=call.name,
-                    tool_kwargs=call.arguments or {},
-                    error_message=error_text,
-                    call_id=call.id,
-                ))
+                messages.extend(
+                    make_failed_tool_call_items(
+                        tool_name=call.name,
+                        tool_kwargs=call.arguments or {},
+                        error_message=error_text,
+                        call_id=call.id,
+                    )
+                )
             else:
                 # Get the latest screenshot
                 screenshots = [
@@ -325,23 +339,27 @@ class MCPComputerAgent(MCPAgent):
                 if screenshots:
                     self._log_image(screenshots[0])
                     self.last_screenshot_b64 = screenshots[0]
-                    messages.append({
-                        "type": "computer_call_output",
-                        "call_id": call.id,
-                        "output": {
-                            "type": "input_image",
-                            "image_url": f"data:image/png;base64,{screenshots[0]}"
-                        },
-                    })
+                    messages.append(
+                        {
+                            "type": "computer_call_output",
+                            "call_id": call.id,
+                            "output": {
+                                "type": "input_image",
+                                "image_url": f"data:image/png;base64,{screenshots[0]}",
+                            },
+                        }
+                    )
                 else:
                     # Otherwise, replace computer call with failed tool call
                     messages.pop()
-                    messages.extend(make_failed_tool_call_items(
-                        tool_name=call.name,
-                        tool_kwargs=call.arguments or {},
-                        error_message="No screenshots returned.",
-                        call_id=call.id,
-                    ))
+                    messages.extend(
+                        make_failed_tool_call_items(
+                            tool_name=call.name,
+                            tool_kwargs=call.arguments or {},
+                            error_message="No screenshots returned.",
+                            call_id=call.id,
+                        )
+                    )
 
         return messages
 

agent/integrations/hud/proxy.py

@@ -7,30 +7,33 @@ OpenAI-like response blocks. We intentionally only support a single-step call
 by consuming the first yielded result from `ComputerAgent.run()`.
 """
 
-import traceback
 import time
+import traceback
 import uuid
 from typing import Any, Dict, List, Optional
 
 from agent.agent import ComputerAgent as BaseComputerAgent
 from agent.callbacks import PromptInstructionsCallback
-from hud.tools.computer.settings import computer_settings
-from PIL import Image
 from hud.agents import OperatorAgent
+from hud.tools.computer.settings import computer_settings
 
 # OpenAI Responses typed models (required)
 from openai.types.responses import (
     Response,
+    ResponseComputerToolCall,
     ResponseInputParam,
     ResponseOutputItem,
-    ResponseComputerToolCall,
     ResponseOutputMessage,
     ResponseOutputText,
     ResponseReasoningItem,
     ResponseUsage,
 )
+from PIL import Image
 
-def _map_agent_output_to_openai_blocks(output_items: List[Dict[str, Any]]) -> List[ResponseOutputItem]:
+
+def _map_agent_output_to_openai_blocks(
+    output_items: List[Dict[str, Any]],
+) -> List[ResponseOutputItem]:
     """Map our agent output items to OpenAI ResponseOutputItem typed models.
 
     Only a subset is supported: computer_call, assistant message (text), and reasoning.
@@ -40,14 +43,16 @@ def _map_agent_output_to_openai_blocks(output_items: List[Dict[str, Any]]) -> Li
     for item in output_items or []:
         t = item.get("type")
         if t == "computer_call":
-            comp = ResponseComputerToolCall.model_validate({
-                "id": item.get("id") or f"cu_{uuid.uuid4().hex}",
-                "type": "computer_call",
-                "call_id": item["call_id"],
-                "action": item["action"],
-                "pending_safety_checks": item.get("pending_safety_checks", []),
-                "status": "completed",
-            })
+            comp = ResponseComputerToolCall.model_validate(
+                {
+                    "id": item.get("id") or f"cu_{uuid.uuid4().hex}",
+                    "type": "computer_call",
+                    "call_id": item["call_id"],
+                    "action": item["action"],
+                    "pending_safety_checks": item.get("pending_safety_checks", []),
+                    "status": "completed",
+                }
+            )
             blocks.append(comp)
             # we will exit early here as the responses api only supports a single step
             break
@@ -55,31 +60,38 @@ def _map_agent_output_to_openai_blocks(output_items: List[Dict[str, Any]]) -> Li
             content_blocks: List[ResponseOutputText] = []
             for c in item.get("content", []) or []:
                 content_blocks.append(
-                    ResponseOutputText.model_validate({
-                        "type": "output_text",
-                        "text": c["text"],
-                        "annotations": [],
-                    })
+                    ResponseOutputText.model_validate(
+                        {
+                            "type": "output_text",
+                            "text": c["text"],
+                            "annotations": [],
+                        }
+                    )
                 )
             if content_blocks:
-                msg = ResponseOutputMessage.model_validate({
-                    "id": item.get("id") or f"msg_{uuid.uuid4()}",
-                    "type": "message",
-                    "role": "assistant",
-                    "status": "completed",
-                    "content": [ct.model_dump() for ct in content_blocks],
-                })
+                msg = ResponseOutputMessage.model_validate(
+                    {
+                        "id": item.get("id") or f"msg_{uuid.uuid4()}",
+                        "type": "message",
+                        "role": "assistant",
+                        "status": "completed",
+                        "content": [ct.model_dump() for ct in content_blocks],
+                    }
+                )
                 blocks.append(msg)
         elif t == "reasoning":
-            reasoning = ResponseReasoningItem.model_validate({
-                "id": item.get("id") or f"rsn_{uuid.uuid4()}",
-                "type": "reasoning",
-                "summary": item["summary"],
-            })
+            reasoning = ResponseReasoningItem.model_validate(
+                {
+                    "id": item.get("id") or f"rsn_{uuid.uuid4()}",
+                    "type": "reasoning",
+                    "summary": item["summary"],
+                }
+            )
             blocks.append(reasoning)
         # Unhandled types are ignored
     return blocks
 
+
 def _to_plain_dict_list(items: Any) -> List[Dict[str, Any]]:
     out: List[Dict[str, Any]] = []
     for it in list(items):
@@ -92,6 +104,7 @@ def _to_plain_dict_list(items: Any) -> List[Dict[str, Any]]:
             out.append(dict(it))  # may raise if not mapping
     return out
 
+
 class FakeAsyncOpenAI:
     """Minimal fake OpenAI client with only `responses.create` implemented.
 
@@ -132,10 +145,12 @@ class FakeAsyncOpenAI:
                 # Pre-pend instructions message
                 effective_input = full_input
                 if instructions:
-                    effective_input = [{
-                        "role": "user",
-                        "content": instructions,
-                    }] + full_input
+                    effective_input = [
+                        {
+                            "role": "user",
+                            "content": instructions,
+                        }
+                    ] + full_input
 
                 # Run a single iteration of the ComputerAgent
                 agent_result: Optional[Dict[str, Any]] = None
@@ -152,32 +167,43 @@ class FakeAsyncOpenAI:
                 blocks_to_cache = full_input + output
                 for b in blocks_to_cache:
                     bid = getattr(b, "id", None) or f"tmp-{hash(repr(b))}"
-                    self.blocks_cache[bid] = b # type: ignore[assignment]
+                    self.blocks_cache[bid] = b  # type: ignore[assignment]
                     block_ids.append(bid)
                 response_id = agent_result.get("id") or f"fake-{int(time.time()*1000)}"
                 self.context_cache[response_id] = block_ids
 
                 try:
-                    return Response.model_validate({
-                        "id": response_id,
-                        "created_at": time.time(),
-                        "object": "response",
-                        "model": model,
-                        "output": output,
-                        "parallel_tool_calls": False,
-                        "tool_choice": "auto",
-                        "tools": [],
-                        "previous_response_id": previous_response_id,
-                        "usage": ResponseUsage.model_validate({
-                            "input_tokens": usage.get("input_tokens", 0),
-                            "output_tokens": usage.get("output_tokens", 0),
-                            "total_tokens": usage.get("total_tokens", 0),
-                            "input_tokens_details": usage.get("input_tokens_details", { "cached_tokens": 0 }),
-                            "output_tokens_details": usage.get("output_tokens_details", { "reasoning_tokens": 0 }),
-                        }),
-                    })
+                    return Response.model_validate(
+                        {
+                            "id": response_id,
+                            "created_at": time.time(),
+                            "object": "response",
+                            "model": model,
+                            "output": output,
+                            "parallel_tool_calls": False,
+                            "tool_choice": "auto",
+                            "tools": [],
+                            "previous_response_id": previous_response_id,
+                            "usage": ResponseUsage.model_validate(
+                                {
+                                    "input_tokens": usage.get("input_tokens", 0),
+                                    "output_tokens": usage.get("output_tokens", 0),
+                                    "total_tokens": usage.get("total_tokens", 0),
+                                    "input_tokens_details": usage.get(
+                                        "input_tokens_details", {"cached_tokens": 0}
+                                    ),
+                                    "output_tokens_details": usage.get(
+                                        "output_tokens_details", {"reasoning_tokens": 0}
+                                    ),
+                                }
+                            ),
+                        }
+                    )
                 except Exception as e:
-                    print(f"Error while validating agent response (attempt {attempt + 1}/{max_retries}): ", e)
+                    print(
+                        f"Error while validating agent response (attempt {attempt + 1}/{max_retries}): ",
+                        e,
+                    )
                     if attempt == max_retries - 1:
                         print(traceback.format_exc())
                         raise e
@@ -221,9 +247,15 @@ class ProxyOperatorAgent(OperatorAgent):
         allowed_tools = allowed_tools or ["openai_computer"]
 
         computer_shim = {
-            'screenshot': lambda: Image.new('RGB', (computer_settings.OPENAI_COMPUTER_WIDTH, computer_settings.OPENAI_COMPUTER_HEIGHT)),
-            'environment': 'linux',
-            'dimensions': (computer_settings.OPENAI_COMPUTER_WIDTH, computer_settings.OPENAI_COMPUTER_HEIGHT)
+            "screenshot": lambda: Image.new(
+                "RGB",
+                (computer_settings.OPENAI_COMPUTER_WIDTH, computer_settings.OPENAI_COMPUTER_HEIGHT),
+            ),
+            "environment": "linux",
+            "dimensions": (
+                computer_settings.OPENAI_COMPUTER_WIDTH,
+                computer_settings.OPENAI_COMPUTER_HEIGHT,
+            ),
         }
         # Build tools ensuring the computer_shim is included
         agent_tools: list[Any] = [computer_shim]
@@ -258,6 +290,7 @@ class ProxyOperatorAgent(OperatorAgent):
             **kwargs,
         )
 
+
 __all__ = [
     "FakeAsyncOpenAI",
     "ProxyOperatorAgent",