cua-agent 0.1.6__py3-none-any.whl → 0.1.18__py3-none-any.whl

agent/__init__.py

@@ -49,6 +49,7 @@ except Exception as e:
     logger.warning(f"Error initializing telemetry: {e}")
 
 from .providers.omni.types import LLMProvider, LLM
-from .types.base import AgentLoop
+from .core.factory import AgentLoop
+from .core.agent import ComputerAgent
 
-__all__ = ["AgentLoop", "LLMProvider", "LLM"]
+__all__ = ["AgentLoop", "LLMProvider", "LLM", "ComputerAgent"]

agent/core/__init__.py

@@ -1,12 +1,7 @@
 """Core agent components."""
 
-from .loop import BaseLoop
+from .factory import BaseLoop
 from .messages import (
-    create_user_message,
-    create_assistant_message,
-    create_system_message,
-    create_image_message,
-    create_screen_message,
     BaseMessageManager,
     ImageRetentionConfig,
 )

agent/core/{computer_agent.py → agent.py}

@@ -3,31 +3,18 @@
 import asyncio
 import logging
 import os
-from typing import Any, AsyncGenerator, Dict, Optional, cast
-from dataclasses import dataclass
+from typing import AsyncGenerator, Optional
 
 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 ..providers.omni.types import LLM
 from .. import AgentLoop
+from .types import AgentResponse
+from .factory import LoopFactory
+from .provider_config import DEFAULT_MODELS, ENV_VARS
 
 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."""
@@ -44,7 +31,6 @@ class ComputerAgent:
         save_trajectory: bool = True,
         trajectory_dir: str = "trajectories",
         only_n_most_recent_images: Optional[int] = None,
-        parser: Optional[OmniParser] = None,
         verbosity: int = logging.INFO,
     ):
         """Initialize the ComputerAgent.
@@ -61,12 +47,11 @@ class ComputerAgent:
             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.
-            parser: Parser instance for the OmniLoop. Only used if provider is not ANTHROPIC.
             verbosity: Logging level.
         """
         # Basic agent configuration
         self.max_retries = max_retries
-        self.computer = computer or Computer()
+        self.computer = computer
         self.queue = asyncio.Queue()
         self.screenshot_dir = screenshot_dir
         self.log_dir = log_dir
@@ -99,39 +84,30 @@ class ComputerAgent:
                     f"No model specified for provider {self.provider} and no default found"
                 )
 
-        # Ensure computer is properly cast for typing purposes
-        computer_instance = cast(Computer, 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:
-            # Default to OmniLoop for other loop types
-            # Initialize parser if not provided
-            actual_parser = parser or OmniParser()
-
-            self._loop = OmniLoop(
+        # Create the appropriate loop using the factory
+        try:
+            # Let the factory create the appropriate loop with needed components
+            self._loop = LoopFactory.create_loop(
+                loop_type=loop,
                 provider=self.provider,
+                computer=self.computer,
+                model_name=actual_model_name,
                 api_key=actual_api_key,
-                model=actual_model_name,
-                computer=computer_instance,
                 save_trajectory=save_trajectory,
-                base_dir=trajectory_dir,
+                trajectory_dir=trajectory_dir,
                 only_n_most_recent_images=only_n_most_recent_images,
-                parser=actual_parser,
             )
+        except ValueError as e:
+            logger.error(f"Failed to create loop: {str(e)}")
+            raise
+
+        # 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}"
@@ -154,21 +130,6 @@ class ComputerAgent:
             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
@@ -201,36 +162,30 @@ class ComputerAgent:
                 await self.computer.run()
             self._initialized = True
 
-    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
-                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 run(self, task: str) -> AsyncGenerator[Dict[str, Any], None]:
+    async def run(self, task: str) -> AsyncGenerator[AgentResponse, None]:
         """Run a task using the computer agent.
 
         Args:
             task: Task description
 
         Yields:
-            Task execution updates
+            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()
 
-            # Format task as a message
-            messages = [{"role": "user", "content": task}]
+            # 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:
@@ -239,7 +194,7 @@ class ComputerAgent:
                 return
 
             # Execute the task and yield results
-            async for result in self._loop.run(messages):
+            async for result in self._loop.run(self.message_manager.messages):
                 yield result
 
         except Exception as e:

agent/core/{loop.py → base.py}

@@ -1,15 +1,13 @@
-"""Base agent loop implementation."""
+"""Base loop definitions."""
 
 import logging
 import asyncio
-import json
-import os
 from abc import ABC, abstractmethod
-from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
-from datetime import datetime
-import base64
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from computer import Computer
+from .messages import StandardMessageManager, ImageRetentionConfig
+from .types import AgentResponse
 from .experiment import ExperimentManager
 
 logger = logging.getLogger(__name__)
@@ -55,8 +53,11 @@ 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 message manager
+        self.message_manager = StandardMessageManager(
+            config=ImageRetentionConfig(num_images_to_keep=only_n_most_recent_images)
+        )
 
         # Initialize experiment manager
         if self.save_trajectory and self.base_dir:
@@ -75,6 +76,60 @@ 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
+    def run(self, messages: List[Dict[str, Any]]) -> AsyncGenerator[AgentResponse, None]:
+        """Run the agent loop with provided messages.
+
+        Args:
+            messages: List of message objects
+
+        Returns:
+            An async generator that yields agent responses
+        """
+        raise NotImplementedError
+
+    ###########################################
+    # EXPERIMENT AND TRAJECTORY MANAGEMENT
+    ###########################################
+
     def _setup_experiment_dirs(self) -> None:
         """Setup the experiment directory structure."""
         if self.experiment_manager:
@@ -100,10 +155,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,120 +187,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()
-
-                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, bytearray, memoryview)):
-                # Raw bytes screenshot
-                base64_image = base64.b64encode(screenshot).decode("utf-8")
-            elif hasattr(screenshot, "base64_image"):
-                # Object-style screenshot with attributes
-                # Type checking can't infer these attributes, but they exist at runtime
-                # on certain screenshot return types
-                base64_image = getattr(screenshot, "base64_image")
-                width = (
-                    getattr(screenshot, "width", width) if hasattr(screenshot, "width") else width
-                )
-                height = (
-                    getattr(screenshot, "height", height)
-                    if hasattr(screenshot, "height")
-                    else 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

agent/core/factory.py

@@ -0,0 +1,104 @@
+"""Base agent loop implementation."""
+
+import logging
+import importlib.util
+from typing import Dict, Optional, Type, TYPE_CHECKING, Any, cast, Callable, Awaitable
+
+from computer import Computer
+from .types import AgentLoop
+from .base import BaseLoop
+
+# For type checking only
+if TYPE_CHECKING:
+    from ..providers.omni.types import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+
+class LoopFactory:
+    """Factory class for creating agent loops."""
+
+    # Registry to store loop implementations
+    _loop_registry: Dict[AgentLoop, Type[BaseLoop]] = {}
+
+    @classmethod
+    def create_loop(
+        cls,
+        loop_type: AgentLoop,
+        api_key: str,
+        model_name: str,
+        computer: Computer,
+        provider: Any = None,
+        save_trajectory: bool = True,
+        trajectory_dir: str = "trajectories",
+        only_n_most_recent_images: Optional[int] = None,
+        acknowledge_safety_check_callback: Optional[Callable[[str], Awaitable[bool]]] = None,
+    ) -> BaseLoop:
+        """Create and return an appropriate loop instance based on type."""
+        if loop_type == AgentLoop.ANTHROPIC:
+            # Lazy import AnthropicLoop only when needed
+            try:
+                from ..providers.anthropic.loop import AnthropicLoop
+            except ImportError:
+                raise ImportError(
+                    "The 'anthropic' provider is not installed. "
+                    "Install it with 'pip install cua-agent[anthropic]'"
+                )
+
+            return AnthropicLoop(
+                api_key=api_key,
+                model=model_name,
+                computer=computer,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+            )
+        elif loop_type == AgentLoop.OPENAI:
+            # Lazy import OpenAILoop only when needed
+            try:
+                from ..providers.openai.loop import OpenAILoop
+            except ImportError:
+                raise ImportError(
+                    "The 'openai' provider is not installed. "
+                    "Install it with 'pip install cua-agent[openai]'"
+                )
+
+            return OpenAILoop(
+                api_key=api_key,
+                model=model_name,
+                computer=computer,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+                acknowledge_safety_check_callback=acknowledge_safety_check_callback,
+            )
+        elif loop_type == AgentLoop.OMNI:
+            # Lazy import OmniLoop and related classes only when needed
+            try:
+                from ..providers.omni.loop import OmniLoop
+                from ..providers.omni.parser import OmniParser
+                from ..providers.omni.types import LLMProvider
+            except ImportError:
+                raise ImportError(
+                    "The 'omni' provider is not installed. "
+                    "Install it with 'pip install cua-agent[all]'"
+                )
+
+            if provider is None:
+                raise ValueError("Provider is required for OMNI loop type")
+
+            # We know provider is the correct type at this point, so cast it
+            provider_instance = cast(LLMProvider, provider)
+
+            return OmniLoop(
+                provider=provider_instance,
+                api_key=api_key,
+                model=model_name,
+                computer=computer,
+                save_trajectory=save_trajectory,
+                base_dir=trajectory_dir,
+                only_n_most_recent_images=only_n_most_recent_images,
+                parser=OmniParser(),
+            )
+        else:
+            raise ValueError(f"Unsupported loop type: {loop_type}")