cua-agent 0.1.0__py3-none-any.whl

agent/README.md

@@ -0,0 +1,63 @@
+# Agent Package Structure
+
+## Overview
+The agent package provides a modular and extensible framework for AI-powered computer agents.
+
+## Directory Structure
+```
+agent/
+├── __init__.py           # Package exports
+├── core/                 # Core functionality
+│   ├── __init__.py
+│   ├── computer_agent.py # Main entry point
+│   └── factory.py        # Provider factory
+├── base/                 # Base implementations
+│   ├── __init__.py
+│   ├── agent.py         # Base agent class
+│   ├── core/            # Core components
+│   │   ├── callbacks.py
+│   │   ├── loop.py
+│   │   └── messages.py
+│   └── tools/           # Tool implementations
+├── providers/           # Provider implementations
+│   ├── __init__.py
+│   ├── anthropic/      # Anthropic provider
+│   │   ├── agent.py
+│   │   ├── loop.py
+│   │   └── tool_manager.py
+│   └── omni/           # Omni provider
+│       ├── agent.py
+│       ├── loop.py
+│       └── tool_manager.py
+└── types/              # Type definitions
+    ├── __init__.py
+    ├── base.py        # Core types
+    ├── messages.py    # Message types
+    ├── tools.py       # Tool types
+    └── providers/     # Provider-specific types
+        ├── anthropic.py
+        └── omni.py
+```
+
+## Key Components
+
+### Core
+- `computer_agent.py`: Main entry point for creating and using agents
+- `factory.py`: Factory for creating provider-specific implementations
+
+### Base
+- `agent.py`: Base agent implementation with shared functionality
+- `core/`: Core components used across providers
+- `tools/`: Shared tool implementations
+
+### Providers
+Each provider follows the same structure:
+- `agent.py`: Provider-specific agent implementation
+- `loop.py`: Provider-specific message loop
+- `tool_manager.py`: Tool management for provider
+
+### Types
+- `base.py`: Core type definitions
+- `messages.py`: Message-related types
+- `tools.py`: Tool-related types
+- `providers/`: Provider-specific type definitions

agent/__init__.py

@@ -0,0 +1,10 @@
+"""CUA (Computer Use) Agent for AI-driven computer interaction."""
+
+__version__ = "0.1.0"
+
+from .core.factory import AgentFactory
+from .core.agent import ComputerAgent
+from .types.base import Provider, AgenticLoop
+from .providers.omni.types import APIProvider
+
+__all__ = ["AgentFactory", "Provider", "ComputerAgent", "AgenticLoop", "APIProvider"]

agent/core/README.md

@@ -0,0 +1,101 @@
+# Unified ComputerAgent
+
+The `ComputerAgent` class provides a unified implementation that consolidates the previously separate agent implementations (AnthropicComputerAgent and OmniComputerAgent) into a single, configurable class.
+
+## Features
+
+- **Multiple Loop Types**: Switch between different agentic loop implementations using the `loop_type` parameter (Anthropic or Omni).
+- **Provider Support**: Use different AI providers (OpenAI, Anthropic, etc.) with the appropriate loop.
+- **Trajectory Saving**: Control whether to save screenshots and logs with the `save_trajectory` parameter.
+- **Consistent Interface**: Maintains a consistent interface regardless of the underlying loop implementation.
+
+## API Key Requirements
+
+To use the ComputerAgent, you'll need API keys for the providers you want to use:
+
+- For **OpenAI**: Set the `OPENAI_API_KEY` environment variable or pass it directly as `api_key`.
+- For **Anthropic**: Set the `ANTHROPIC_API_KEY` environment variable or pass it directly as `api_key`.
+- For **Groq**: Set the `GROQ_API_KEY` environment variable or pass it directly as `api_key`.
+
+You can set environment variables in several ways:
+
+```bash
+# In your terminal before running the code
+export OPENAI_API_KEY=your_api_key_here
+
+# Or in a .env file
+OPENAI_API_KEY=your_api_key_here
+```
+
+## Usage
+
+Here's how to use the unified ComputerAgent:
+
+```python
+from agent.core.agent import ComputerAgent
+from agent.types.base import AgenticLoop
+from agent.providers.omni.types import APIProvider
+from computer import Computer
+
+# Create a Computer instance
+computer = Computer()
+
+# Create an agent with the OMNI loop and OpenAI provider
+agent = ComputerAgent(
+    computer=computer,
+    loop_type=AgenticLoop.OMNI,
+    provider=APIProvider.OPENAI,
+    model="gpt-4o",
+    api_key="your_api_key_here",  # Can also use OPENAI_API_KEY environment variable
+    save_trajectory=True,
+    only_n_most_recent_images=5
+)
+
+# Create an agent with the ANTHROPIC loop
+agent = ComputerAgent(
+    computer=computer,
+    loop_type=AgenticLoop.ANTHROPIC,
+    model="claude-3-7-sonnet-20250219",
+    api_key="your_api_key_here",  # Can also use ANTHROPIC_API_KEY environment variable
+    save_trajectory=True,
+    only_n_most_recent_images=5
+)
+
+# Use the agent
+async with agent:
+    async for result in agent.run("Your task description here"):
+        # Process the result
+        title = result["metadata"].get("title", "Screen Analysis")
+        content = result["content"]
+        print(f"\n{title}")
+        print(content)
+```
+
+## Parameters
+
+- `computer`: Computer instance to control
+- `loop_type`: The type of loop to use (AgenticLoop.ANTHROPIC or AgenticLoop.OMNI)
+- `provider`: AI provider to use (required for Omni loop)
+- `api_key`: Optional API key (will use environment variable if not provided)
+- `model`: Optional model name (will use provider default if not specified)
+- `save_trajectory`: Whether to save screenshots and logs
+- `only_n_most_recent_images`: Only keep N most recent images
+- `max_retries`: Maximum number of retry attempts
+
+## Directory Structure
+
+When `save_trajectory` is enabled, the agent will create the following directory structure:
+
+```
+experiments/
+  ├── screenshots/   # Screenshots captured during agent execution
+  └── logs/          # API call logs and other logging information
+```
+
+## Extending with New Loop Types
+
+To add a new loop type:
+
+1. Implement a new loop class
+2. Add a new value to the `AgenticLoop` enum
+3. Update the `_initialize_loop` method in `ComputerAgent` to handle the new loop type 

agent/core/__init__.py

@@ -0,0 +1,34 @@
+"""Core agent components."""
+
+from .base_agent import BaseComputerAgent
+from .loop import BaseLoop
+from .messages import (
+    create_user_message,
+    create_assistant_message,
+    create_system_message,
+    create_image_message,
+    create_screen_message,
+    BaseMessageManager,
+    ImageRetentionConfig,
+)
+from .callbacks import (
+    CallbackManager, 
+    CallbackHandler,
+    BaseCallbackManager,
+    ContentCallback,
+    ToolCallback,
+    APICallback,
+)
+
+__all__ = [
+    "BaseComputerAgent", 
+    "BaseLoop", 
+    "CallbackManager", 
+    "CallbackHandler",
+    "BaseMessageManager",
+    "ImageRetentionConfig",
+    "BaseCallbackManager",
+    "ContentCallback",
+    "ToolCallback",
+    "APICallback",
+]

agent/core/agent.py

@@ -0,0 +1,284 @@
+"""Unified computer agent implementation that supports multiple loops."""
+
+import os
+import logging
+import asyncio
+from typing import Any, AsyncGenerator, Dict, List, Optional, TYPE_CHECKING
+from datetime import datetime
+
+from computer import Computer
+
+from ..types.base import Provider, AgenticLoop
+from .base_agent import BaseComputerAgent
+
+# Only import types for type checking to avoid circular imports
+if TYPE_CHECKING:
+    from ..providers.anthropic.loop import AnthropicLoop
+    from ..providers.omni.loop import OmniLoop
+    from ..providers.omni.parser import OmniParser
+
+# Import the APIProvider enum without importing the whole module
+from ..providers.omni.types import APIProvider
+
+logger = logging.getLogger(__name__)
+
+# Default models for different providers
+DEFAULT_MODELS = {
+    APIProvider.OPENAI: "gpt-4o",
+    APIProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+    APIProvider.GROQ: "llama3-70b-8192",
+}
+
+# Map providers to their environment variable names
+ENV_VARS = {
+    APIProvider.OPENAI: "OPENAI_API_KEY",
+    APIProvider.GROQ: "GROQ_API_KEY",
+    APIProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
+}
+
+
+class ComputerAgent(BaseComputerAgent):
+    """Unified implementation of the computer agent supporting multiple loop types.
+
+    This class consolidates the previous AnthropicComputerAgent and OmniComputerAgent
+    into a single implementation with configurable loop type.
+    """
+
+    def __init__(
+        self,
+        computer: Computer,
+        loop_type: AgenticLoop = AgenticLoop.OMNI,
+        ai_provider: APIProvider = APIProvider.OPENAI,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        save_trajectory: bool = True,
+        trajectory_dir: Optional[str] = "trajectories",
+        only_n_most_recent_images: Optional[int] = None,
+        max_retries: int = 3,
+        verbosity: int = logging.INFO,
+        **kwargs,
+    ):
+        """Initialize the computer agent.
+
+        Args:
+            computer: Computer instance to control
+            loop_type: The type of loop to use (Anthropic or Omni)
+            ai_provider: AI provider to use (required for Cua loop)
+            api_key: Optional API key (will use environment variable if not provided)
+            model: Optional model name (will use provider default if not specified)
+            save_trajectory: Whether to save screenshots and logs
+            trajectory_dir: Directory to save trajectories (defaults to "trajectories")
+            only_n_most_recent_images: Limit history to N most recent images
+            max_retries: Maximum number of retry attempts for failed operations
+            verbosity: Logging level (standard Python logging levels: logging.DEBUG, logging.INFO, etc.)
+            **kwargs: Additional keyword arguments to pass to the loop
+        """
+        # Set up trajectory directories based on save_trajectory
+        base_dir = trajectory_dir if save_trajectory else None
+        # Don't create a redundant screenshots directory - directly use the timestamp folder
+        screenshot_dir = None  # This was previously set to os.path.join(base_dir, "screenshots")
+        log_dir = None
+
+        super().__init__(
+            max_retries=max_retries,
+            computer=computer,
+            screenshot_dir=screenshot_dir,
+            log_dir=log_dir,
+            **kwargs,
+        )
+
+        self.loop_type = loop_type
+        self.provider = ai_provider
+        self.save_trajectory = save_trajectory
+        self.trajectory_dir = trajectory_dir
+        self.only_n_most_recent_images = only_n_most_recent_images
+        self.verbosity = verbosity
+        self._kwargs = kwargs  # Keep this for loop initialization
+
+        # Configure logging based on verbosity
+        self._configure_logging(verbosity)
+
+        # Get API key from environment if not provided
+        if api_key is None:
+            env_var = (
+                ENV_VARS.get(ai_provider) if loop_type == AgenticLoop.OMNI else "ANTHROPIC_API_KEY"
+            )
+            if not env_var:
+                raise ValueError(
+                    f"Unsupported provider: {ai_provider}. Please use one of: {list(ENV_VARS.keys())}"
+                )
+
+            api_key = os.environ.get(env_var)
+            if not api_key:
+                raise ValueError(
+                    f"No API key provided and {env_var} environment variable is not set.\n"
+                    f"Please set the {env_var} environment variable or pass the api_key directly:\n"
+                    f"  - Export in terminal: export {env_var}=your_api_key_here\n"
+                    f"  - Add to .env file: {env_var}=your_api_key_here\n"
+                    f"  - Pass directly: api_key='your_api_key_here'"
+                )
+        self.api_key = api_key
+
+        # Set model based on provider if not specified
+        if model is None:
+            if loop_type == AgenticLoop.OMNI:
+                self.model = DEFAULT_MODELS[ai_provider]
+            else:  # Anthropic loop
+                self.model = DEFAULT_MODELS[APIProvider.ANTHROPIC]
+        else:
+            self.model = model
+
+        # Initialize the appropriate loop based on loop_type
+        self.loop = self._init_loop()
+
+    def _configure_logging(self, verbosity: int):
+        """Configure logging based on verbosity level."""
+        # Use the logging level directly without mapping
+        logger.setLevel(verbosity)
+        logging.getLogger("agent").setLevel(verbosity)
+
+        # Log the verbosity level that was set
+        if verbosity <= logging.DEBUG:
+            logger.info("Agent logging set to DEBUG level (full debug information)")
+        elif verbosity <= logging.INFO:
+            logger.info("Agent logging set to INFO level (standard output)")
+        elif verbosity <= logging.WARNING:
+            logger.warning("Agent logging set to WARNING level (warnings and errors only)")
+        elif verbosity <= logging.ERROR:
+            logger.warning("Agent logging set to ERROR level (errors only)")
+        elif verbosity <= logging.CRITICAL:
+            logger.warning("Agent logging set to CRITICAL level (critical errors only)")
+
+    def _init_loop(self) -> Any:
+        """Initialize the loop based on the loop_type.
+
+        Returns:
+            Initialized loop instance
+        """
+        # Lazy import OmniLoop and OmniParser to avoid circular imports
+        from ..providers.omni.loop import OmniLoop
+        from ..providers.omni.parser import OmniParser
+
+        if self.loop_type == AgenticLoop.ANTHROPIC:
+            from ..providers.anthropic.loop import AnthropicLoop
+
+            return AnthropicLoop(
+                api_key=self.api_key,
+                model=self.model,
+                computer=self.computer,
+                save_trajectory=self.save_trajectory,
+                base_dir=self.trajectory_dir,
+                only_n_most_recent_images=self.only_n_most_recent_images,
+                **self._kwargs,
+            )
+
+        # Initialize parser for OmniLoop with appropriate device
+        if "parser" not in self._kwargs:
+            self._kwargs["parser"] = OmniParser()
+
+        return OmniLoop(
+            provider=self.provider,
+            api_key=self.api_key,
+            model=self.model,
+            computer=self.computer,
+            save_trajectory=self.save_trajectory,
+            base_dir=self.trajectory_dir,
+            only_n_most_recent_images=self.only_n_most_recent_images,
+            **self._kwargs,
+        )
+
+    async def _execute_task(self, task: str) -> AsyncGenerator[Dict[str, Any], None]:
+        """Execute a task using the appropriate loop.
+
+        Args:
+            task: Task description to execute
+
+        Yields:
+            Dict containing response content and metadata
+        """
+        try:
+            # Format the messages based on loop type
+            if self.loop_type == AgenticLoop.ANTHROPIC:
+                # Anthropic format
+                messages = [{"role": "user", "content": [{"type": "text", "text": task}]}]
+            else:
+                # Cua format
+                messages = [{"role": "user", "content": task}]
+
+            # Run the loop
+            try:
+                async for result in self.loop.run(messages):
+                    if result is None:
+                        break
+
+                    # Handle error case
+                    if "error" in result:
+                        yield {
+                            "role": "assistant",
+                            "content": result["error"],
+                            "metadata": {"title": "❌ Error"},
+                        }
+                        continue
+
+                    # Extract content and metadata based on loop type
+                    if self.loop_type == AgenticLoop.ANTHROPIC:
+                        # Handle Anthropic format
+                        if "content" in result:
+                            content_text = ""
+                            for content_block in result["content"]:
+                                try:
+                                    # Try to access the text attribute directly
+                                    content_text += content_block.text
+                                except (AttributeError, TypeError):
+                                    # If it's a dictionary instead of an object
+                                    if isinstance(content_block, dict) and "text" in content_block:
+                                        content_text += content_block["text"]
+
+                            yield {
+                                "role": "assistant",
+                                "content": content_text,
+                                "metadata": result.get("parsed_screen", {}),
+                            }
+                        else:
+                            yield {
+                                "role": "assistant",
+                                "content": str(result),
+                                "metadata": {"title": "Screen Analysis"},
+                            }
+                    else:
+                        # Handle Omni format
+                        content = ""
+                        metadata = {"title": "Screen Analysis"}
+
+                        # If result has content (normal case)
+                        if "content" in result:
+                            content = result["content"]
+
+                            # Ensure metadata has a title
+                            if isinstance(content, dict) and "metadata" in content:
+                                metadata = content["metadata"]
+                                if "title" not in metadata:
+                                    metadata["title"] = "Screen Analysis"
+
+                            # For string content, convert to proper format
+                            if isinstance(content, str):
+                                content = content
+                            elif isinstance(content, dict) and "content" in content:
+                                content = content.get("content", "")
+
+                        yield {"role": "assistant", "content": content, "metadata": metadata}
+            except Exception as e:
+                logger.error(f"Error running the loop: {str(e)}")
+                yield {
+                    "role": "assistant",
+                    "content": f"Error running the agent loop: {str(e)}",
+                    "metadata": {"title": "❌ Loop Error"},
+                }
+
+        except Exception as e:
+            logger.error(f"Error in _execute_task: {str(e)}")
+            yield {
+                "role": "assistant",
+                "content": f"Error: {str(e)}",
+                "metadata": {"title": "❌ Error"},
+            }

agent/core/base_agent.py

@@ -0,0 +1,164 @@
+"""Base computer agent implementation."""
+
+import asyncio
+import logging
+import os
+from abc import ABC, abstractmethod
+from typing import Any, AsyncGenerator, Dict, Optional
+
+from computer import Computer
+
+from ..types.base import Provider
+
+logger = logging.getLogger(__name__)
+
+
+class BaseComputerAgent(ABC):
+    """Base class for computer agents."""
+
+    def __init__(
+        self,
+        max_retries: int = 3,
+        computer: Optional[Computer] = None,
+        screenshot_dir: Optional[str] = None,
+        log_dir: Optional[str] = None,
+        **kwargs,
+    ):
+        """Initialize the base computer agent.
+
+        Args:
+            max_retries: Maximum number of retry attempts
+            computer: Optional Computer instance
+            screenshot_dir: Directory to save screenshots
+            log_dir: Directory to save logs (set to None to disable logging to files)
+            **kwargs: Additional provider-specific arguments
+        """
+        self.max_retries = max_retries
+        self.computer = computer or Computer()
+        self.queue = asyncio.Queue()
+        self.screenshot_dir = screenshot_dir
+        self.log_dir = log_dir
+        self._retry_count = 0
+        self.provider = Provider.UNKNOWN
+
+        # Setup logging
+        if self.log_dir:
+            os.makedirs(self.log_dir, exist_ok=True)
+            logger.info(f"Created logs directory: {self.log_dir}")
+
+        # Setup screenshots directory
+        if self.screenshot_dir:
+            os.makedirs(self.screenshot_dir, exist_ok=True)
+            logger.info(f"Created screenshots directory: {self.screenshot_dir}")
+
+        logger.info("BaseComputerAgent initialized")
+
+    async def run(self, task: str) -> AsyncGenerator[Dict[str, Any], None]:
+        """Run a task using the computer agent.
+
+        Args:
+            task: Task description
+
+        Yields:
+            Task execution updates
+        """
+        try:
+            logger.info(f"Running task: {task}")
+
+            # Initialize the computer if needed
+            await self._init_if_needed()
+
+            # Execute the task and yield results
+            # The _execute_task method should be implemented to yield results
+            async for result in self._execute_task(task):
+                yield result
+
+        except Exception as e:
+            logger.error(f"Error in agent run method: {str(e)}")
+            yield {
+                "role": "assistant",
+                "content": f"Error: {str(e)}",
+                "metadata": {"title": "❌ Error"},
+            }
+
+    async def _init_if_needed(self):
+        """Initialize the computer interface if it hasn't been initialized yet."""
+        if not self.computer._initialized:
+            logger.info("Computer not initialized, initializing now...")
+            try:
+                # Call run directly without setting the flag first
+                await self.computer.run()
+                logger.info("Computer interface initialized successfully")
+            except Exception as e:
+                logger.error(f"Error initializing computer interface: {str(e)}")
+                raise
+
+    async def __aenter__(self):
+        """Initialize the agent when used as a context manager."""
+        logger.info("Entering BaseComputerAgent context")
+
+        # In case the computer wasn't initialized
+        try:
+            # Initialize the computer only if not already initialized
+            logger.info("Checking if computer is already initialized...")
+            if not self.computer._initialized:
+                logger.info("Initializing computer in __aenter__...")
+                # Use the computer's __aenter__ directly instead of calling run()
+                # This avoids the circular dependency
+                await self.computer.__aenter__()
+                logger.info("Computer initialized in __aenter__")
+            else:
+                logger.info("Computer already initialized, skipping initialization")
+
+            # Take a test screenshot to verify the computer is working
+            logger.info("Testing computer with a screenshot...")
+            try:
+                test_screenshot = await self.computer.screenshot()
+                # Determine the screenshot size based on its type
+                if isinstance(test_screenshot, bytes):
+                    size = len(test_screenshot)
+                else:
+                    # Assume it's an object with base64_image attribute
+                    try:
+                        size = len(test_screenshot.base64_image)
+                    except AttributeError:
+                        size = "unknown"
+                logger.info(f"Screenshot test successful, size: {size}")
+            except Exception as e:
+                logger.error(f"Screenshot test failed: {str(e)}")
+                # Even though screenshot failed, we continue since some tests might not need it
+        except Exception as e:
+            logger.error(f"Error initializing computer in __aenter__: {str(e)}")
+            raise
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Cleanup computer resources if needed."""
+        logger.info("Cleaning up agent resources")
+
+        # Do any necessary cleanup
+        # We're not shutting down the computer here as it might be shared
+        # Just log that we're exiting
+        if exc_type:
+            logger.error(f"Exiting agent context with error: {exc_type.__name__}: {exc_val}")
+        else:
+            logger.info("Exiting agent context normally")
+
+        # If we have a queue, make sure to signal it's done
+        if hasattr(self, "queue") and self.queue:
+            await self.queue.put(None)  # Signal that we're done
+
+    @abstractmethod
+    async def _execute_task(self, task: str) -> AsyncGenerator[Dict[str, Any], None]:
+        """Execute a task. Must be implemented by subclasses.
+
+        This is an async method that returns an AsyncGenerator. Implementations
+        should use 'yield' statements to produce results asynchronously.
+        """
+        yield {
+            "role": "assistant",
+            "content": "Base class method called",
+            "metadata": {"title": "Error"},
+        }
+        raise NotImplementedError("Subclasses must implement _execute_task")