cua-agent 0.3.2__py3-none-any.whl → 0.4.0b1__py3-none-any.whl

agent/loops/omniparser.py

@@ -0,0 +1,339 @@
+"""
+OpenAI computer-use-preview agent loop implementation using liteLLM
+"""
+
+import asyncio
+import json
+from typing import Dict, List, Any, AsyncGenerator, Union, Optional, Tuple
+import litellm
+import inspect
+import base64
+
+from ..decorators import agent_loop
+from ..types import Messages, AgentResponse, Tools
+
+SOM_TOOL_SCHEMA = {
+  "type": "function",
+  "name": "computer",
+  "description": "Control a computer by taking screenshots and interacting with UI elements. This tool shows screenshots with numbered elements overlaid on them. Each UI element has been assigned a unique ID number that you can see in the image. Use the element's ID number to interact with any element instead of pixel coordinates.",
+  "parameters": {
+    "type": "object",
+    "properties": {
+      "action": {
+        "type": "string",
+        "enum": [
+          "screenshot",
+          "click",
+          "double_click",
+          "drag",
+          "type",
+          "keypress",
+          "scroll",
+          "move",
+          "wait",
+          "get_current_url",
+          "get_dimensions",
+          "get_environment"
+        ],
+        "description": "The action to perform"
+      },
+      "element_id": {
+        "type": "integer",
+        "description": "The ID of the element to interact with (required for click, double_click, move, scroll actions, and as start/end for drag)"
+      },
+      "start_element_id": {
+        "type": "integer",
+        "description": "The ID of the element to start dragging from (required for drag action)"
+      },
+      "end_element_id": {
+        "type": "integer",
+        "description": "The ID of the element to drag to (required for drag action)"
+      },
+      "text": {
+        "type": "string",
+        "description": "The text to type (required for type action)"
+      },
+      "keys": {
+        "type": "string",
+        "description": "Key combination to press (required for keypress action). Single key for individual key press, multiple keys for combinations (e.g., 'ctrl+c')"
+      },
+      "button": {
+        "type": "string",
+        "description": "The mouse button to use for click action (left, right, wheel, back, forward) Default: left",
+      },
+      "scroll_x": {
+        "type": "integer",
+        "description": "Horizontal scroll amount for scroll action (positive for right, negative for left)",
+      },
+      "scroll_y": {
+        "type": "integer",
+        "description": "Vertical scroll amount for scroll action (positive for down, negative for up)",
+      },
+    },
+    "required": [
+      "action"
+    ]
+  }
+}
+
+OMNIPARSER_AVAILABLE = False
+try:
+    from som import OmniParser
+    OMNIPARSER_AVAILABLE = True
+except ImportError:
+    pass
+OMNIPARSER_SINGLETON = None
+
+def get_parser() -> OmniParser:
+    global OMNIPARSER_SINGLETON
+    if OMNIPARSER_SINGLETON is None:
+        OMNIPARSER_SINGLETON = OmniParser()
+    return OMNIPARSER_SINGLETON
+    
+def get_last_computer_call_output(messages: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Get the last computer_call_output message from a messages list.
+    
+    Args:
+        messages: List of messages to search through
+        
+    Returns:
+        The last computer_call_output message dict, or None if not found
+    """
+    for message in reversed(messages):
+        if isinstance(message, dict) and message.get("type") == "computer_call_output":
+            return message
+    return None
+
+def _prepare_tools_for_omniparser(tool_schemas: List[Dict[str, Any]]) -> Tuple[Tools, dict]:
+    """Prepare tools for OpenAI API format"""
+    omniparser_tools = []
+    id2xy = dict()
+    
+    for schema in tool_schemas:
+        if schema["type"] == "computer":
+            omniparser_tools.append(SOM_TOOL_SCHEMA)
+            if "id2xy" in schema:
+                id2xy = schema["id2xy"]
+            else:
+                schema["id2xy"] = id2xy
+        elif schema["type"] == "function":
+            # Function tools use OpenAI-compatible schema directly (liteLLM expects this format)
+            # Schema should be: {type, name, description, parameters}
+            omniparser_tools.append({ "type": "function", **schema["function"] })
+    
+    return omniparser_tools, id2xy
+
+async def replace_function_with_computer_call(item: Dict[str, Any], id2xy: Dict[int, Tuple[float, float]]):
+  item_type = item.get("type")
+
+  def _get_xy(element_id: Optional[int]) -> Union[Tuple[float, float], Tuple[None, None]]:
+    if element_id is None:
+      return (None, None)
+    return id2xy.get(element_id, (None, None))
+
+  if item_type == "function_call":
+    fn_name = item.get("name")
+    fn_args = json.loads(item.get("arguments", "{}"))
+
+    item_id = item.get("id")
+    call_id = item.get("call_id")
+    
+    if fn_name == "computer":
+      action = fn_args.get("action")
+      element_id = fn_args.get("element_id")
+      start_element_id = fn_args.get("start_element_id")
+      end_element_id = fn_args.get("end_element_id")
+      text = fn_args.get("text")
+      keys = fn_args.get("keys")
+      button = fn_args.get("button")
+      scroll_x = fn_args.get("scroll_x")
+      scroll_y = fn_args.get("scroll_y")
+
+      x, y = _get_xy(element_id)
+      start_x, start_y = _get_xy(start_element_id)
+      end_x, end_y = _get_xy(end_element_id)
+
+      action_args = {
+          "type": action,
+          "x": x,
+          "y": y,
+          "start_x": start_x,
+          "start_y": start_y,
+          "end_x": end_x,
+          "end_y": end_y,
+          "text": text,
+          "keys": keys,
+          "button": button,
+          "scroll_x": scroll_x,
+          "scroll_y": scroll_y
+        }
+      # Remove None values to keep the JSON clean
+      action_args = {k: v for k, v in action_args.items() if v is not None}
+
+      return [{
+        "type": "computer_call",
+        "action": action_args,
+        "id": item_id,
+        "call_id": call_id,
+        "status": "completed"
+      }]
+
+  return [item]
+
+async def replace_computer_call_with_function(item: Dict[str, Any], xy2id: Dict[Tuple[float, float], int]):
+    """
+    Convert computer_call back to function_call format.
+    Also handles computer_call_output -> function_call_output conversion.
+    
+    Args:
+        item: The item to convert
+        xy2id: Mapping from (x, y) coordinates to element IDs
+    """
+    item_type = item.get("type")
+
+    def _get_element_id(x: Optional[float], y: Optional[float]) -> Optional[int]:
+        """Get element ID from coordinates, return None if coordinates are None"""
+        if x is None or y is None:
+            return None
+        return xy2id.get((x, y))
+
+    if item_type == "computer_call":
+        action_data = item.get("action", {})
+        
+        # Extract coordinates and convert back to element IDs
+        element_id = _get_element_id(action_data.get("x"), action_data.get("y"))
+        start_element_id = _get_element_id(action_data.get("start_x"), action_data.get("start_y"))
+        end_element_id = _get_element_id(action_data.get("end_x"), action_data.get("end_y"))
+        
+        # Build function arguments
+        fn_args = {
+            "action": action_data.get("type"),
+            "element_id": element_id,
+            "start_element_id": start_element_id,
+            "end_element_id": end_element_id,
+            "text": action_data.get("text"),
+            "keys": action_data.get("keys"),
+            "button": action_data.get("button"),
+            "scroll_x": action_data.get("scroll_x"),
+            "scroll_y": action_data.get("scroll_y")
+        }
+        
+        # Remove None values to keep the JSON clean
+        fn_args = {k: v for k, v in fn_args.items() if v is not None}
+        
+        return [{
+            "type": "function_call",
+            "name": "computer",
+            "arguments": json.dumps(fn_args),
+            "id": item.get("id"),
+            "call_id": item.get("call_id"),
+            "status": "completed",
+
+            # Fall back to string representation
+            "content": f"Used tool: {action_data.get("type")}({json.dumps(fn_args)})"
+        }]
+    
+    elif item_type == "computer_call_output":
+        # Simple conversion: computer_call_output -> function_call_output
+        return [{
+            "type": "function_call_output",
+            "call_id": item.get("call_id"),
+            "content": [item.get("output")],
+            "id": item.get("id"),
+            "status": "completed"
+        }]
+
+    return [item]
+
+
+@agent_loop(models=r"omniparser\+.*|omni\+.*", priority=10)
+async def omniparser_loop(
+    messages: Messages,
+    model: str,
+    tools: Optional[List[Dict[str, Any]]] = None,
+    max_retries: Optional[int] = None,
+    stream: bool = False,
+    computer_handler=None,
+    use_prompt_caching: Optional[bool] = False,
+    _on_api_start=None,
+    _on_api_end=None,
+    _on_usage=None,
+    _on_screenshot=None,
+    **kwargs
+) -> Union[AgentResponse, AsyncGenerator[Dict[str, Any], None]]:
+    """
+    OpenAI computer-use-preview agent loop using liteLLM responses.
+    
+    Supports OpenAI's computer use preview models.
+    """
+    if not OMNIPARSER_AVAILABLE:
+        raise ValueError("omniparser loop requires som to be installed. Install it with `pip install cua-som`.")
+      
+    tools = tools or []
+    
+    llm_model = model.split('+')[-1]
+
+    # Prepare tools for OpenAI API
+    openai_tools, id2xy = _prepare_tools_for_omniparser(tools)
+
+    # Find last computer_call_output
+    last_computer_call_output = get_last_computer_call_output(messages)
+    if last_computer_call_output:
+        image_url = last_computer_call_output.get("output", {}).get("image_url", "")
+        image_data = image_url.split(",")[-1]
+        if image_data:
+            parser = get_parser()
+            result = parser.parse(image_data)
+            if _on_screenshot:
+                await _on_screenshot(result.annotated_image_base64, "annotated_image")
+            for element in result.elements:
+                id2xy[element.id] = ((element.bbox.x1 + element.bbox.x2) / 2, (element.bbox.y1 + element.bbox.y2) / 2)
+    
+    # handle computer calls -> function calls
+    new_messages = []
+    for message in messages:
+        if not isinstance(message, dict):
+            message = message.__dict__
+        new_messages += await replace_computer_call_with_function(message, id2xy)
+    messages = new_messages
+
+    # Prepare API call kwargs
+    api_kwargs = {
+        "model": llm_model,
+        "input": messages,
+        "tools": openai_tools if openai_tools else None,
+        "stream": stream,
+        "reasoning": {"summary": "concise"},
+        "truncation": "auto",
+        "num_retries": max_retries,
+        **kwargs
+    }
+    
+    # Call API start hook
+    if _on_api_start:
+        await _on_api_start(api_kwargs)
+    
+    print(str(api_kwargs)[:1000])
+
+    # Use liteLLM responses
+    response = await litellm.aresponses(**api_kwargs)
+
+    # Call API end hook
+    if _on_api_end:
+        await _on_api_end(api_kwargs, response)
+
+    # Extract usage information
+    response.usage = {
+        **response.usage.model_dump(),
+        "response_cost": response._hidden_params.get("response_cost", 0.0),
+    }
+    if _on_usage:
+        await _on_usage(response.usage)
+
+    # handle som function calls -> xy computer calls
+    new_output = []
+    for i in range(len(response.output)):
+      new_output += await replace_function_with_computer_call(response.output[i].model_dump(), id2xy)
+    response.output = new_output
+
+    return response

agent/loops/openai.py

@@ -0,0 +1,95 @@
+"""
+OpenAI computer-use-preview agent loop implementation using liteLLM
+"""
+
+import asyncio
+import json
+from typing import Dict, List, Any, AsyncGenerator, Union, Optional
+import litellm
+
+from ..decorators import agent_loop
+from ..types import Messages, AgentResponse, Tools
+
+def _map_computer_tool_to_openai(computer_tool: Any) -> Dict[str, Any]:
+    """Map a computer tool to OpenAI's computer-use-preview tool schema"""
+    return {
+        "type": "computer_use_preview",
+        "display_width": getattr(computer_tool, 'display_width', 1024),
+        "display_height": getattr(computer_tool, 'display_height', 768),
+        "environment": getattr(computer_tool, 'environment', "linux")  # mac, windows, linux, browser
+    }
+
+
+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
+            openai_tools.append(_map_computer_tool_to_openai(schema["computer"]))
+        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"] })
+    
+    return openai_tools
+
+
+@agent_loop(models=r".*computer-use-preview.*", priority=10)
+async def openai_computer_use_loop(
+    messages: Messages,
+    model: str,
+    tools: Optional[List[Dict[str, Any]]] = None,
+    max_retries: Optional[int] = None,
+    stream: bool = False,
+    computer_handler=None,
+    use_prompt_caching: Optional[bool] = False,
+    _on_api_start=None,
+    _on_api_end=None,
+    _on_usage=None,
+    _on_screenshot=None,
+    **kwargs
+) -> Union[AgentResponse, AsyncGenerator[Dict[str, Any], None]]:
+    """
+    OpenAI computer-use-preview agent loop using liteLLM responses.
+    
+    Supports OpenAI's computer use preview models.
+    """
+    tools = tools or []
+    
+    # Prepare tools for OpenAI API
+    openai_tools = _prepare_tools_for_openai(tools)
+
+    # Prepare API call kwargs
+    api_kwargs = {
+        "model": model,
+        "input": messages,
+        "tools": openai_tools if openai_tools else None,
+        "stream": stream,
+        "reasoning": {"summary": "concise"},
+        "truncation": "auto",
+        "num_retries": max_retries,
+        **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)
+
+    # Extract usage information
+    response.usage = {
+        **response.usage.model_dump(),
+        "response_cost": response._hidden_params.get("response_cost", 0.0),
+    }
+    if _on_usage:
+        await _on_usage(response.usage)
+    
+    return response