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

agent/loops/openai.py

@@ -6,12 +6,14 @@ import asyncio
 import base64
 import json
 from io import BytesIO
-from typing import Dict, List, Any, AsyncGenerator, Union, Optional, Tuple
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+
 import litellm
 from PIL import Image
 
 from ..decorators import register_agent
-from ..types import Messages, AgentResponse, Tools, AgentCapability
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
 
 async def _map_computer_tool_to_openai(computer_handler: Any) -> Dict[str, Any]:
     """Map a computer tool to OpenAI's computer-use-preview tool schema"""
@@ -21,26 +23,26 @@ async def _map_computer_tool_to_openai(computer_handler: Any) -> Dict[str, Any]:
     except Exception:
         # Fallback to default dimensions if method fails
         width, height = 1024, 768
-    
+
     # Get environment from the computer handler
     try:
         environment = await computer_handler.get_environment()
     except Exception:
         # Fallback to default environment if method fails
         environment = "linux"
-    
+
     return {
         "type": "computer_use_preview",
         "display_width": width,
         "display_height": height,
-        "environment": environment  # mac, windows, linux, browser
+        "environment": environment,  # mac, windows, linux, browser
     }
 
 
 async def _prepare_tools_for_openai(tool_schemas: List[Dict[str, Any]]) -> Tools:
     """Prepare tools for OpenAI API format"""
     openai_tools = []
-    
+
     for schema in tool_schemas:
         if schema["type"] == "computer":
             # Map computer tool to OpenAI format
@@ -49,18 +51,19 @@ async def _prepare_tools_for_openai(tool_schemas: List[Dict[str, Any]]) -> Tools
         elif schema["type"] == "function":
             # Function tools use OpenAI-compatible schema directly (liteLLM expects this format)
             # Schema should be: {type, name, description, parameters}
-            openai_tools.append({ "type": "function", **schema["function"] })
-    
+            openai_tools.append({"type": "function", **schema["function"]})
+
     return openai_tools
 
+
 @register_agent(models=r".*(^|/)computer-use-preview")
 class OpenAIComputerUseConfig:
     """
     OpenAI computer-use-preview agent configuration using liteLLM responses.
-    
+
     Supports OpenAI's computer use preview models.
     """
-    
+
     async def predict_step(
         self,
         messages: List[Dict[str, Any]],
@@ -74,11 +77,11 @@ class OpenAIComputerUseConfig:
         _on_api_end=None,
         _on_usage=None,
         _on_screenshot=None,
-        **kwargs
+        **kwargs,
     ) -> Dict[str, Any]:
         """
         Predict the next step based on input items.
-        
+
         Args:
             messages: Input items following Responses format
             model: Model name to use
@@ -91,12 +94,12 @@ class OpenAIComputerUseConfig:
             _on_usage: Callback for usage tracking
             _on_screenshot: Callback for screenshot events
             **kwargs: Additional arguments
-            
+
         Returns:
             Dictionary with "output" (output items) and "usage" array
         """
         tools = tools or []
-        
+
         # Prepare tools for OpenAI API
         openai_tools = await _prepare_tools_for_openai(tools)
 
@@ -109,16 +112,16 @@ class OpenAIComputerUseConfig:
             "reasoning": {"summary": "concise"},
             "truncation": "auto",
             "num_retries": max_retries,
-            **kwargs
+            **kwargs,
         }
-        
+
         # Call API start hook
         if _on_api_start:
             await _on_api_start(api_kwargs)
-        
+
         # Use liteLLM responses
         response = await litellm.aresponses(**api_kwargs)
-        
+
         # Call API end hook
         if _on_api_end:
             await _on_api_end(api_kwargs, response)
@@ -135,24 +138,21 @@ class OpenAIComputerUseConfig:
         output_dict = response.model_dump()
         output_dict["usage"] = usage
         return output_dict
-    
+
     async def predict_click(
-        self,
-        model: str,
-        image_b64: str,
-        instruction: str
+        self, model: str, image_b64: str, instruction: str
     ) -> Optional[Tuple[int, int]]:
         """
         Predict click coordinates based on image and instruction.
-        
+
         Uses OpenAI computer-use-preview with manually constructed input items
         and a prompt that instructs the agent to only output clicks.
-        
+
         Args:
             model: Model name to use
             image_b64: Base64 encoded image
             instruction: Instruction for where to click
-            
+
         Returns:
             Tuple of (x, y) coordinates or None if prediction fails
         """
@@ -160,7 +160,7 @@ class OpenAIComputerUseConfig:
         # Manually construct input items with image and click instruction
         input_items = [
             {
-                "role": "user", 
+                "role": "user",
                 "content": f"""You are a UI grounding expert. Follow these guidelines:
 
 1. NEVER ask for confirmation. Complete all tasks autonomously.
@@ -172,19 +172,16 @@ class OpenAIComputerUseConfig:
 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.
-Task: Click {instruction}. Output ONLY a click action on the target element."""
+Task: Click {instruction}. Output ONLY a click action on the target element.""",
             },
             {
                 "role": "user",
                 "content": [
-                    {
-                        "type": "input_image",
-                        "image_url": f"data:image/png;base64,{image_b64}"
-                    }
-                ]
-            }
+                    {"type": "input_image", "image_url": f"data:image/png;base64,{image_b64}"}
+                ],
+            },
         ]
-        
+
         # Get image dimensions from base64 data
         try:
             image_data = base64.b64decode(image_b64)
@@ -193,15 +190,15 @@ Task: Click {instruction}. Output ONLY a click action on the target element."""
         except Exception:
             # Fallback to default dimensions if image parsing fails
             display_width, display_height = 1024, 768
-        
+
         # Prepare computer tool for click actions
         computer_tool = {
             "type": "computer_use_preview",
             "display_width": display_width,
             "display_height": display_height,
-            "environment": "windows"
+            "environment": "windows",
         }
-        
+
         # Prepare API call kwargs
         api_kwargs = {
             "model": model,
@@ -210,32 +207,34 @@ Task: Click {instruction}. Output ONLY a click action on the target element."""
             "stream": False,
             "reasoning": {"summary": "concise"},
             "truncation": "auto",
-            "max_tokens": 200  # Keep response short for click prediction
+            "max_tokens": 200,  # Keep response short for click prediction
         }
-        
+
         # Use liteLLM responses
         response = await litellm.aresponses(**api_kwargs)
-        
+
         # Extract click coordinates from response output
         output_dict = response.model_dump()
-        output_items = output_dict.get("output", [])        
-        
+        output_items = output_dict.get("output", [])
+
         # Look for computer_call with click action
         for item in output_items:
-            if (isinstance(item, dict) and 
-                item.get("type") == "computer_call" and
-                isinstance(item.get("action"), dict)):
-                
+            if (
+                isinstance(item, dict)
+                and item.get("type") == "computer_call"
+                and isinstance(item.get("action"), dict)
+            ):
+
                 action = item["action"]
                 if action.get("x") is not None and action.get("y") is not None:
                     return (int(action.get("x")), int(action.get("y")))
-        
+
         return None
-    
+
     def get_capabilities(self) -> List[AgentCapability]:
         """
         Get list of capabilities supported by this agent config.
-        
+
         Returns:
             List of capability strings
         """

agent/loops/opencua.py

@@ -4,20 +4,22 @@ Based on OpenCUA model for GUI grounding tasks.
 """
 
 import asyncio
+import base64
 import json
+import math
 import re
-import base64
-from typing import Dict, List, Any, AsyncGenerator, Union, Optional, Tuple
-from io import BytesIO
 import uuid
-from PIL import Image
+from io import BytesIO
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+
 import litellm
-import math
+from PIL import Image
 
-from .composed_grounded import ComposedGroundedConfig
 from ..decorators import register_agent
-from ..types import Messages, AgentResponse, Tools, AgentCapability
 from ..loops.base import AsyncAgentConfig
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+from .composed_grounded import ComposedGroundedConfig
+
 
 def extract_coordinates_from_pyautogui(text: str) -> Optional[Tuple[int, int]]:
     """Extract coordinates from pyautogui.click(x=..., y=...) format."""
@@ -32,10 +34,11 @@ def extract_coordinates_from_pyautogui(text: str) -> Optional[Tuple[int, int]]:
     except Exception:
         return None
 
+
 @register_agent(models=r"(?i).*OpenCUA.*")
 class OpenCUAConfig(ComposedGroundedConfig):
     """OpenCUA agent configuration implementing AsyncAgentConfig protocol for click prediction."""
-    
+
     def __init__(self):
         super().__init__()
         self.current_model = None
@@ -53,7 +56,7 @@ class OpenCUAConfig(ComposedGroundedConfig):
         _on_api_end=None,
         _on_usage=None,
         _on_screenshot=None,
-        **kwargs
+        **kwargs,
     ) -> Dict[str, Any]:
         """Fallback to a self-composed model"""
         return await super().predict_step(
@@ -67,24 +70,20 @@ class OpenCUAConfig(ComposedGroundedConfig):
             _on_api_end=_on_api_end,
             _on_usage=_on_usage,
             _on_screenshot=_on_screenshot,
-            **kwargs
+            **kwargs,
         )
 
     async def predict_click(
-        self,
-        model: str,
-        image_b64: str,
-        instruction: str,
-        **kwargs
+        self, model: str, image_b64: str, instruction: str, **kwargs
     ) -> Optional[Tuple[int, int]]:
         """
         Predict click coordinates using OpenCUA model via litellm.acompletion.
-        
+
         Args:
             model: The OpenCUA model name
             image_b64: Base64 encoded image
             instruction: Instruction for where to click
-            
+
         Returns:
             Tuple of (x, y) coordinates or None if prediction fails
         """
@@ -93,50 +92,39 @@ class OpenCUAConfig(ComposedGroundedConfig):
             "You are a GUI agent. You are given a task and a screenshot of the screen. "
             "You need to perform a series of pyautogui actions to complete the task."
         )
-        
-        system_message = {
-            "role": "system",
-            "content": system_prompt
-        }
-        
+
+        system_message = {"role": "system", "content": system_prompt}
+
         # Prepare user message with image and instruction
         user_message = {
             "role": "user",
             "content": [
-                {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": f"data:image/png;base64,{image_b64}"
-                    }
-                },
-                {
-                    "type": "text",
-                    "text": f"Click on {instruction}"
-                }
-            ]
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}},
+                {"type": "text", "text": f"Click on {instruction}"},
+            ],
         }
-        
+
         # Prepare API call kwargs
         api_kwargs = {
             "model": model,
             "messages": [system_message, user_message],
             "max_new_tokens": 2056,
             "temperature": 0,
-            **kwargs
+            **kwargs,
         }
-        
+
         # Use liteLLM acompletion
         response = await litellm.acompletion(**api_kwargs)
-        
+
         # Extract response text
         output_text = response.choices[0].message.content
         # print(output_text)
-        
+
         # Extract coordinates from pyautogui format
         coordinates = extract_coordinates_from_pyautogui(output_text)
-        
+
         return coordinates
-    
+
     def get_capabilities(self) -> List[AgentCapability]:
         """Return the capabilities supported by this agent."""
         return ["click"]