cua-agent 0.4.6__py3-none-any.whl → 0.4.8__py3-none-any.whl

agent/loops/composed_grounded.py

@@ -0,0 +1,318 @@
+"""
+Composed-grounded agent loop implementation that combines grounding and thinking models.
+Uses a two-stage approach: grounding model for element detection, thinking model for reasoning.
+"""
+
+import uuid
+import asyncio
+import json
+import base64
+from typing import Dict, List, Any, Optional, Tuple
+from io import BytesIO
+from PIL import Image
+import litellm
+
+from ..decorators import register_agent
+from ..types import Messages, AgentResponse, Tools, AgentCapability
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_computer_calls_xy2desc,
+    convert_responses_items_to_completion_messages,
+    convert_completion_messages_to_responses_items,
+    convert_computer_calls_desc2xy,
+    get_all_element_descriptions
+)
+from ..agent import find_agent_config
+
+GROUNDED_COMPUTER_TOOL_SCHEMA = {
+  "type": "function",
+  "function": {
+    "name": "computer",
+    "description": "Control a computer by taking screenshots and interacting with UI elements. This tool uses element descriptions to locate and interact with UI elements on the screen (e.g., 'red submit button', 'search text field', 'hamburger menu icon', 'close button in top right corner').",
+    "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_description": {
+            "type": "string",
+            "description": "Description of the element to interact with (required for click, double_click, move, scroll actions, and as start/end for drag)"
+        },
+        "start_element_description": {
+            "type": "string",
+            "description": "Description of the element to start dragging from (required for drag action)"
+        },
+        "end_element_description": {
+            "type": "string",
+            "description": "Description 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"
+        ]
+    }
+  }
+}
+
+def _prepare_tools_for_grounded(tool_schemas: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Prepare tools for grounded API format"""
+    grounded_tools = []
+    
+    for schema in tool_schemas:
+        if schema["type"] == "computer":
+            grounded_tools.append(GROUNDED_COMPUTER_TOOL_SCHEMA)
+        else:
+            grounded_tools.append(schema)
+    
+    return grounded_tools
+
+def get_last_computer_call_image(messages: List[Dict[str, Any]]) -> Optional[str]:
+    """Get the last computer call output image from messages."""
+    for message in reversed(messages):
+        if (isinstance(message, dict) and 
+            message.get("type") == "computer_call_output" and
+            isinstance(message.get("output"), dict) and
+            message["output"].get("type") == "input_image"):
+            image_url = message["output"].get("image_url", "")
+            if image_url.startswith("data:image/png;base64,"):
+                return image_url.split(",", 1)[1]
+    return None
+
+
+@register_agent(r".*\+.*", priority=1)
+class ComposedGroundedConfig:
+    """
+    Composed-grounded agent configuration that uses both grounding and thinking models.
+    
+    The model parameter should be in format: "grounding_model+thinking_model"
+    e.g., "huggingface-local/HelloKKMe/GTA1-7B+gemini/gemini-1.5-pro"
+    """
+    
+    def __init__(self):
+        self.desc2xy: Dict[str, Tuple[float, float]] = {}
+    
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        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
+    ) -> Dict[str, Any]:
+        """
+        Composed-grounded predict step implementation.
+        
+        Process:
+        0. Store last computer call image, if none then take a screenshot
+        1. Convert computer calls from xy to descriptions
+        2. Convert responses items to completion messages
+        3. Call thinking model with litellm.acompletion
+        4. Convert completion messages to responses items
+        5. Get all element descriptions and populate desc2xy mapping
+        6. Convert computer calls from descriptions back to xy coordinates
+        7. Return output and usage
+        """
+        # Parse the composed model
+        if "+" not in model:
+            raise ValueError(f"Composed model must be in format 'grounding_model+thinking_model', got: {model}")
+        grounding_model, thinking_model = model.split("+", 1)
+        
+        pre_output_items = []
+        
+        # Step 0: Store last computer call image, if none then take a screenshot
+        last_image_b64 = get_last_computer_call_image(messages)
+        if last_image_b64 is None:
+            # Take a screenshot
+            screenshot_b64 = await computer_handler.screenshot() # type: ignore
+            if screenshot_b64:
+                
+                call_id = uuid.uuid4().hex
+                pre_output_items += [
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "output_text",
+                                "text": "Taking a screenshot to see the current computer screen."
+                            }
+                        ]
+                    },
+                    {
+                        "action": {
+                            "type": "screenshot"
+                        },
+                        "call_id": call_id,
+                        "status": "completed",
+                        "type": "computer_call"
+                    },
+                    {
+                        "type": "computer_call_output",
+                        "call_id": call_id,
+                        "output": {
+                            "type": "input_image",
+                            "image_url": f"data:image/png;base64,{screenshot_b64}"
+                        }
+                    },
+                ]
+                last_image_b64 = screenshot_b64
+                
+                # Call screenshot callback if provided
+                if _on_screenshot:
+                    await _on_screenshot(screenshot_b64)
+        
+        tool_schemas = _prepare_tools_for_grounded(tools) # type: ignore
+
+        # Step 1: Convert computer calls from xy to descriptions
+        input_messages = messages + pre_output_items
+        messages_with_descriptions = convert_computer_calls_xy2desc(input_messages, self.desc2xy)
+        
+        # Step 2: Convert responses items to completion messages
+        completion_messages = convert_responses_items_to_completion_messages(
+            messages_with_descriptions, 
+            allow_images_in_tool_results=False
+        )
+        
+        # Step 3: Call thinking model with litellm.acompletion
+        api_kwargs = {
+            "model": thinking_model,
+            "messages": completion_messages,
+            "tools": tool_schemas,
+            "max_retries": max_retries,
+            "stream": stream,
+            **kwargs
+        }
+
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+        
+        # Call API start hook
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+        
+        # Make the completion call
+        response = await litellm.acompletion(**api_kwargs)
+        
+        # Call API end hook
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+        
+        # Extract usage information
+        usage = {
+            **response.usage.model_dump(), # type: ignore
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+        
+        # Step 4: Convert completion messages back to responses items format
+        response_dict = response.model_dump() # type: ignore
+        choice_messages = [choice["message"] for choice in response_dict["choices"]]
+        thinking_output_items = []
+        
+        for choice_message in choice_messages:
+            thinking_output_items.extend(convert_completion_messages_to_responses_items([choice_message]))
+        
+        # Step 5: Get all element descriptions and populate desc2xy mapping
+        element_descriptions = get_all_element_descriptions(thinking_output_items)
+        
+        if element_descriptions and last_image_b64:
+            # Use grounding model to predict coordinates for each description
+            grounding_agent_conf = find_agent_config(grounding_model)
+            if grounding_agent_conf:
+                grounding_agent = grounding_agent_conf.agent_class()
+                
+                for desc in element_descriptions:
+                    coords = await grounding_agent.predict_click(
+                        model=grounding_model,
+                        image_b64=last_image_b64,
+                        instruction=desc
+                    )
+                    if coords:
+                        self.desc2xy[desc] = coords
+        
+        # Step 6: Convert computer calls from descriptions back to xy coordinates
+        final_output_items = convert_computer_calls_desc2xy(thinking_output_items, self.desc2xy)
+        
+        # Step 7: Return output and usage
+        return {
+            "output": pre_output_items + final_output_items,
+            "usage": usage
+        }
+    
+    async def predict_click(
+        self,
+        model: str,
+        image_b64: str,
+        instruction: str,
+        **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using the grounding model.
+        
+        For composed models, uses only the grounding model part for click prediction.
+        """
+        # Parse the composed model to get grounding model
+        if "+" not in model:
+            raise ValueError(f"Composed model must be in format 'grounding_model+thinking_model', got: {model}")
+        grounding_model, thinking_model = model.split("+", 1)
+        
+        # Find and use the grounding agent
+        grounding_agent_conf = find_agent_config(grounding_model)
+        if grounding_agent_conf:
+            grounding_agent = grounding_agent_conf.agent_class()
+            return await grounding_agent.predict_click(
+                model=grounding_model,
+                image_b64=image_b64,
+                instruction=instruction,
+                **kwargs
+            )
+        
+        return None
+    
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["click", "step"]

agent/loops/gta1.py

@@ -0,0 +1,178 @@
+"""
+GTA1 agent loop implementation for click prediction using litellm.acompletion
+Paper: https://arxiv.org/pdf/2507.05791
+Code: https://github.com/Yan98/GTA1
+"""
+
+import asyncio
+import json
+import re
+import base64
+from typing import Dict, List, Any, AsyncGenerator, Union, Optional, Tuple
+from io import BytesIO
+import uuid
+from PIL import Image
+import litellm
+import math
+
+from ..decorators import register_agent
+from ..types import Messages, AgentResponse, Tools, AgentCapability
+from ..loops.base import AsyncAgentConfig
+
+SYSTEM_PROMPT = '''
+You are an expert UI element locator. Given a GUI image and a user's element description, provide the coordinates of the specified element as a single (x,y) point. The image resolution is height {height} and width {width}. For elements with area, return the center point.
+
+Output the coordinate pair exactly:
+(x,y)
+'''.strip()
+
+def extract_coordinates(raw_string: str) -> Tuple[float, float]:
+    """Extract coordinates from model output."""
+    try:
+        matches = re.findall(r"\((-?\d*\.?\d+),\s*(-?\d*\.?\d+)\)", raw_string)
+        return tuple(map(float, matches[0])) # type: ignore
+    except:
+        return (0.0, 0.0)
+
+def smart_resize(height: int, width: int, factor: int = 28, min_pixels: int = 3136, max_pixels: int = 8847360) -> Tuple[int, int]:
+    """Smart resize function similar to qwen_vl_utils."""
+    # Calculate the total pixels
+    total_pixels = height * width
+    
+    # If already within bounds, return original dimensions
+    if min_pixels <= total_pixels <= max_pixels:
+        # Round to nearest factor
+        new_height = (height // factor) * factor
+        new_width = (width // factor) * factor
+        return new_height, new_width
+    
+    # Calculate scaling factor
+    if total_pixels > max_pixels:
+        scale = (max_pixels / total_pixels) ** 0.5
+    else:
+        scale = (min_pixels / total_pixels) ** 0.5
+    
+    # Apply scaling
+    new_height = int(height * scale)
+    new_width = int(width * scale)
+    
+    # Round to nearest factor
+    new_height = (new_height // factor) * factor
+    new_width = (new_width // factor) * factor
+    
+    # Ensure minimum size
+    new_height = max(new_height, factor)
+    new_width = max(new_width, factor)
+    
+    return new_height, new_width
+
+@register_agent(models=r".*GTA1.*")
+class GTA1Config(AsyncAgentConfig):
+    """GTA1 agent configuration implementing AsyncAgentConfig protocol for click prediction."""
+    
+    def __init__(self):
+        self.current_model = None
+        self.last_screenshot_b64 = None
+    
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        raise NotImplementedError()
+
+    async def predict_click(
+        self,
+        model: str,
+        image_b64: str,
+        instruction: str,
+        **kwargs
+    ) -> Optional[Tuple[float, float]]:
+        """
+        Predict click coordinates using GTA1 model via litellm.acompletion.
+        
+        Args:
+            model: The GTA1 model name
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+            
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        # Decode base64 image
+        image_data = base64.b64decode(image_b64)
+        image = Image.open(BytesIO(image_data))
+        width, height = image.width, image.height
+        
+        # Smart resize the image (similar to qwen_vl_utils)
+        resized_height, resized_width = smart_resize(
+            height, width, 
+            factor=28,  # Default factor for Qwen models
+            min_pixels=3136,
+            max_pixels=4096 * 2160
+        )
+        resized_image = image.resize((resized_width, resized_height))
+        scale_x, scale_y = width / resized_width, height / resized_height
+        
+        # Convert resized image back to base64
+        buffered = BytesIO()
+        resized_image.save(buffered, format="PNG")
+        resized_image_b64 = base64.b64encode(buffered.getvalue()).decode()
+        
+        # Prepare system and user messages
+        system_message = {
+            "role": "system",
+            "content": SYSTEM_PROMPT.format(height=resized_height, width=resized_width)
+        }
+        
+        user_message = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{resized_image_b64}"
+                    }
+                },
+                {
+                    "type": "text",
+                    "text": instruction
+                }
+            ]
+        }
+        
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "messages": [system_message, user_message],
+            "max_tokens": 32,
+            "temperature": 0.0,
+            **kwargs
+        }
+        
+        # Use liteLLM acompletion
+        response = await litellm.acompletion(**api_kwargs)
+        
+        # Extract response text
+        output_text = response.choices[0].message.content # type: ignore
+        
+        # Extract and rescale coordinates
+        pred_x, pred_y = extract_coordinates(output_text) # type: ignore
+        pred_x *= scale_x
+        pred_y *= scale_y
+        
+        return (math.floor(pred_x), math.floor(pred_y))
+    
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["click"]

agent/loops/model_types.csv

@@ -0,0 +1,6 @@
+model,predict_step,predict_point
+anthropic,✅,✅
+openai,✅,✅
+uitars,✅,✅
+omniparser,❌,✅
+gta1,❌,✅