cua-agent 0.4.17__tar.gz → 0.4.19__tar.gz

{cua_agent-0.4.17 → cua_agent-0.4.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.4.17
+Version: 0.4.19
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: >=3.12
@@ -38,7 +38,7 @@ Requires-Dist: python-dotenv>=1.0.1; extra == "ui"
 Provides-Extra: cli
 Requires-Dist: yaspin>=3.1.0; extra == "cli"
 Provides-Extra: hud
-Requires-Dist: hud-python==0.2.10; extra == "hud"
+Requires-Dist: hud-python<0.5.0,>=0.4.12; extra == "hud"
 Provides-Extra: all
 Requires-Dist: ultralytics>=8.0.0; extra == "all"
 Requires-Dist: cua-som<0.2.0,>=0.1.0; extra == "all"
@@ -49,7 +49,7 @@ Requires-Dist: transformers>=4.54.0; extra == "all"
 Requires-Dist: gradio>=5.23.3; extra == "all"
 Requires-Dist: python-dotenv>=1.0.1; extra == "all"
 Requires-Dist: yaspin>=3.1.0; extra == "all"
-Requires-Dist: hud-python==0.2.10; extra == "all"
+Requires-Dist: hud-python<0.5.0,>=0.4.12; extra == "all"
 Description-Content-Type: text/markdown
 
 <div align="center">

{cua_agent-0.4.17 → cua_agent-0.4.19}/agent/adapters/__init__.py

@@ -4,8 +4,10 @@ Adapters package for agent - Custom LLM adapters for LiteLLM
 
 from .huggingfacelocal_adapter import HuggingFaceLocalAdapter
 from .human_adapter import HumanAdapter
+from .mlxvlm_adapter import MLXVLMAdapter
 
 __all__ = [
     "HuggingFaceLocalAdapter",
     "HumanAdapter",
+    "MLXVLMAdapter",
 ]

cua_agent-0.4.19/agent/adapters/mlxvlm_adapter.py

@@ -0,0 +1,359 @@
+import asyncio
+import functools
+import warnings
+import io
+import base64
+import math
+import re
+from concurrent.futures import ThreadPoolExecutor
+from typing import Iterator, AsyncIterator, Dict, List, Any, Optional, Tuple, cast
+from PIL import Image
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+from litellm.llms.custom_llm import CustomLLM
+from litellm import completion, acompletion
+
+# Try to import MLX dependencies
+try:
+    import mlx.core as mx
+    from mlx_vlm import load, generate
+    from mlx_vlm.prompt_utils import apply_chat_template
+    from mlx_vlm.utils import load_config
+    from transformers.tokenization_utils import PreTrainedTokenizer
+    MLX_AVAILABLE = True
+except ImportError:
+    MLX_AVAILABLE = False
+
+# Constants for smart_resize
+IMAGE_FACTOR = 28
+MIN_PIXELS = 100 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+def round_by_factor(number: float, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+def ceil_by_factor(number: float, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+def floor_by_factor(number: float, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+def smart_resize(
+    height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+class MLXVLMAdapter(CustomLLM):
+    """MLX VLM Adapter for running vision-language models locally using MLX."""
+    
+    def __init__(self, **kwargs):
+        """Initialize the adapter.
+        
+        Args:
+            **kwargs: Additional arguments
+        """
+        super().__init__()
+        
+        self.models = {}  # Cache for loaded models
+        self.processors = {}  # Cache for loaded processors
+        self.configs = {}  # Cache for loaded configs
+        self._executor = ThreadPoolExecutor(max_workers=1)  # Single thread pool
+        
+    def _load_model_and_processor(self, model_name: str):
+        """Load model and processor if not already cached.
+        
+        Args:
+            model_name: Name of the model to load
+            
+        Returns:
+            Tuple of (model, processor, config)
+        """
+        if not MLX_AVAILABLE:
+            raise ImportError("MLX VLM dependencies not available. Please install mlx-vlm.")
+        
+        if model_name not in self.models:
+            # Load model and processor
+            model_obj, processor = load(
+                model_name, 
+                processor_kwargs={"min_pixels": MIN_PIXELS, "max_pixels": MAX_PIXELS}
+            )
+            config = load_config(model_name)
+            
+            # Cache them
+            self.models[model_name] = model_obj
+            self.processors[model_name] = processor
+            self.configs[model_name] = config
+            
+        return self.models[model_name], self.processors[model_name], self.configs[model_name]
+    
+    def _process_coordinates(self, text: str, original_size: Tuple[int, int], model_size: Tuple[int, int]) -> str:
+        """Process coordinates in box tokens based on image resizing using smart_resize approach.
+        
+        Args:
+            text: Text containing box tokens
+            original_size: Original image size (width, height)
+            model_size: Model processed image size (width, height)
+            
+        Returns:
+            Text with processed coordinates
+        """
+        # Find all box tokens
+        box_pattern = r"<\|box_start\|>\((\d+),\s*(\d+)\)<\|box_end\|>"
+        
+        def process_coords(match):
+            model_x, model_y = int(match.group(1)), int(match.group(2))
+            # Scale coordinates from model space to original image space
+            # Both original_size and model_size are in (width, height) format
+            new_x = int(model_x * original_size[0] / model_size[0])  # Width
+            new_y = int(model_y * original_size[1] / model_size[1])  # Height
+            return f"<|box_start|>({new_x},{new_y})<|box_end|>"
+        
+        return re.sub(box_pattern, process_coords, text)
+    
+    def _convert_messages(self, messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Image.Image], Dict[int, Tuple[int, int]], Dict[int, Tuple[int, int]]]:
+        """Convert OpenAI format messages to MLX VLM format and extract images.
+        
+        Args:
+            messages: Messages in OpenAI format
+            
+        Returns:
+            Tuple of (processed_messages, images, original_sizes, model_sizes)
+        """
+        processed_messages = []
+        images = []
+        original_sizes = {}  # Track original sizes of images for coordinate mapping
+        model_sizes = {}  # Track model processed sizes
+        image_index = 0
+        
+        for message in messages:
+            processed_message = {
+                "role": message["role"],
+                "content": []
+            }
+            
+            content = message.get("content", [])
+            if isinstance(content, str):
+                # Simple text content
+                processed_message["content"] = content
+            elif isinstance(content, list):
+                # Multi-modal content
+                processed_content = []
+                for item in content:
+                    if item.get("type") == "text":
+                        processed_content.append({
+                            "type": "text",
+                            "text": item.get("text", "")
+                        })
+                    elif item.get("type") == "image_url":
+                        image_url = item.get("image_url", {}).get("url", "")
+                        pil_image = None
+                        
+                        if image_url.startswith("data:image/"):
+                            # Extract base64 data
+                            base64_data = image_url.split(',')[1]
+                            # Convert base64 to PIL Image
+                            image_data = base64.b64decode(base64_data)
+                            pil_image = Image.open(io.BytesIO(image_data))
+                        else:
+                            # Handle file path or URL
+                            pil_image = Image.open(image_url)
+                        
+                        # Store original image size for coordinate mapping
+                        original_size = pil_image.size
+                        original_sizes[image_index] = original_size
+                        
+                        # Use smart_resize to determine model size
+                        # Note: smart_resize expects (height, width) but PIL gives (width, height)
+                        height, width = original_size[1], original_size[0]
+                        new_height, new_width = smart_resize(height, width)
+                        # Store model size in (width, height) format for consistent coordinate processing
+                        model_sizes[image_index] = (new_width, new_height)
+                        
+                        # Resize the image using the calculated dimensions from smart_resize
+                        resized_image = pil_image.resize((new_width, new_height))
+                        images.append(resized_image)
+                        
+                        # Add image placeholder to content
+                        processed_content.append({
+                            "type": "image"
+                        })
+                        
+                        image_index += 1
+                
+                processed_message["content"] = processed_content
+            
+            processed_messages.append(processed_message)
+        
+        return processed_messages, images, original_sizes, model_sizes
+    
+    def _generate(self, **kwargs) -> str:
+        """Generate response using the local MLX VLM model.
+        
+        Args:
+            **kwargs: Keyword arguments containing messages and model info
+            
+        Returns:
+            Generated text response
+        """
+        messages = kwargs.get('messages', [])
+        model_name = kwargs.get('model', 'mlx-community/UI-TARS-1.5-7B-4bit')
+        max_tokens = kwargs.get('max_tokens', 128)
+        
+        # Warn about ignored kwargs
+        ignored_kwargs = set(kwargs.keys()) - {'messages', 'model', 'max_tokens'}
+        if ignored_kwargs:
+            warnings.warn(f"Ignoring unsupported kwargs: {ignored_kwargs}")
+        
+        # Load model and processor
+        model, processor, config = self._load_model_and_processor(model_name)
+        
+        # Convert messages and extract images
+        processed_messages, images, original_sizes, model_sizes = self._convert_messages(messages)
+        
+        # Process user text input with box coordinates after image processing
+        # Swap original_size and model_size arguments for inverse transformation
+        for msg_idx, msg in enumerate(processed_messages):
+            if msg.get("role") == "user" and isinstance(msg.get("content"), str):
+                content = msg.get("content", "")
+                if "<|box_start|>" in content and original_sizes and model_sizes and 0 in original_sizes and 0 in model_sizes:
+                    orig_size = original_sizes[0]
+                    model_size = model_sizes[0]
+                    # Swap arguments to perform inverse transformation for user input
+                    processed_messages[msg_idx]["content"] = self._process_coordinates(content, model_size, orig_size)
+        
+        try:
+            # Format prompt according to model requirements using the processor directly
+            prompt = processor.apply_chat_template(
+                processed_messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                return_tensors='pt'
+            )
+            tokenizer = cast(PreTrainedTokenizer, processor)
+            
+            # Generate response
+            text_content, usage = generate(
+                model, 
+                tokenizer, 
+                str(prompt), 
+                images, # type: ignore
+                verbose=False,
+                max_tokens=max_tokens
+            )
+            
+        except Exception as e:
+            raise RuntimeError(f"Error generating response: {str(e)}") from e
+        
+        # Process coordinates in the response back to original image space
+        if original_sizes and model_sizes and 0 in original_sizes and 0 in model_sizes:
+            # Get original image size and model size (using the first image)
+            orig_size = original_sizes[0]
+            model_size = model_sizes[0]
+            
+            # Check if output contains box tokens that need processing
+            if "<|box_start|>" in text_content:
+                # Process coordinates from model space back to original image space
+                text_content = self._process_coordinates(text_content, orig_size, model_size)
+        
+        return text_content
+    
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        """Synchronous completion method.
+        
+        Returns:
+            ModelResponse with generated text
+        """
+        generated_text = self._generate(**kwargs)
+        
+        result = completion(
+            model=f"mlx/{kwargs.get('model', 'mlx-community/UI-TARS-1.5-7B-4bit')}",
+            mock_response=generated_text,
+        )
+        return cast(ModelResponse, result)
+    
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        """Asynchronous completion method.
+        
+        Returns:
+            ModelResponse with generated text
+        """
+        # Run _generate in thread pool to avoid blocking
+        loop = asyncio.get_event_loop()
+        generated_text = await loop.run_in_executor(
+            self._executor, 
+            functools.partial(self._generate, **kwargs)
+        )
+        
+        result = await acompletion(
+            model=f"mlx/{kwargs.get('model', 'mlx-community/UI-TARS-1.5-7B-4bit')}",
+            mock_response=generated_text,
+        )
+        return cast(ModelResponse, result)
+    
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        """Synchronous streaming method.
+        
+        Returns:
+            Iterator of GenericStreamingChunk
+        """
+        generated_text = self._generate(**kwargs)
+        
+        generic_streaming_chunk: GenericStreamingChunk = {
+            "finish_reason": "stop",
+            "index": 0,
+            "is_finished": True,
+            "text": generated_text,
+            "tool_use": None,
+            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+        }
+        
+        yield generic_streaming_chunk
+    
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        """Asynchronous streaming method.
+        
+        Returns:
+            AsyncIterator of GenericStreamingChunk
+        """
+        # Run _generate in thread pool to avoid blocking
+        loop = asyncio.get_event_loop()
+        generated_text = await loop.run_in_executor(
+            self._executor, 
+            functools.partial(self._generate, **kwargs)
+        )
+        
+        generic_streaming_chunk: GenericStreamingChunk = {
+            "finish_reason": "stop",
+            "index": 0,
+            "is_finished": True,
+            "text": generated_text,
+            "tool_use": None,
+            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+        }
+        
+        yield generic_streaming_chunk

{cua_agent-0.4.17 → cua_agent-0.4.19}/agent/agent.py

@@ -3,6 +3,7 @@ ComputerAgent - Main agent class that selects and runs agent loops
 """
 
 import asyncio
+from pathlib import Path
 from typing import Dict, List, Any, Optional, AsyncGenerator, Union, cast, Callable, Set, Tuple
 
 from litellm.responses.utils import Usage
@@ -22,6 +23,7 @@ import inspect
 from .adapters import (
     HuggingFaceLocalAdapter,
     HumanAdapter,
+    MLXVLMAdapter,
 )
 from .callbacks import (
     ImageRetentionCallback, 
@@ -29,6 +31,7 @@ from .callbacks import (
     TrajectorySaverCallback, 
     BudgetManagerCallback,
     TelemetryCallback,
+    OperatorNormalizerCallback
 )
 from .computers import (
     AsyncComputerHandler,
@@ -160,7 +163,7 @@ class ComputerAgent:
         only_n_most_recent_images: Optional[int] = None,
         callbacks: Optional[List[Any]] = None,
         verbosity: Optional[int] = None,
-        trajectory_dir: Optional[str] = None,
+        trajectory_dir: Optional[str | Path | dict] = None,
         max_retries: Optional[int] = 3,
         screenshot_delay: Optional[float | int] = 0.5,
         use_prompt_caching: Optional[bool] = False,
@@ -201,6 +204,9 @@ class ComputerAgent:
 
         # == Add built-in callbacks ==
 
+        # Prepend operator normalizer callback
+        self.callbacks.insert(0, OperatorNormalizerCallback())
+
         # Add telemetry callback if telemetry_enabled is set
         if self.telemetry_enabled:
             if isinstance(self.telemetry_enabled, bool):
@@ -218,7 +224,10 @@ class ComputerAgent:
         
         # Add trajectory saver callback if trajectory_dir is set
         if self.trajectory_dir:
-            self.callbacks.append(TrajectorySaverCallback(self.trajectory_dir))
+            if isinstance(self.trajectory_dir, dict):
+                self.callbacks.append(TrajectorySaverCallback(**self.trajectory_dir))
+            elif isinstance(self.trajectory_dir, (str, Path)):
+                self.callbacks.append(TrajectorySaverCallback(str(self.trajectory_dir)))
         
         # Add budget manager if max_trajectory_budget is set
         if max_trajectory_budget:
@@ -234,9 +243,11 @@ class ComputerAgent:
             device="auto"
         )
         human_adapter = HumanAdapter()
+        mlx_adapter = MLXVLMAdapter()
         litellm.custom_provider_map = [
             {"provider": "huggingface-local", "custom_handler": hf_adapter},
-            {"provider": "human", "custom_handler": human_adapter}
+            {"provider": "human", "custom_handler": human_adapter},
+            {"provider": "mlx", "custom_handler": mlx_adapter}
         ]
         litellm.suppress_debug_info = True
 

{cua_agent-0.4.17 → cua_agent-0.4.19}/agent/callbacks/__init__.py

@@ -8,6 +8,7 @@ from .logging import LoggingCallback
 from .trajectory_saver import TrajectorySaverCallback
 from .budget_manager import BudgetManagerCallback
 from .telemetry import TelemetryCallback
+from .operator_validator import OperatorNormalizerCallback
 
 __all__ = [
     "AsyncCallbackHandler",
@@ -16,4 +17,5 @@ __all__ = [
     "TrajectorySaverCallback",
     "BudgetManagerCallback",
     "TelemetryCallback",
+    "OperatorNormalizerCallback",
 ]

cua_agent-0.4.19/agent/callbacks/operator_validator.py

@@ -0,0 +1,138 @@
+"""
+OperatorValidatorCallback
+
+Ensures agent output actions conform to expected schemas by fixing common issues:
+- click: add default button='left' if missing
+- keypress: wrap keys string into a list
+- etc.
+
+This runs in on_llm_end, which receives the output array (AgentMessage[] as dicts).
+The purpose is to avoid spending another LLM call to fix broken computer call syntax when possible.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+from .base import AsyncCallbackHandler
+
+
+class OperatorNormalizerCallback(AsyncCallbackHandler):
+    """Normalizes common computer call hallucinations / errors in computer call syntax."""
+
+    async def on_llm_end(self, output: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        # Mutate in-place as requested, but still return the list for chaining
+        for item in output or []:
+            if item.get("type") != "computer_call":
+                continue
+            action = item.get("action")
+            if not isinstance(action, dict):
+                continue
+
+            # rename mouse click actions to "click"
+            for mouse_btn in ["left", "right", "wheel", "back", "forward"]:
+                if action.get("type", "") == f"{mouse_btn}_click":
+                    action["type"] = "click"
+                    action["button"] = mouse_btn
+            # rename hotkey actions to "keypress"
+            for alias in ["hotkey", "key", "press", "key_press"]:
+                if action.get("type", "") == alias:
+                    action["type"] = "keypress"
+            # assume click actions
+            if "button" in action and "type" not in action:
+                action["type"] = "click"
+            if "click" in action and "type" not in action:
+                action["type"] = "click"
+            if ("scroll_x" in action or "scroll_y" in action) and "type" not in action:
+                action["type"] = "scroll"
+            if "text" in action and "type" not in action:
+                action["type"] = "type"
+
+            action_type = action.get("type")
+            def _keep_keys(action: Dict[str, Any], keys_to_keep: List[str]):
+                """Keep only the provided keys on action; delete everything else.
+                Always ensures required 'type' is present if listed in keys_to_keep.
+                """
+                for key in list(action.keys()):
+                    if key not in keys_to_keep:
+                        del action[key]
+            # rename "coordinate" to "x", "y"
+            if "coordinate" in action:
+                action["x"] = action["coordinate"][0]
+                action["y"] = action["coordinate"][1]
+                del action["coordinate"]
+            if action_type == "click":
+                # convert "click" to "button"
+                if "button" not in action and "click" in action:
+                    action["button"] = action["click"]
+                    del action["click"]
+                # default button to "left"
+                action["button"] = action.get("button", "left")
+            # add default scroll x, y if missing
+            if action_type == "scroll":
+                action["scroll_x"] = action.get("scroll_x", 0)
+                action["scroll_y"] = action.get("scroll_y", 0)
+            # ensure keys arg is a list (normalize aliases first)
+            if action_type == "keypress":
+                keys = action.get("keys")
+                for keys_alias in ["keypress", "key", "press", "key_press", "text"]:
+                    if keys_alias in action:
+                        action["keys"] = action[keys_alias]
+                        del action[keys_alias]
+                keys = action.get("keys")
+                if isinstance(keys, str):
+                    action["keys"] = keys.replace("-", "+").split("+") if len(keys) > 1 else [keys]
+            required_keys_by_type = {
+                # OpenAI actions
+                "click": ["type", "button", "x", "y"],
+                "double_click": ["type", "x", "y"],
+                "drag": ["type", "path"],
+                "keypress": ["type", "keys"],
+                "move": ["type", "x", "y"],
+                "screenshot": ["type"],
+                "scroll": ["type", "scroll_x", "scroll_y", "x", "y"],
+                "type": ["type", "text"],
+                "wait": ["type"],
+                # Anthropic actions
+                "left_mouse_down": ["type", "x", "y"],
+                "left_mouse_up": ["type", "x", "y"],
+                "triple_click": ["type", "button", "x", "y"],
+            }
+            keep = required_keys_by_type.get(action_type or "")
+            if keep:
+                _keep_keys(action, keep)
+            
+
+        # Second pass: if an assistant message is immediately followed by a computer_call,
+        # replace the assistant message itself with a reasoning message with summary text.
+        if isinstance(output, list):
+            for i, item in enumerate(output):
+                # AssistantMessage shape: { type: 'message', role: 'assistant', content: OutputContent[] }
+                if item.get("type") == "message" and item.get("role") == "assistant":
+                    next_idx = i + 1
+                    if next_idx >= len(output):
+                        continue
+                    next_item = output[next_idx]
+                    if not isinstance(next_item, dict):
+                        continue
+                    if next_item.get("type") != "computer_call":
+                        continue
+                    contents = item.get("content") or []
+                    # Extract text from OutputContent[]
+                    text_parts: List[str] = []
+                    if isinstance(contents, list):
+                        for c in contents:
+                            if isinstance(c, dict) and c.get("type") == "output_text" and isinstance(c.get("text"), str):
+                                text_parts.append(c["text"])
+                    text_content = "\n".join(text_parts).strip()
+                    # Replace assistant message with reasoning message
+                    output[i] = {
+                        "type": "reasoning",
+                        "summary": [
+                            {
+                                "type": "summary_text",
+                                "text": text_content,
+                            }
+                        ],
+                    }
+
+        return output