cua-agent 0.1.5__tar.gz → 0.1.17__tar.gz

{cua_agent-0.1.5 → cua_agent-0.1.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.1.5
+Version: 0.1.17
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: <3.13,>=3.10
@@ -63,43 +63,13 @@ Description-Content-Type: text/markdown
 
 **Agent** is a Computer Use (CUA) framework for running multi-app agentic workflows targeting macOS and Linux sandbox, supporting local (Ollama) and cloud model providers (OpenAI, Anthropic, Groq, DeepSeek, Qwen). The framework integrates with Microsoft's OmniParser for enhanced UI understanding and interaction.
 
-### Get started with Agent
-
-```python
-from agent import ComputerAgent, AgentLoop, LLMProvider
-from computer import Computer
-
-computer = Computer(verbosity=logging.INFO)
-
-agent = ComputerAgent(
-    computer=computer,
-    loop=AgentLoop.ANTHROPIC,
-    # loop=AgentLoop.OMNI,
-    model=LLM(provider=LLMProvider.ANTHROPIC, name="claude-3-7-sonnet-20250219"),
-    # model=LLM(provider=LLMProvider.OPENAI, name="gpt-4.5-preview"),
-    save_trajectory=True,
-    trajectory_dir=str(Path("trajectories")),
-    only_n_most_recent_images=3,
-    verbosity=logging.INFO,
-)
+> While our north star is to create a 1-click experience, this preview of Agent might be still a bit rough around the edges. We appreciate your patience as we work to improve the experience.
 
-tasks = [
-"""
-Please help me with the following task:
-1. Open Safari browser
-2. Go to Wikipedia.org
-3. Search for "Claude AI" 
-4. Summarize the main points you find about Claude AI
-"""
-]
+### Get started with Agent
 
-async with agent:
-    for i, task in enumerate(tasks, 1):
-        print(f"\nExecuting task {i}/{len(tasks)}: {task}")
-        async for result in agent.run(task):
-            print(result)
-        print(f"Task {i} completed")
-```
+<div align="center">
+    <img src="../../img/agent.png"/>
+</div>
 
 ## Install
 

{cua_agent-0.1.5 → cua_agent-0.1.17}/README.md

@@ -17,43 +17,13 @@
 
 **Agent** is a Computer Use (CUA) framework for running multi-app agentic workflows targeting macOS and Linux sandbox, supporting local (Ollama) and cloud model providers (OpenAI, Anthropic, Groq, DeepSeek, Qwen). The framework integrates with Microsoft's OmniParser for enhanced UI understanding and interaction.
 
-### Get started with Agent
-
-```python
-from agent import ComputerAgent, AgentLoop, LLMProvider
-from computer import Computer
-
-computer = Computer(verbosity=logging.INFO)
-
-agent = ComputerAgent(
-    computer=computer,
-    loop=AgentLoop.ANTHROPIC,
-    # loop=AgentLoop.OMNI,
-    model=LLM(provider=LLMProvider.ANTHROPIC, name="claude-3-7-sonnet-20250219"),
-    # model=LLM(provider=LLMProvider.OPENAI, name="gpt-4.5-preview"),
-    save_trajectory=True,
-    trajectory_dir=str(Path("trajectories")),
-    only_n_most_recent_images=3,
-    verbosity=logging.INFO,
-)
+> While our north star is to create a 1-click experience, this preview of Agent might be still a bit rough around the edges. We appreciate your patience as we work to improve the experience.
 
-tasks = [
-"""
-Please help me with the following task:
-1. Open Safari browser
-2. Go to Wikipedia.org
-3. Search for "Claude AI" 
-4. Summarize the main points you find about Claude AI
-"""
-]
+### Get started with Agent
 
-async with agent:
-    for i, task in enumerate(tasks, 1):
-        print(f"\nExecuting task {i}/{len(tasks)}: {task}")
-        async for result in agent.run(task):
-            print(result)
-        print(f"Task {i} completed")
-```
+<div align="center">
+    <img src="../../img/agent.png"/>
+</div>
 
 ## Install
 

{cua_agent-0.1.5 → cua_agent-0.1.17}/agent/__init__.py

@@ -48,9 +48,8 @@ except Exception as e:
     # Other issues with telemetry
     logger.warning(f"Error initializing telemetry: {e}")
 
-from .core.factory import AgentFactory
-from .core.agent import ComputerAgent
 from .providers.omni.types import LLMProvider, LLM
-from .types.base import Provider, AgentLoop
+from .core.loop import AgentLoop
+from .core.computer_agent import ComputerAgent
 
-__all__ = ["AgentFactory", "Provider", "ComputerAgent", "AgentLoop", "LLMProvider", "LLM"]
+__all__ = ["AgentLoop", "LLMProvider", "LLM", "ComputerAgent"]

{cua_agent-0.1.5 → cua_agent-0.1.17}/agent/core/__init__.py

@@ -1,18 +1,12 @@
 """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, 
+    CallbackManager,
     CallbackHandler,
     BaseCallbackManager,
     ContentCallback,
@@ -21,9 +15,8 @@ from .callbacks import (
 )
 
 __all__ = [
-    "BaseComputerAgent", 
-    "BaseLoop", 
-    "CallbackManager", 
+    "BaseLoop",
+    "CallbackManager",
     "CallbackHandler",
     "BaseMessageManager",
     "ImageRetentionConfig",

cua_agent-0.1.17/agent/core/computer_agent.py

@@ -0,0 +1,244 @@
+"""Main entry point for computer agents."""
+
+import asyncio
+import logging
+import os
+from typing import Any, AsyncGenerator, Dict, Optional, cast, List
+
+from computer import Computer
+from ..providers.anthropic.loop import AnthropicLoop
+from ..providers.omni.loop import OmniLoop
+from ..providers.omni.parser import OmniParser
+from ..providers.omni.types import LLMProvider, LLM
+from .. import AgentLoop
+from .messages import StandardMessageManager, ImageRetentionConfig
+from .types import AgentResponse
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Default models for different providers
+DEFAULT_MODELS = {
+    LLMProvider.OPENAI: "gpt-4o",
+    LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+}
+
+# Map providers to their environment variable names
+ENV_VARS = {
+    LLMProvider.OPENAI: "OPENAI_API_KEY",
+    LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
+}
+
+
+class ComputerAgent:
+    """A computer agent that can perform automated tasks using natural language instructions."""
+
+    def __init__(
+        self,
+        computer: Computer,
+        model: LLM,
+        loop: AgentLoop,
+        max_retries: int = 3,
+        screenshot_dir: Optional[str] = None,
+        log_dir: Optional[str] = None,
+        api_key: Optional[str] = None,
+        save_trajectory: bool = True,
+        trajectory_dir: str = "trajectories",
+        only_n_most_recent_images: Optional[int] = None,
+        verbosity: int = logging.INFO,
+    ):
+        """Initialize the ComputerAgent.
+
+        Args:
+            computer: Computer instance. If not provided, one will be created with default settings.
+            max_retries: Maximum number of retry attempts.
+            screenshot_dir: Directory to save screenshots.
+            log_dir: Directory to save logs (set to None to disable logging to files).
+            model: LLM object containing provider and model name. Takes precedence over provider/model_name.
+            provider: The AI provider to use (e.g., LLMProvider.ANTHROPIC). Only used if model is None.
+            api_key: The API key for the provider. If not provided, will look for environment variable.
+            model_name: The model name to use. Only used if model is None.
+            save_trajectory: Whether to save the trajectory.
+            trajectory_dir: Directory to save the trajectory.
+            only_n_most_recent_images: Maximum number of recent screenshots to include in API requests.
+            verbosity: Logging level.
+        """
+        # Basic agent configuration
+        self.max_retries = max_retries
+        self.computer = computer
+        self.queue = asyncio.Queue()
+        self.screenshot_dir = screenshot_dir
+        self.log_dir = log_dir
+        self._retry_count = 0
+        self._initialized = False
+        self._in_context = False
+
+        # Set logging level
+        logger.setLevel(verbosity)
+
+        # 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}")
+
+        # Use the provided LLM object
+        self.provider = model.provider
+        actual_model_name = model.name or DEFAULT_MODELS.get(self.provider, "")
+
+        # Ensure we have a valid model name
+        if not actual_model_name:
+            actual_model_name = DEFAULT_MODELS.get(self.provider, "")
+            if not actual_model_name:
+                raise ValueError(
+                    f"No model specified for provider {self.provider} and no default found"
+                )
+
+        # Ensure computer is properly cast for typing purposes
+        computer_instance = self.computer
+
+        # Get API key from environment if not provided
+        actual_api_key = api_key or os.environ.get(ENV_VARS[self.provider], "")
+        if not actual_api_key:
+            raise ValueError(f"No API key provided for {self.provider}")
+
+        # Initialize the appropriate loop based on the loop parameter
+        if loop == AgentLoop.ANTHROPIC:
+            self._loop = AnthropicLoop(
+                api_key=actual_api_key,
+                model=actual_model_name,
+                computer=computer_instance,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+            )
+        else:
+            self._loop = OmniLoop(
+                provider=self.provider,
+                api_key=actual_api_key,
+                model=actual_model_name,
+                computer=computer_instance,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+                parser=OmniParser(),
+            )
+
+        # Initialize the message manager from the loop
+        self.message_manager = self._loop.message_manager
+
+        logger.info(
+            f"ComputerAgent initialized with provider: {self.provider}, model: {actual_model_name}"
+        )
+
+    async def __aenter__(self):
+        """Initialize the agent when used as a context manager."""
+        logger.info("Entering ComputerAgent context")
+        self._in_context = True
+
+        # 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()
+                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.interface.screenshot()
+                # Determine the screenshot size based on its type
+                if isinstance(test_screenshot, (bytes, bytearray, memoryview)):
+                    size = len(test_screenshot)
+                elif hasattr(test_screenshot, "base64_image"):
+                    size = len(test_screenshot.base64_image)
+                else:
+                    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
+
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Cleanup agent resources if needed."""
+        logger.info("Cleaning up agent resources")
+        self._in_context = False
+
+        # 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
+
+    async def initialize(self) -> None:
+        """Initialize the agent and its components."""
+        if not self._initialized:
+            # Always initialize the computer if available
+            if self.computer and not self.computer._initialized:
+                await self.computer.run()
+            self._initialized = True
+
+    async def run(self, task: str) -> AsyncGenerator[AgentResponse, None]:
+        """Run a task using the computer agent.
+
+        Args:
+            task: Task description
+
+        Yields:
+            Agent response format
+        """
+        try:
+            logger.info(f"Running task: {task}")
+            logger.info(
+                f"Message history before task has {len(self.message_manager.messages)} messages"
+            )
+
+            # Initialize the computer if needed
+            if not self._initialized:
+                await self.initialize()
+
+            # Add task as a user message using the message manager
+            self.message_manager.add_user_message([{"type": "text", "text": task}])
+            logger.info(
+                f"Added task message. Message history now has {len(self.message_manager.messages)} messages"
+            )
+
+            # Pass properly formatted messages to the loop
+            if self._loop is None:
+                logger.error("Loop not initialized properly")
+                yield {"error": "Loop not initialized properly"}
+                return
+
+            # Execute the task and yield results
+            async for result in self._loop.run(self.message_manager.messages):
+                # Yield the result to the caller
+                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"},
+            }

{cua_agent-0.1.5 → cua_agent-0.1.17}/agent/core/experiment.py

@@ -84,7 +84,21 @@ class ExperimentManager:
         if isinstance(data, dict):
             result = {}
             for k, v in data.items():
-                result[k] = self.sanitize_log_data(v)
+                # Special handling for 'data' field in Anthropic message source
+                if k == "data" and isinstance(v, str) and len(v) > 1000:
+                    result[k] = f"[BASE64_DATA_LENGTH_{len(v)}]"
+                # Special handling for the 'media_type' key which indicates we're in an image block
+                elif k == "media_type" and "image" in str(v):
+                    result[k] = v
+                    # If we're in an image block, look for a sibling 'data' field with base64 content
+                    if (
+                        "data" in result
+                        and isinstance(result["data"], str)
+                        and len(result["data"]) > 1000
+                    ):
+                        result["data"] = f"[BASE64_DATA_LENGTH_{len(result['data'])}]"
+                else:
+                    result[k] = self.sanitize_log_data(v)
             return result
         elif isinstance(data, list):
             return [self.sanitize_log_data(item) for item in data]
@@ -93,15 +107,18 @@ class ExperimentManager:
         else:
             return data
 
-    def save_screenshot(self, img_base64: str, action_type: str = "") -> None:
+    def save_screenshot(self, img_base64: str, action_type: str = "") -> Optional[str]:
         """Save a screenshot to the experiment directory.
 
         Args:
             img_base64: Base64 encoded screenshot
             action_type: Type of action that triggered the screenshot
+
+        Returns:
+            Path to the saved screenshot or None if there was an error
         """
         if not self.current_turn_dir:
-            return
+            return None
 
         try:
             # Increment screenshot counter

{cua_agent-0.1.5 → cua_agent-0.1.17}/agent/core/loop.py

@@ -2,22 +2,34 @@
 
 import logging
 import asyncio
-import json
-import os
 from abc import ABC, abstractmethod
+from enum import Enum, auto
 from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
 from datetime import datetime
-import base64
 
 from computer import Computer
 from .experiment import ExperimentManager
+from .messages import StandardMessageManager, ImageRetentionConfig
+from .types import AgentResponse
 
 logger = logging.getLogger(__name__)
 
 
+class AgentLoop(Enum):
+    """Enumeration of available loop types."""
+
+    ANTHROPIC = auto()  # Anthropic implementation
+    OMNI = auto()  # OmniLoop implementation
+    # Add more loop types as needed
+
+
 class BaseLoop(ABC):
     """Base class for agent loops that handle message processing and tool execution."""
 
+    ###########################################
+    # INITIALIZATION AND CONFIGURATION
+    ###########################################
+
     def __init__(
         self,
         computer: Computer,
@@ -55,8 +67,6 @@ class BaseLoop(ABC):
         self.save_trajectory = save_trajectory
         self.only_n_most_recent_images = only_n_most_recent_images
         self._kwargs = kwargs
-        self.message_history = []
-        # self.tool_manager = BaseToolManager(computer)
 
         # Initialize experiment manager
         if self.save_trajectory and self.base_dir:
@@ -75,6 +85,64 @@ class BaseLoop(ABC):
         # Initialize basic tracking
         self.turn_count = 0
 
+    async def initialize(self) -> None:
+        """Initialize both the API client and computer interface with retries."""
+        for attempt in range(self.max_retries):
+            try:
+                logger.info(
+                    f"Starting initialization (attempt {attempt + 1}/{self.max_retries})..."
+                )
+
+                # Initialize API client
+                await self.initialize_client()
+
+                logger.info("Initialization complete.")
+                return
+            except Exception as e:
+                if attempt < self.max_retries - 1:
+                    logger.warning(
+                        f"Initialization failed (attempt {attempt + 1}/{self.max_retries}): {str(e)}. Retrying..."
+                    )
+                    await asyncio.sleep(self.retry_delay)
+                else:
+                    logger.error(
+                        f"Initialization failed after {self.max_retries} attempts: {str(e)}"
+                    )
+                    raise RuntimeError(f"Failed to initialize: {str(e)}")
+
+        ###########################################
+
+    # ABSTRACT METHODS TO BE IMPLEMENTED BY SUBCLASSES
+    ###########################################
+
+    @abstractmethod
+    async def initialize_client(self) -> None:
+        """Initialize the API client and any provider-specific components.
+
+        This method must be implemented by subclasses to set up
+        provider-specific clients and tools.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def run(self, messages: List[Dict[str, Any]]) -> AsyncGenerator[AgentResponse, None]:
+        """Run the agent loop with provided messages.
+
+        This method handles the main agent loop including message processing,
+        API calls, response handling, and action execution.
+
+        Args:
+            messages: List of message objects
+
+        Yields:
+            Agent response format
+        """
+        raise NotImplementedError
+
+    ###########################################
+    # EXPERIMENT AND TRAJECTORY MANAGEMENT
+    ###########################################
+
     def _setup_experiment_dirs(self) -> None:
         """Setup the experiment directory structure."""
         if self.experiment_manager:
@@ -100,10 +168,13 @@ class BaseLoop(ABC):
     ) -> None:
         """Log API call details to file.
 
+        Preserves provider-specific formats for requests and responses to ensure
+        accurate logging for debugging and analysis purposes.
+
         Args:
             call_type: Type of API call (e.g., 'request', 'response', 'error')
-            request: The API request data
-            response: Optional API response data
+            request: The API request data in provider-specific format
+            response: Optional API response data in provider-specific format
             error: Optional error information
         """
         if self.experiment_manager:
@@ -129,116 +200,3 @@ class BaseLoop(ABC):
         """
         if self.experiment_manager:
             self.experiment_manager.save_screenshot(img_base64, action_type)
-
-    async def initialize(self) -> None:
-        """Initialize both the API client and computer interface with retries."""
-        for attempt in range(self.max_retries):
-            try:
-                logger.info(
-                    f"Starting initialization (attempt {attempt + 1}/{self.max_retries})..."
-                )
-
-                # Initialize API client
-                await self.initialize_client()
-
-                # Initialize computer
-                await self.computer.initialize()
-
-                logger.info("Initialization complete.")
-                return
-            except Exception as e:
-                if attempt < self.max_retries - 1:
-                    logger.warning(
-                        f"Initialization failed (attempt {attempt + 1}/{self.max_retries}): {str(e)}. Retrying..."
-                    )
-                    await asyncio.sleep(self.retry_delay)
-                else:
-                    logger.error(
-                        f"Initialization failed after {self.max_retries} attempts: {str(e)}"
-                    )
-                    raise RuntimeError(f"Failed to initialize: {str(e)}")
-
-    async def _get_parsed_screen_som(self) -> Dict[str, Any]:
-        """Get parsed screen information.
-
-        Returns:
-            Dict containing screen information
-        """
-        try:
-            # Take screenshot
-            screenshot = await self.computer.interface.screenshot()
-
-            # Initialize with default values
-            width, height = 1024, 768
-            base64_image = ""
-
-            # Handle different types of screenshot returns
-            if isinstance(screenshot, bytes):
-                # Raw bytes screenshot
-                base64_image = base64.b64encode(screenshot).decode("utf-8")
-            elif hasattr(screenshot, "base64_image"):
-                # Object-style screenshot with attributes
-                base64_image = screenshot.base64_image
-                if hasattr(screenshot, "width") and hasattr(screenshot, "height"):
-                    width = screenshot.width
-                    height = screenshot.height
-
-            # Create parsed screen data
-            parsed_screen = {
-                "width": width,
-                "height": height,
-                "parsed_content_list": [],
-                "timestamp": datetime.now().isoformat(),
-                "screenshot_base64": base64_image,
-            }
-
-            # Save screenshot if requested
-            if self.save_trajectory and self.experiment_manager:
-                try:
-                    img_data = base64_image
-                    if "," in img_data:
-                        img_data = img_data.split(",")[1]
-                    self._save_screenshot(img_data, action_type="state")
-                except Exception as e:
-                    logger.error(f"Error saving screenshot: {str(e)}")
-
-            return parsed_screen
-        except Exception as e:
-            logger.error(f"Error taking screenshot: {str(e)}")
-            return {
-                "width": 1024,
-                "height": 768,
-                "parsed_content_list": [],
-                "timestamp": datetime.now().isoformat(),
-                "error": f"Error taking screenshot: {str(e)}",
-                "screenshot_base64": "",
-            }
-
-    @abstractmethod
-    async def initialize_client(self) -> None:
-        """Initialize the API client and any provider-specific components."""
-        raise NotImplementedError
-
-    @abstractmethod
-    async def run(self, messages: List[Dict[str, Any]]) -> AsyncGenerator[Dict[str, Any], None]:
-        """Run the agent loop with provided messages.
-
-        Args:
-            messages: List of message objects
-
-        Yields:
-            Dict containing response data
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def _process_screen(
-        self, parsed_screen: Dict[str, Any], messages: List[Dict[str, Any]]
-    ) -> None:
-        """Process screen information and add to messages.
-
-        Args:
-            parsed_screen: Dictionary containing parsed screen info
-            messages: List of messages to update
-        """
-        raise NotImplementedError