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

agent/__init__.py

@@ -48,9 +48,7 @@ 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 .types.base import AgentLoop
 
-__all__ = ["AgentFactory", "Provider", "ComputerAgent", "AgentLoop", "LLMProvider", "LLM"]
+__all__ = ["AgentLoop", "LLMProvider", "LLM"]

agent/core/__init__.py

@@ -1,6 +1,5 @@
 """Core agent components."""
 
-from .base_agent import BaseComputerAgent
 from .loop import BaseLoop
 from .messages import (
     create_user_message,
@@ -12,7 +11,7 @@ from .messages import (
     ImageRetentionConfig,
 )
 from .callbacks import (
-    CallbackManager, 
+    CallbackManager,
     CallbackHandler,
     BaseCallbackManager,
     ContentCallback,
@@ -21,9 +20,8 @@ from .callbacks import (
 )
 
 __all__ = [
-    "BaseComputerAgent", 
-    "BaseLoop", 
-    "CallbackManager", 
+    "BaseLoop",
+    "CallbackManager",
     "CallbackHandler",
     "BaseMessageManager",
     "ImageRetentionConfig",

agent/core/computer_agent.py

@@ -1,69 +1,251 @@
 """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
+from dataclasses import dataclass
 
 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
 
 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,
+        parser: Optional[OmniParser] = 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.
+            parser: Parser instance for the OmniLoop. Only used if provider is not ANTHROPIC.
+            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 or 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 = 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(
+                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=actual_parser,
+            )
+
+        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 _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]:
-        """Run the agent with a given task."""
-        if not self._initialized:
-            await self.initialize()
+        """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
+            if not self._initialized:
+                await self.initialize()
+
+            # Format task as a message
+            messages = [{"role": "user", "content": task}]
 
-        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(messages):
+                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

@@ -141,9 +141,6 @@ class BaseLoop(ABC):
                 # Initialize API client
                 await self.initialize_client()
 
-                # Initialize computer
-                await self.computer.initialize()
-
                 logger.info("Initialization complete.")
                 return
             except Exception as e:
@@ -173,15 +170,22 @@ class BaseLoop(ABC):
             base64_image = ""
 
             # Handle different types of screenshot returns
-            if isinstance(screenshot, bytes):
+            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
-                base64_image = screenshot.base64_image
-                if hasattr(screenshot, "width") and hasattr(screenshot, "height"):
-                    width = screenshot.width
-                    height = screenshot.height
+                # 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 = {

agent/core/telemetry.py

@@ -4,58 +4,70 @@ import logging
 import os
 import platform
 import sys
-from typing import Dict, Any
+from typing import Dict, Any, Callable
 
 # Import the core telemetry module
 TELEMETRY_AVAILABLE = False
 
+
+# Local fallbacks in case core telemetry isn't available
+def _noop(*args: Any, **kwargs: Any) -> None:
+    """No-op function for when telemetry is not available."""
+    pass
+
+
+# Define default functions with unique names to avoid shadowing
+_default_record_event = _noop
+_default_increment_counter = _noop
+_default_set_dimension = _noop
+_default_get_telemetry_client = lambda: None
+_default_flush = _noop
+_default_is_telemetry_enabled = lambda: False
+_default_is_telemetry_globally_disabled = lambda: True
+
+# Set the actual functions to the defaults initially
+record_event = _default_record_event
+increment_counter = _default_increment_counter
+set_dimension = _default_set_dimension
+get_telemetry_client = _default_get_telemetry_client
+flush = _default_flush
+is_telemetry_enabled = _default_is_telemetry_enabled
+is_telemetry_globally_disabled = _default_is_telemetry_globally_disabled
+
+logger = logging.getLogger("cua.agent.telemetry")
+
 try:
+    # Import from core telemetry
     from core.telemetry import (
-        record_event,
-        increment,
-        get_telemetry_client,
-        flush,
-        is_telemetry_enabled,
-        is_telemetry_globally_disabled,
+        record_event as core_record_event,
+        increment as core_increment,
+        get_telemetry_client as core_get_telemetry_client,
+        flush as core_flush,
+        is_telemetry_enabled as core_is_telemetry_enabled,
+        is_telemetry_globally_disabled as core_is_telemetry_globally_disabled,
     )
 
+    # Override the default functions with actual implementations
+    record_event = core_record_event
+    get_telemetry_client = core_get_telemetry_client
+    flush = core_flush
+    is_telemetry_enabled = core_is_telemetry_enabled
+    is_telemetry_globally_disabled = core_is_telemetry_globally_disabled
+
     def increment_counter(counter_name: str, value: int = 1) -> None:
         """Wrapper for increment to maintain backward compatibility."""
         if is_telemetry_enabled():
-            increment(counter_name, value)
+            core_increment(counter_name, value)
 
     def set_dimension(name: str, value: Any) -> None:
         """Set a dimension that will be attached to all events."""
-        logger = logging.getLogger("cua.agent.telemetry")
         logger.debug(f"Setting dimension {name}={value}")
 
     TELEMETRY_AVAILABLE = True
-    logger = logging.getLogger("cua.agent.telemetry")
     logger.info("Successfully imported telemetry")
 except ImportError as e:
-    logger = logging.getLogger("cua.agent.telemetry")
     logger.warning(f"Could not import telemetry: {e}")
-    TELEMETRY_AVAILABLE = False
-
-
-# Local fallbacks in case core telemetry isn't available
-def _noop(*args: Any, **kwargs: Any) -> None:
-    """No-op function for when telemetry is not available."""
-    pass
-
-
-logger = logging.getLogger("cua.agent.telemetry")
-
-# If telemetry isn't available, use no-op functions
-if not TELEMETRY_AVAILABLE:
     logger.debug("Telemetry not available, using no-op functions")
-    record_event = _noop  # type: ignore
-    increment_counter = _noop  # type: ignore
-    set_dimension = _noop  # type: ignore
-    get_telemetry_client = lambda: None  # type: ignore
-    flush = _noop  # type: ignore
-    is_telemetry_enabled = lambda: False  # type: ignore
-    is_telemetry_globally_disabled = lambda: True  # type: ignore
 
 # Get system info once to use in telemetry
 SYSTEM_INFO = {
@@ -71,7 +83,7 @@ def enable_telemetry() -> bool:
     Returns:
         bool: True if telemetry was successfully enabled, False otherwise
     """
-    global TELEMETRY_AVAILABLE
+    global TELEMETRY_AVAILABLE, record_event, increment_counter, get_telemetry_client, flush, is_telemetry_enabled, is_telemetry_globally_disabled
 
     # Check if globally disabled using core function
     if TELEMETRY_AVAILABLE and is_telemetry_globally_disabled():