cua-agent 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

agent/callbacks/trajectory_saver.py

@@ -0,0 +1,305 @@
+"""
+Trajectory saving callback handler for ComputerAgent.
+"""
+
+import os
+import json
+import uuid
+from datetime import datetime
+import base64
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Union, override
+from PIL import Image, ImageDraw
+import io
+from .base import AsyncCallbackHandler
+
+def sanitize_image_urls(data: Any) -> Any:
+    """
+    Recursively search for 'image_url' keys and set their values to '[omitted]'.
+    
+    Args:
+        data: Any data structure (dict, list, or primitive type)
+        
+    Returns:
+        A deep copy of the data with all 'image_url' values replaced with '[omitted]'
+    """
+    if isinstance(data, dict):
+        # Create a copy of the dictionary
+        sanitized = {}
+        for key, value in data.items():
+            if key == "image_url":
+                sanitized[key] = "[omitted]"
+            else:
+                # Recursively sanitize the value
+                sanitized[key] = sanitize_image_urls(value)
+        return sanitized
+    
+    elif isinstance(data, list):
+        # Recursively sanitize each item in the list
+        return [sanitize_image_urls(item) for item in data]
+    
+    else:
+        # For primitive types (str, int, bool, None, etc.), return as-is
+        return data
+
+
+class TrajectorySaverCallback(AsyncCallbackHandler):
+    """
+    Callback handler that saves agent trajectories to disk.
+    
+    Saves each run as a separate trajectory with unique ID, and each turn
+    within the trajectory gets its own folder with screenshots and responses.
+    """
+    
+    def __init__(self, trajectory_dir: str):
+        """
+        Initialize trajectory saver.
+        
+        Args:
+            trajectory_dir: Base directory to save trajectories
+        """
+        self.trajectory_dir = Path(trajectory_dir)
+        self.trajectory_id: Optional[str] = None
+        self.current_turn: int = 0
+        self.current_artifact: int = 0
+        self.model: Optional[str] = None
+        self.total_usage: Dict[str, Any] = {}
+        
+        # Ensure trajectory directory exists
+        self.trajectory_dir.mkdir(parents=True, exist_ok=True)
+    
+    def _get_turn_dir(self) -> Path:
+        """Get the directory for the current turn."""
+        if not self.trajectory_id:
+            raise ValueError("Trajectory not initialized - call _on_run_start first")
+        
+        # format: trajectory_id/turn_000
+        turn_dir = self.trajectory_dir / self.trajectory_id / f"turn_{self.current_turn:03d}"
+        turn_dir.mkdir(parents=True, exist_ok=True)
+        return turn_dir
+
+    def _save_artifact(self, name: str, artifact: Union[str, bytes, Dict[str, Any]]) -> None:
+        """Save an artifact to the current turn directory."""
+        turn_dir = self._get_turn_dir()
+        if isinstance(artifact, bytes):
+            # format: turn_000/0000_name.png
+            artifact_filename = f"{self.current_artifact:04d}_{name}"
+            artifact_path = turn_dir / f"{artifact_filename}.png"
+            with open(artifact_path, "wb") as f:
+                f.write(artifact)
+        else:
+            # format: turn_000/0000_name.json
+            artifact_filename = f"{self.current_artifact:04d}_{name}"
+            artifact_path = turn_dir / f"{artifact_filename}.json"
+            with open(artifact_path, "w") as f:
+                json.dump(sanitize_image_urls(artifact), f, indent=2)
+        self.current_artifact += 1
+
+    def _update_usage(self, usage: Dict[str, Any]) -> None:
+        """Update total usage statistics."""
+        def add_dicts(target: Dict[str, Any], source: Dict[str, Any]) -> None:
+            for key, value in source.items():
+                if isinstance(value, dict):
+                    if key not in target:
+                        target[key] = {}
+                    add_dicts(target[key], value)
+                else:
+                    if key not in target:
+                        target[key] = 0
+                    target[key] += value
+        add_dicts(self.total_usage, usage)
+    
+    @override
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Initialize trajectory tracking for a new run."""
+        model = kwargs.get("model", "unknown")
+        model_name_short = model.split("+")[-1].split("/")[-1].lower()[:16]
+        if "+" in model:
+            model_name_short = model.split("+")[0].lower()[:4] + "_" + model_name_short
+
+        # id format: yyyy-mm-dd_model_hhmmss_uuid[:4]
+        now = datetime.now()
+        self.trajectory_id = f"{now.strftime('%Y-%m-%d')}_{model_name_short}_{now.strftime('%H%M%S')}_{str(uuid.uuid4())[:4]}"
+        self.current_turn = 0
+        self.current_artifact = 0
+        self.model = model
+        self.total_usage = {}
+        
+        # Create trajectory directory
+        trajectory_path = self.trajectory_dir / self.trajectory_id
+        trajectory_path.mkdir(parents=True, exist_ok=True)
+        
+        # Save trajectory metadata
+        metadata = {
+            "trajectory_id": self.trajectory_id,
+            "created_at": str(uuid.uuid1().time),
+            "status": "running",
+            "kwargs": kwargs,
+        }
+        
+        with open(trajectory_path / "metadata.json", "w") as f:
+            json.dump(metadata, f, indent=2)
+
+    @override
+    async def on_run_end(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]], new_items: List[Dict[str, Any]]) -> None:
+        """Finalize run tracking by updating metadata with completion status, usage, and new items."""
+        if not self.trajectory_id:
+            return
+        
+        # Update metadata with completion status, total usage, and new items
+        trajectory_path = self.trajectory_dir / self.trajectory_id
+        metadata_path = trajectory_path / "metadata.json"
+        
+        # Read existing metadata
+        if metadata_path.exists():
+            with open(metadata_path, "r") as f:
+                metadata = json.load(f)
+        else:
+            metadata = {}
+        
+        # Update metadata with completion info
+        metadata.update({
+            "status": "completed",
+            "completed_at": str(uuid.uuid1().time),
+            "total_usage": self.total_usage,
+            "new_items": sanitize_image_urls(new_items),
+            "total_turns": self.current_turn
+        })
+        
+        # Save updated metadata
+        with open(metadata_path, "w") as f:
+            json.dump(metadata, f, indent=2)
+    
+    @override 
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        if not self.trajectory_id:
+            return
+        
+        self._save_artifact("api_start", { "kwargs": kwargs })
+    
+    @override
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """Save API call result."""
+        if not self.trajectory_id:
+            return
+        
+        self._save_artifact("api_result", { "kwargs": kwargs, "result": result })
+
+    @override
+    async def on_screenshot(self, screenshot: Union[str, bytes], name: str = "screenshot") -> None:
+        """Save a screenshot."""
+        if isinstance(screenshot, str):
+            screenshot = base64.b64decode(screenshot)
+        self._save_artifact(name, screenshot)
+
+    @override
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        self._update_usage(usage)
+
+    @override
+    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
+        """Save responses to the current turn directory and update usage statistics."""
+        if not self.trajectory_id:
+            return
+        
+        # Save responses
+        turn_dir = self._get_turn_dir()
+        response_data = {
+            "timestamp": str(uuid.uuid1().time),
+            "model": self.model,
+            "kwargs": kwargs,
+            "response": responses
+        }
+        
+        self._save_artifact("agent_response", response_data)
+        
+        # Increment turn counter
+        self.current_turn += 1
+
+    def _draw_crosshair_on_image(self, image_bytes: bytes, x: int, y: int) -> bytes:
+        """
+        Draw a red dot and crosshair at the specified coordinates on the image.
+        
+        Args:
+            image_bytes: The original image as bytes
+            x: X coordinate for the crosshair
+            y: Y coordinate for the crosshair
+            
+        Returns:
+            Modified image as bytes with red dot and crosshair
+        """
+        # Open the image
+        image = Image.open(io.BytesIO(image_bytes))
+        draw = ImageDraw.Draw(image)
+        
+        # Draw crosshair lines (red, 2px thick)
+        crosshair_size = 20
+        line_width = 2
+        color = "red"
+        
+        # Horizontal line
+        draw.line([(x - crosshair_size, y), (x + crosshair_size, y)], fill=color, width=line_width)
+        # Vertical line
+        draw.line([(x, y - crosshair_size), (x, y + crosshair_size)], fill=color, width=line_width)
+        
+        # Draw center dot (filled circle)
+        dot_radius = 3
+        draw.ellipse([(x - dot_radius, y - dot_radius), (x + dot_radius, y + dot_radius)], fill=color)
+        
+        # Convert back to bytes
+        output = io.BytesIO()
+        image.save(output, format='PNG')
+        return output.getvalue()
+
+    @override
+    async def on_computer_call_end(self, item: Dict[str, Any], result: List[Dict[str, Any]]) -> None:
+        """
+        Called when a computer call has completed.
+        Saves screenshots and computer call output.
+        """
+        if not self.trajectory_id:
+            return
+        
+        self._save_artifact("computer_call_result", { "item": item, "result": result })
+        
+        # Check if action has x/y coordinates and there's a screenshot in the result
+        action = item.get("action", {})
+        if "x" in action and "y" in action:
+            # Look for screenshot in the result
+            for result_item in result:
+                if (result_item.get("type") == "computer_call_output" and 
+                    result_item.get("output", {}).get("type") == "input_image"):
+                    
+                    image_url = result_item["output"]["image_url"]
+                
+                    # Extract base64 image data
+                    if image_url.startswith("data:image/"):
+                        # Format: data:image/png;base64,<base64_data>
+                        base64_data = image_url.split(",", 1)[1]
+                    else:
+                        # Assume it's just base64 data
+                        base64_data = image_url
+                    
+                    try:
+                        # Decode the image
+                        image_bytes = base64.b64decode(base64_data)
+                        
+                        # Draw crosshair at the action coordinates
+                        annotated_image = self._draw_crosshair_on_image(
+                            image_bytes, 
+                            int(action["x"]), 
+                            int(action["y"])
+                        )
+                        
+                        # Save as screenshot_action
+                        self._save_artifact("screenshot_action", annotated_image)
+                        
+                    except Exception as e:
+                        # If annotation fails, just log and continue
+                        print(f"Failed to annotate screenshot: {e}")
+                    
+                    break  # Only process the first screenshot found
+
+        # Increment turn counter
+        self.current_turn += 1

agent/cli.py

@@ -0,0 +1,297 @@
+"""
+CLI chat interface for agent - Computer Use Agent
+
+Usage:
+    python -m agent.cli <model_string>
+    
+Examples:
+    python -m agent.cli openai/computer-use-preview
+    python -m agent.cli anthropic/claude-3-5-sonnet-20241022
+    python -m agent.cli omniparser+anthropic/claude-3-5-sonnet-20241022
+"""
+
+try:
+    import asyncio
+    import argparse
+    import os
+    import sys
+    import json
+    from typing import List, Dict, Any
+    import dotenv
+    from yaspin import yaspin
+except ImportError:
+    if __name__ == "__main__":
+        raise ImportError(
+            "CLI dependencies not found. "
+            "Please install with: pip install \"cua-agent[cli]\""
+        )
+
+# Load environment variables
+dotenv.load_dotenv()
+
+# Color codes for terminal output
+class Colors:
+    RESET = '\033[0m'
+    BOLD = '\033[1m'
+    DIM = '\033[2m'
+    
+    # Text colors
+    RED = '\033[31m'
+    GREEN = '\033[32m'
+    YELLOW = '\033[33m'
+    BLUE = '\033[34m'
+    MAGENTA = '\033[35m'
+    CYAN = '\033[36m'
+    WHITE = '\033[37m'
+    GRAY = '\033[90m'
+    
+    # Background colors
+    BG_RED = '\033[41m'
+    BG_GREEN = '\033[42m'
+    BG_YELLOW = '\033[43m'
+    BG_BLUE = '\033[44m'
+
+
+def print_colored(text: str, color: str = "", bold: bool = False, dim: bool = False, end: str = "\n"):
+    """Print colored text to terminal."""
+    prefix = ""
+    if bold:
+        prefix += Colors.BOLD
+    if dim:
+        prefix += Colors.DIM
+    if color:
+        prefix += color
+    
+    print(f"{prefix}{text}{Colors.RESET}", end=end)
+
+
+def print_action(action_type: str, details: Dict[str, Any]):
+    """Print computer action with nice formatting."""
+    # Format action details
+    args_str = ""
+    if action_type == "click" and "x" in details and "y" in details:
+        args_str = f"({details['x']}, {details['y']})"
+    elif action_type == "type" and "text" in details:
+        text = details["text"]
+        if len(text) > 50:
+            text = text[:47] + "..."
+        args_str = f'"{text}"'
+    elif action_type == "key" and "key" in details:
+        args_str = f"'{details['key']}'"
+    elif action_type == "scroll" and "x" in details and "y" in details:
+        args_str = f"({details['x']}, {details['y']})"
+    
+    print_colored(f"🛠️  {action_type}{args_str}", dim=True)
+
+
+def print_welcome(model: str, agent_loop: str, container_name: str):
+    """Print welcome message."""
+    print_colored(f"Connected to {container_name} ({model}, {agent_loop})")
+    print_colored("Type 'exit' to quit.", dim=True)
+
+async def ainput(prompt: str = ""):
+    return await asyncio.to_thread(input, prompt)
+
+async def chat_loop(agent, model: str, container_name: str):
+    """Main chat loop with the agent."""
+    print_welcome(model, agent.agent_loop.__name__, container_name)
+    
+    history = []
+    
+    while True:
+        # Get user input with prompt
+        print_colored("> ", end="")
+        user_input = await ainput()
+        
+        if user_input.lower() in ['exit', 'quit', 'q']:
+            print_colored("\n👋 Goodbye!")
+            break
+            
+        if not user_input:
+            continue
+            
+        # Add user message to history
+        history.append({"role": "user", "content": user_input})
+        
+        # Stream responses from the agent with spinner
+        with yaspin(text="Thinking...", spinner="line", attrs=["dark"]) as spinner:
+            spinner.hide()
+            
+            async for result in agent.run(history):
+                # Add agent responses to history
+                history.extend(result.get("output", []))
+                
+                # Process and display the output
+                for item in result.get("output", []):
+                    if item.get("type") == "message":
+                        # Display agent text response
+                        content = item.get("content", [])
+                        for content_part in content:
+                            if content_part.get("text"):
+                                text = content_part.get("text", "").strip()
+                                if text:
+                                    spinner.hide()
+                                    print_colored(text)
+                    
+                    elif item.get("type") == "computer_call":
+                        # Display computer action
+                        action = item.get("action", {})
+                        action_type = action.get("type", "")
+                        if action_type:
+                            spinner.hide()
+                            print_action(action_type, action)
+                            spinner.text = f"Performing {action_type}..."
+                            spinner.show()
+                    
+                    elif item.get("type") == "function_call":
+                        # Display function call
+                        function_name = item.get("name", "")
+                        spinner.hide()
+                        print_colored(f"🔧 Calling function: {function_name}", dim=True)
+                        spinner.text = f"Calling {function_name}..."
+                        spinner.show()
+                    
+                    elif item.get("type") == "function_call_output":
+                        # Display function output (dimmed)
+                        output = item.get("output", "")
+                        if output and len(output.strip()) > 0:
+                            spinner.hide()
+                            print_colored(f"📤 {output}", dim=True)
+            
+            spinner.hide()
+        
+
+async def main():
+    """Main CLI function."""
+    parser = argparse.ArgumentParser(
+        description="CUA Agent CLI - Interactive computer use assistant",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python -m agent.cli openai/computer-use-preview
+  python -m agent.cli anthropic/claude-3-5-sonnet-20241022
+  python -m agent.cli omniparser+anthropic/claude-3-5-sonnet-20241022
+  python -m agent.cli huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B
+        """
+    )
+    
+    parser.add_argument(
+        "model",
+        help="Model string (e.g., 'openai/computer-use-preview', 'anthropic/claude-3-5-sonnet-20241022')"
+    )
+    
+    parser.add_argument(
+        "--images",
+        type=int,
+        default=3,
+        help="Number of recent images to keep in context (default: 3)"
+    )
+    
+    parser.add_argument(
+        "--trajectory",
+        action="store_true",
+        help="Save trajectory for debugging"
+    )
+    
+    parser.add_argument(
+        "--budget",
+        type=float,
+        help="Maximum budget for the session (in dollars)"
+    )
+    
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Enable verbose logging"
+    )
+    
+    args = parser.parse_args()
+    
+    # Check for required environment variables
+    container_name = os.getenv("CUA_CONTAINER_NAME")
+    cua_api_key = os.getenv("CUA_API_KEY")
+    
+    # Prompt for missing environment variables
+    if not container_name:
+        print_colored("CUA_CONTAINER_NAME not set.", dim=True)
+        print_colored("You can get a CUA container at https://www.trycua.com/", dim=True)
+        container_name = input("Enter your CUA container name: ").strip()
+        if not container_name:
+            print_colored("❌ Container name is required.")
+            sys.exit(1)
+    
+    if not cua_api_key:
+        print_colored("CUA_API_KEY not set.", dim=True)
+        cua_api_key = input("Enter your CUA API key: ").strip()
+        if not cua_api_key:
+            print_colored("❌ API key is required.")
+            sys.exit(1)
+    
+    # Check for provider-specific API keys based on model
+    provider_api_keys = {
+        "openai/": "OPENAI_API_KEY",
+        "anthropic/": "ANTHROPIC_API_KEY",
+        "omniparser+": "OPENAI_API_KEY",
+        "omniparser+": "ANTHROPIC_API_KEY",
+    }
+    
+    # Find matching provider and check for API key
+    for prefix, env_var in provider_api_keys.items():
+        if args.model.startswith(prefix):
+            if not os.getenv(env_var):
+                print_colored(f"{env_var} not set.", dim=True)
+                api_key = input(f"Enter your {env_var.replace('_', ' ').title()}: ").strip()
+                if not api_key:
+                    print_colored(f"❌ {env_var.replace('_', ' ').title()} is required.")
+                    sys.exit(1)
+                # Set the environment variable for the session
+                os.environ[env_var] = api_key
+            break
+    
+    # Import here to avoid import errors if dependencies are missing
+    try:
+        from agent import ComputerAgent
+        from computer import Computer
+    except ImportError as e:
+        print_colored(f"❌ Import error: {e}", Colors.RED, bold=True)
+        print_colored("Make sure agent and computer libraries are installed.", Colors.YELLOW)
+        sys.exit(1)
+    
+    # Create computer instance
+    async with Computer(
+        os_type="linux",
+        provider_type="cloud",
+        name=container_name,
+        api_key=cua_api_key
+    ) as computer:
+        
+        # Create agent
+        agent_kwargs = {
+            "model": args.model,
+            "tools": [computer],
+            "only_n_most_recent_images": args.images,
+            "verbosity": 20 if args.verbose else 30,  # DEBUG vs WARNING
+        }
+        
+        if args.trajectory:
+            agent_kwargs["trajectory_dir"] = "trajectories"
+        
+        if args.budget:
+            agent_kwargs["max_trajectory_budget"] = {
+                "max_budget": args.budget,
+                "raise_error": True,
+                "reset_after_each_run": False
+            }
+        
+        agent = ComputerAgent(**agent_kwargs)
+        
+        # Start chat loop
+        await chat_loop(agent, args.model, container_name)
+
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except (KeyboardInterrupt, EOFError) as _:
+        print_colored("\n\n👋 Goodbye!")