cua-agent 0.1.0__py3-none-any.whl

agent/providers/omni/tools/computer.py

@@ -0,0 +1,216 @@
+"""Provider-agnostic implementation of the ComputerTool."""
+
+import logging
+import base64
+import io
+from typing import Any, Dict
+
+from PIL import Image
+from computer.computer import Computer
+
+from ....core.tools.computer import BaseComputerTool
+from ....core.tools import ToolResult, ToolError
+
+
+class OmniComputerTool(BaseComputerTool):
+    """A provider-agnostic implementation of the computer tool."""
+
+    name = "computer"
+    logger = logging.getLogger(__name__)
+
+    def __init__(self, computer: Computer):
+        """Initialize the ComputerTool.
+
+        Args:
+            computer: Computer instance for screen interactions
+        """
+        super().__init__(computer)
+        # Initialize dimensions to None, will be set in initialize_dimensions
+        self.width = None
+        self.height = None
+        self.display_num = None
+
+    def to_params(self) -> Dict[str, Any]:
+        """Convert tool to provider-agnostic parameters.
+
+        Returns:
+            Dictionary with tool parameters
+        """
+        return {
+            "name": self.name,
+            "description": "A tool that allows the agent to interact with the screen, keyboard, and mouse",
+            "parameters": {
+                "action": {
+                    "type": "string",
+                    "enum": [
+                        "key",
+                        "type",
+                        "mouse_move",
+                        "left_click",
+                        "left_click_drag",
+                        "right_click",
+                        "middle_click",
+                        "double_click",
+                        "screenshot",
+                        "cursor_position",
+                        "scroll",
+                    ],
+                    "description": "The action to perform on the computer",
+                },
+                "text": {
+                    "type": "string",
+                    "description": "Text to type or key to press, required for 'key' and 'type' actions",
+                },
+                "coordinate": {
+                    "type": "array",
+                    "items": {"type": "integer"},
+                    "description": "X,Y coordinates for mouse actions like click and move",
+                },
+                "direction": {
+                    "type": "string",
+                    "enum": ["up", "down"],
+                    "description": "Direction to scroll, used with the 'scroll' action",
+                },
+                "amount": {
+                    "type": "integer",
+                    "description": "Amount to scroll, used with the 'scroll' action",
+                },
+            },
+            **self.options,
+        }
+
+    async def __call__(self, **kwargs) -> ToolResult:
+        """Execute the computer tool with the provided arguments.
+
+        Args:
+            action: The action to perform
+            text: Text to type or key to press (for key/type actions)
+            coordinate: X,Y coordinates (for mouse actions)
+            direction: Direction to scroll (for scroll action)
+            amount: Amount to scroll (for scroll action)
+
+        Returns:
+            ToolResult with the action output and optional screenshot
+        """
+        # Ensure dimensions are initialized
+        if self.width is None or self.height is None:
+            await self.initialize_dimensions()
+
+        action = kwargs.get("action")
+        text = kwargs.get("text")
+        coordinate = kwargs.get("coordinate")
+        direction = kwargs.get("direction", "down")
+        amount = kwargs.get("amount", 10)
+
+        self.logger.info(f"Executing computer action: {action}")
+
+        try:
+            if action == "screenshot":
+                return await self.screenshot()
+            elif action == "left_click" and coordinate:
+                x, y = coordinate
+                self.logger.info(f"Clicking at ({x}, {y})")
+                await self.computer.interface.move_cursor(x, y)
+                await self.computer.interface.left_click()
+
+                # Take screenshot after action
+                screenshot = await self.computer.interface.screenshot()
+                screenshot = await self.resize_screenshot_if_needed(screenshot)
+                return ToolResult(
+                    output=f"Performed left click at ({x}, {y})",
+                    base64_image=base64.b64encode(screenshot).decode(),
+                )
+            elif action == "right_click" and coordinate:
+                x, y = coordinate
+                self.logger.info(f"Right clicking at ({x}, {y})")
+                await self.computer.interface.move_cursor(x, y)
+                await self.computer.interface.right_click()
+
+                # Take screenshot after action
+                screenshot = await self.computer.interface.screenshot()
+                screenshot = await self.resize_screenshot_if_needed(screenshot)
+                return ToolResult(
+                    output=f"Performed right click at ({x}, {y})",
+                    base64_image=base64.b64encode(screenshot).decode(),
+                )
+            elif action == "double_click" and coordinate:
+                x, y = coordinate
+                self.logger.info(f"Double clicking at ({x}, {y})")
+                await self.computer.interface.move_cursor(x, y)
+                await self.computer.interface.double_click()
+
+                # Take screenshot after action
+                screenshot = await self.computer.interface.screenshot()
+                screenshot = await self.resize_screenshot_if_needed(screenshot)
+                return ToolResult(
+                    output=f"Performed double click at ({x}, {y})",
+                    base64_image=base64.b64encode(screenshot).decode(),
+                )
+            elif action == "mouse_move" and coordinate:
+                x, y = coordinate
+                self.logger.info(f"Moving cursor to ({x}, {y})")
+                await self.computer.interface.move_cursor(x, y)
+
+                # Take screenshot after action
+                screenshot = await self.computer.interface.screenshot()
+                screenshot = await self.resize_screenshot_if_needed(screenshot)
+                return ToolResult(
+                    output=f"Moved cursor to ({x}, {y})",
+                    base64_image=base64.b64encode(screenshot).decode(),
+                )
+            elif action == "type" and text:
+                self.logger.info(f"Typing text: {text}")
+                await self.computer.interface.type_text(text)
+
+                # Take screenshot after action
+                screenshot = await self.computer.interface.screenshot()
+                screenshot = await self.resize_screenshot_if_needed(screenshot)
+                return ToolResult(
+                    output=f"Typed text: {text}",
+                    base64_image=base64.b64encode(screenshot).decode(),
+                )
+            elif action == "key" and text:
+                self.logger.info(f"Pressing key: {text}")
+
+                # Handle special key combinations
+                if "+" in text:
+                    keys = text.split("+")
+                    await self.computer.interface.hotkey(*keys)
+                else:
+                    await self.computer.interface.press(text)
+
+                # Take screenshot after action
+                screenshot = await self.computer.interface.screenshot()
+                screenshot = await self.resize_screenshot_if_needed(screenshot)
+                return ToolResult(
+                    output=f"Pressed key: {text}",
+                    base64_image=base64.b64encode(screenshot).decode(),
+                )
+            elif action == "cursor_position":
+                pos = await self.computer.interface.get_cursor_position()
+                return ToolResult(output=f"X={int(pos[0])},Y={int(pos[1])}")
+            elif action == "scroll":
+                if direction == "down":
+                    self.logger.info(f"Scrolling down, amount: {amount}")
+                    for _ in range(amount):
+                        await self.computer.interface.hotkey("fn", "down")
+                else:
+                    self.logger.info(f"Scrolling up, amount: {amount}")
+                    for _ in range(amount):
+                        await self.computer.interface.hotkey("fn", "up")
+
+                # Take screenshot after action
+                screenshot = await self.computer.interface.screenshot()
+                screenshot = await self.resize_screenshot_if_needed(screenshot)
+                return ToolResult(
+                    output=f"Scrolled {direction} by {amount} steps",
+                    base64_image=base64.b64encode(screenshot).decode(),
+                )
+
+            # Default to screenshot for unimplemented actions
+            self.logger.warning(f"Action {action} not fully implemented, taking screenshot")
+            return await self.screenshot()
+
+        except Exception as e:
+            self.logger.error(f"Error during computer action: {str(e)}")
+            return ToolResult(error=f"Failed to perform {action}: {str(e)}")

agent/providers/omni/tools/manager.py

@@ -0,0 +1,83 @@
+"""Omni tool manager implementation."""
+
+from typing import Dict, List, Any
+from enum import Enum
+
+from computer.computer import Computer
+
+from ....core.tools import BaseToolManager
+from ....core.tools.collection import ToolCollection
+
+from .bash import OmniBashTool
+from .computer import OmniComputerTool
+from .edit import OmniEditTool
+
+
+class ProviderType(Enum):
+    """Supported provider types."""
+
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    CLAUDE = "claude"  # Alias for Anthropic
+    GPT = "gpt"  # Alias for OpenAI
+
+
+class OmniToolManager(BaseToolManager):
+    """Tool manager for multi-provider support."""
+
+    def __init__(self, computer: Computer):
+        """Initialize Omni tool manager.
+
+        Args:
+            computer: Computer instance for tools
+        """
+        super().__init__(computer)
+        # Initialize tools
+        self.computer_tool = OmniComputerTool(self.computer)
+        self.bash_tool = OmniBashTool(self.computer)
+        self.edit_tool = OmniEditTool(self.computer)
+
+    def _initialize_tools(self) -> ToolCollection:
+        """Initialize all available tools."""
+        return ToolCollection(self.computer_tool, self.bash_tool, self.edit_tool)
+
+    async def _initialize_tools_specific(self) -> None:
+        """Initialize provider-specific tool requirements."""
+        await self.computer_tool.initialize_dimensions()
+
+    def get_tool_params(self) -> List[Dict[str, Any]]:
+        """Get tool parameters for API calls.
+
+        Returns:
+            List of tool parameters in default format
+        """
+        if self.tools is None:
+            raise RuntimeError("Tools not initialized. Call initialize() first.")
+        return self.tools.to_params()
+
+    def get_provider_tools(self, provider: ProviderType) -> List[Dict[str, Any]]:
+        """Get tools formatted for a specific provider.
+
+        Args:
+            provider: Provider type to format tools for
+
+        Returns:
+            List of tool parameters in provider-specific format
+        """
+        if self.tools is None:
+            raise RuntimeError("Tools not initialized. Call initialize() first.")
+
+        # Default is the base implementation
+        tools = self.tools.to_params()
+
+        # Customize for each provider if needed
+        if provider in [ProviderType.ANTHROPIC, ProviderType.CLAUDE]:
+            # Format for Anthropic API
+            # Additional adjustments can be made here
+            pass
+        elif provider in [ProviderType.OPENAI, ProviderType.GPT]:
+            # Format for OpenAI API
+            # Future implementation
+            pass
+
+        return tools

agent/providers/omni/types.py

@@ -0,0 +1,30 @@
+"""Type definitions for the Omni provider."""
+
+from enum import StrEnum
+from typing import Dict
+
+
+class APIProvider(StrEnum):
+    """Supported API providers."""
+
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    GROQ = "groq"
+    QWEN = "qwen"
+
+
+# Default models for each provider
+PROVIDER_TO_DEFAULT_MODEL: Dict[APIProvider, str] = {
+    APIProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+    APIProvider.OPENAI: "gpt-4o",
+    APIProvider.GROQ: "deepseek-r1-distill-llama-70b",
+    APIProvider.QWEN: "qwen2.5-vl-72b-instruct",
+}
+
+# Environment variable names for each provider
+PROVIDER_TO_ENV_VAR: Dict[APIProvider, str] = {
+    APIProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
+    APIProvider.OPENAI: "OPENAI_API_KEY",
+    APIProvider.GROQ: "GROQ_API_KEY",
+    APIProvider.QWEN: "QWEN_API_KEY",
+}

agent/providers/omni/utils.py

@@ -0,0 +1,155 @@
+"""Utility functions for Omni provider."""
+
+import base64
+import io
+import logging
+from typing import Tuple
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+
+def compress_image_base64(
+    base64_str: str, max_size_bytes: int = 5 * 1024 * 1024, quality: int = 90
+) -> tuple[str, str]:
+    """Compress a base64 encoded image to ensure it's below a certain size.
+
+    Args:
+        base64_str: Base64 encoded image string (with or without data URL prefix)
+        max_size_bytes: Maximum size in bytes (default: 5MB)
+        quality: Initial JPEG quality (0-100)
+
+    Returns:
+        tuple[str, str]: (Compressed base64 encoded image, media_type)
+    """
+    # Handle data URL prefix if present (e.g., "data:image/png;base64,...")
+    original_prefix = ""
+    media_type = "image/png"  # Default media type
+
+    if base64_str.startswith("data:"):
+        parts = base64_str.split(",", 1)
+        if len(parts) == 2:
+            original_prefix = parts[0] + ","
+            base64_str = parts[1]
+            # Try to extract media type from the prefix
+            if "image/jpeg" in original_prefix.lower():
+                media_type = "image/jpeg"
+            elif "image/png" in original_prefix.lower():
+                media_type = "image/png"
+
+    # Check if the base64 string is small enough already
+    if len(base64_str) <= max_size_bytes:
+        logger.info(f"Image already within size limit: {len(base64_str)} bytes")
+        return original_prefix + base64_str, media_type
+
+    try:
+        # Decode base64
+        img_data = base64.b64decode(base64_str)
+        img_size = len(img_data)
+        logger.info(f"Original image size: {img_size} bytes")
+
+        # Open image
+        img = Image.open(io.BytesIO(img_data))
+
+        # First, try to compress as PNG (maintains transparency if present)
+        buffer = io.BytesIO()
+        img.save(buffer, format="PNG", optimize=True)
+        buffer.seek(0)
+        compressed_data = buffer.getvalue()
+        compressed_b64 = base64.b64encode(compressed_data).decode("utf-8")
+
+        if len(compressed_b64) <= max_size_bytes:
+            logger.info(f"Compressed to {len(compressed_data)} bytes as PNG")
+            return compressed_b64, "image/png"
+
+        # Strategy 1: Try reducing quality with JPEG format
+        current_quality = quality
+        while current_quality > 20:
+            buffer = io.BytesIO()
+            # Convert to RGB if image has alpha channel (JPEG doesn't support transparency)
+            if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info):
+                logger.info("Converting transparent image to RGB for JPEG compression")
+                rgb_img = Image.new("RGB", img.size, (255, 255, 255))
+                rgb_img.paste(img, mask=img.split()[3] if img.mode == "RGBA" else None)
+                rgb_img.save(buffer, format="JPEG", quality=current_quality, optimize=True)
+            else:
+                img.save(buffer, format="JPEG", quality=current_quality, optimize=True)
+
+            buffer.seek(0)
+            compressed_data = buffer.getvalue()
+            compressed_b64 = base64.b64encode(compressed_data).decode("utf-8")
+
+            if len(compressed_b64) <= max_size_bytes:
+                logger.info(
+                    f"Compressed to {len(compressed_data)} bytes with JPEG quality {current_quality}"
+                )
+                return compressed_b64, "image/jpeg"
+
+            # Reduce quality and try again
+            current_quality -= 10
+
+        # Strategy 2: If quality reduction isn't enough, reduce dimensions
+        scale_factor = 0.8
+        current_img = img
+
+        while scale_factor > 0.3:
+            # Resize image
+            new_width = int(img.width * scale_factor)
+            new_height = int(img.height * scale_factor)
+            current_img = img.resize((new_width, new_height), Image.LANCZOS)
+
+            # Try with reduced size and quality
+            buffer = io.BytesIO()
+            # Convert to RGB if necessary for JPEG
+            if current_img.mode in ("RGBA", "LA") or (
+                current_img.mode == "P" and "transparency" in current_img.info
+            ):
+                rgb_img = Image.new("RGB", current_img.size, (255, 255, 255))
+                rgb_img.paste(
+                    current_img, mask=current_img.split()[3] if current_img.mode == "RGBA" else None
+                )
+                rgb_img.save(buffer, format="JPEG", quality=70, optimize=True)
+            else:
+                current_img.save(buffer, format="JPEG", quality=70, optimize=True)
+
+            buffer.seek(0)
+            compressed_data = buffer.getvalue()
+            compressed_b64 = base64.b64encode(compressed_data).decode("utf-8")
+
+            if len(compressed_b64) <= max_size_bytes:
+                logger.info(
+                    f"Compressed to {len(compressed_data)} bytes with scale {scale_factor} and JPEG quality 70"
+                )
+                return compressed_b64, "image/jpeg"
+
+            # Reduce scale factor and try again
+            scale_factor -= 0.1
+
+        # If we get here, we couldn't compress enough
+        logger.warning("Could not compress image below required size with quality preservation")
+
+        # Last resort: Use minimum quality and size
+        buffer = io.BytesIO()
+        smallest_img = img.resize((int(img.width * 0.5), int(img.height * 0.5)), Image.LANCZOS)
+        # Convert to RGB if necessary
+        if smallest_img.mode in ("RGBA", "LA") or (
+            smallest_img.mode == "P" and "transparency" in smallest_img.info
+        ):
+            rgb_img = Image.new("RGB", smallest_img.size, (255, 255, 255))
+            rgb_img.paste(
+                smallest_img, mask=smallest_img.split()[3] if smallest_img.mode == "RGBA" else None
+            )
+            rgb_img.save(buffer, format="JPEG", quality=20, optimize=True)
+        else:
+            smallest_img.save(buffer, format="JPEG", quality=20, optimize=True)
+
+        buffer.seek(0)
+        final_data = buffer.getvalue()
+        final_b64 = base64.b64encode(final_data).decode("utf-8")
+
+        logger.warning(f"Final compressed size: {len(final_b64)} bytes (may still exceed limit)")
+        return final_b64, "image/jpeg"
+
+    except Exception as e:
+        logger.error(f"Error compressing image: {str(e)}")
+        raise

agent/providers/omni/visualization.py

@@ -0,0 +1,130 @@
+"""Visualization utilities for the Cua provider."""
+
+import base64
+import logging
+from io import BytesIO
+from typing import Tuple
+from PIL import Image, ImageDraw
+
+logger = logging.getLogger(__name__)
+
+
+def visualize_click(x: int, y: int, img_base64: str) -> Image.Image:
+    """Visualize a click action by drawing on the screenshot.
+
+    Args:
+        x: X coordinate of the click
+        y: Y coordinate of the click
+        img_base64: Base64 encoded image to draw on
+
+    Returns:
+        PIL Image with visualization
+    """
+    try:
+        # Decode the base64 image
+        img_data = base64.b64decode(img_base64)
+        img = Image.open(BytesIO(img_data))
+
+        # Create a drawing context
+        draw = ImageDraw.Draw(img)
+
+        # Draw concentric circles at the click position
+        small_radius = 10
+        large_radius = 30
+
+        # Draw filled inner circle
+        draw.ellipse(
+            [(x - small_radius, y - small_radius), (x + small_radius, y + small_radius)],
+            fill="red",
+        )
+
+        # Draw outlined outer circle
+        draw.ellipse(
+            [(x - large_radius, y - large_radius), (x + large_radius, y + large_radius)],
+            outline="red",
+            width=3,
+        )
+
+        return img
+
+    except Exception as e:
+        logger.error(f"Error visualizing click: {str(e)}")
+        # Return a blank image in case of error
+        return Image.new("RGB", (800, 600), color="white")
+
+
+def visualize_scroll(direction: str, clicks: int, img_base64: str) -> Image.Image:
+    """Visualize a scroll action by drawing arrows on the screenshot.
+
+    Args:
+        direction: 'up' or 'down'
+        clicks: Number of scroll clicks
+        img_base64: Base64 encoded image to draw on
+
+    Returns:
+        PIL Image with visualization
+    """
+    try:
+        # Decode the base64 image
+        img_data = base64.b64decode(img_base64)
+        img = Image.open(BytesIO(img_data))
+
+        # Get image dimensions
+        width, height = img.size
+
+        # Create a drawing context
+        draw = ImageDraw.Draw(img)
+
+        # Determine arrow direction and positions
+        center_x = width // 2
+        arrow_width = 100
+
+        if direction.lower() == "up":
+            # Draw up arrow in the middle of the screen
+            arrow_y = height // 2
+            # Arrow points
+            points = [
+                (center_x, arrow_y - 50),  # Top point
+                (center_x - arrow_width // 2, arrow_y + 50),  # Bottom left
+                (center_x + arrow_width // 2, arrow_y + 50),  # Bottom right
+            ]
+            color = "blue"
+        else:  # down
+            # Draw down arrow in the middle of the screen
+            arrow_y = height // 2
+            # Arrow points
+            points = [
+                (center_x, arrow_y + 50),  # Bottom point
+                (center_x - arrow_width // 2, arrow_y - 50),  # Top left
+                (center_x + arrow_width // 2, arrow_y - 50),  # Top right
+            ]
+            color = "green"
+
+        # Draw filled arrow
+        draw.polygon(points, fill=color)
+
+        # Add text showing number of clicks
+        text_y = arrow_y + 70 if direction.lower() == "down" else arrow_y - 70
+        draw.text((center_x - 40, text_y), f"{clicks} clicks", fill="black")
+
+        return img
+
+    except Exception as e:
+        logger.error(f"Error visualizing scroll: {str(e)}")
+        # Return a blank image in case of error
+        return Image.new("RGB", (800, 600), color="white")
+
+
+def calculate_element_center(box: Tuple[int, int, int, int]) -> Tuple[int, int]:
+    """Calculate the center coordinates of a bounding box.
+
+    Args:
+        box: Tuple of (left, top, right, bottom) coordinates
+
+    Returns:
+        Tuple of (center_x, center_y) coordinates
+    """
+    left, top, right, bottom = box
+    center_x = (left + right) // 2
+    center_y = (top + bottom) // 2
+    return center_x, center_y

agent/types/__init__.py

@@ -0,0 +1,26 @@
+"""Type definitions for the agent package."""
+
+from .base import Provider, HostConfig, TaskResult, Annotation
+from .messages import Message, Request, Response, StepMessage, DisengageMessage
+from .tools import ToolInvocation, ToolInvocationState, ClientAttachment, ToolResult
+
+__all__ = [
+    # Base types
+    "Provider",
+    "HostConfig",
+    "TaskResult",
+    "Annotation",
+    
+    # Message types
+    "Message",
+    "Request",
+    "Response",
+    "StepMessage",
+    "DisengageMessage",
+    
+    # Tool types
+    "ToolInvocation",
+    "ToolInvocationState",
+    "ClientAttachment",
+    "ToolResult",
+]

agent/types/base.py

@@ -0,0 +1,52 @@
+"""Base type definitions."""
+
+from enum import Enum, auto
+from typing import Dict, Any
+from pydantic import BaseModel, ConfigDict
+
+
+class Provider(str, Enum):
+    """Available AI providers."""
+
+    UNKNOWN = "unknown"  # Default provider for base class
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    OLLAMA = "ollama"
+    OMNI = "omni"
+    GROQ = "groq"
+
+
+class HostConfig(BaseModel):
+    """Host configuration."""
+
+    model_config = ConfigDict(extra="forbid")
+    hostname: str
+    port: int
+
+    @property
+    def address(self) -> str:
+        return f"{self.hostname}:{self.port}"
+
+
+class TaskResult(BaseModel):
+    """Result of a task execution."""
+
+    model_config = ConfigDict(extra="forbid")
+    result: str
+    vnc_password: str
+
+
+class Annotation(BaseModel):
+    """Annotation metadata."""
+
+    model_config = ConfigDict(extra="forbid")
+    id: str
+    vm_url: str
+
+
+class AgenticLoop(Enum):
+    """Enumeration of available loop types."""
+
+    ANTHROPIC = auto()  # Anthropic implementation
+    OMNI = auto()  # OmniLoop implementation
+    # Add more loop types as needed