cua-agent 0.1.5__py3-none-any.whl → 0.1.17__py3-none-any.whl

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"]

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",

agent/core/computer_agent.py

@@ -1,69 +1,244 @@
 """Main entry point for computer agents."""
 
+import asyncio
 import logging
-from typing import Any, AsyncGenerator, Dict, Optional
+import os
+from typing import Any, AsyncGenerator, Dict, Optional, cast, List
 
 from computer import Computer
-from ..types.base import Provider
-from .factory import AgentFactory
+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, provider: Provider, computer: Optional[Computer] = None, **kwargs):
+    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:
-            provider: The AI provider to use (e.g., Provider.ANTHROPIC)
-            computer: Optional Computer instance. If not provided, one will be created with default settings.
-            **kwargs: Additional provider-specific arguments
+            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.
         """
-        self.provider = provider
-        self._computer = computer
-        self._kwargs = kwargs
-        self._agent = None
+        # 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
 
-        # Create provider-specific agent using factory
-        self._agent = AgentFactory.create(provider=provider, computer=computer, **kwargs)
+        # 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):
-        """Enter the async context manager."""
+        """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):
-        """Exit the async context manager."""
+        """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:
-            if not self._in_context and self._computer:
-                # If not in context manager but have a computer, initialize it
-                await self._computer.run()
+            # 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[Dict[str, Any], None]:
-        """Run the agent with a given task."""
-        if not self._initialized:
-            await self.initialize()
+    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"
+            )
 
-        if self._agent is None:
-            logger.error("Agent not initialized properly")
-            yield {"error": "Agent not initialized properly"}
-            return
+            # 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
 
-        async for result in self._agent.run(task):
-            yield result
+            # 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
 
-    @property
-    def computer(self) -> Optional[Computer]:
-        """Get the underlying computer instance."""
-        return self._agent.computer if self._agent else None
+        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"},
+            }

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

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