cua-agent 0.3.2__py3-none-any.whl → 0.4.0b1__py3-none-any.whl

agent/core/base.py

@@ -1,217 +0,0 @@
-"""Base loop definitions."""
-
-import logging
-import asyncio
-from abc import ABC, abstractmethod
-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
-from .callbacks import CallbackManager, CallbackHandler
-
-logger = logging.getLogger(__name__)
-
-
-class BaseLoop(ABC):
-    """Base class for agent loops that handle message processing and tool execution."""
-
-    def __init__(
-        self,
-        computer: Computer,
-        model: str,
-        api_key: str,
-        max_tokens: int = 4096,
-        max_retries: int = 3,
-        retry_delay: float = 1.0,
-        base_dir: Optional[str] = "trajectories",
-        save_trajectory: bool = True,
-        only_n_most_recent_images: Optional[int] = 2,
-        callback_handlers: Optional[List[CallbackHandler]] = None,
-        **kwargs,
-    ):
-        """Initialize base agent loop.
-
-        Args:
-            computer: Computer instance to control
-            model: Model name to use
-            api_key: API key for provider
-            max_tokens: Maximum tokens to generate
-            max_retries: Maximum number of retries
-            retry_delay: Delay between retries in seconds
-            base_dir: Base directory for saving experiment data
-            save_trajectory: Whether to save trajectory data
-            only_n_most_recent_images: Maximum number of recent screenshots to include in API requests
-            **kwargs: Additional provider-specific arguments
-        """
-        self.computer = computer
-        self.model = model
-        self.api_key = api_key
-        self.max_tokens = max_tokens
-        self.max_retries = max_retries
-        self.retry_delay = retry_delay
-        self.base_dir = base_dir
-        self.save_trajectory = save_trajectory
-        self.only_n_most_recent_images = only_n_most_recent_images
-        self._kwargs = kwargs
-
-        # 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:
-            self.experiment_manager = ExperimentManager(
-                base_dir=self.base_dir,
-                only_n_most_recent_images=only_n_most_recent_images,
-            )
-            # Track directories for convenience
-            self.run_dir = self.experiment_manager.run_dir
-            self.current_turn_dir = self.experiment_manager.current_turn_dir
-        else:
-            self.experiment_manager = None
-            self.run_dir = None
-            self.current_turn_dir = None
-
-        # Initialize basic tracking
-        self.turn_count = 0
-        
-        # Initialize callback manager
-        self.callback_manager = CallbackManager(handlers=callback_handlers or [])
-
-    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
-        
-    @abstractmethod
-    async def cancel(self) -> None:
-        """Cancel the currently running agent loop task.
-        
-        This method should stop any ongoing processing in the agent loop
-        and clean up resources appropriately.
-        """
-        raise NotImplementedError
-
-    ###########################################
-    # EXPERIMENT AND TRAJECTORY MANAGEMENT
-    ###########################################
-
-    def _setup_experiment_dirs(self) -> None:
-        """Setup the experiment directory structure."""
-        if self.experiment_manager:
-            # Use the experiment manager to set up directories
-            self.experiment_manager.setup_experiment_dirs()
-
-            # Update local tracking variables
-            self.run_dir = self.experiment_manager.run_dir
-            self.current_turn_dir = self.experiment_manager.current_turn_dir
-
-    def _create_turn_dir(self) -> None:
-        """Create a new directory for the current turn."""
-        if self.experiment_manager:
-            # Use the experiment manager to create the turn directory
-            self.experiment_manager.create_turn_dir()
-
-            # Update local tracking variables
-            self.current_turn_dir = self.experiment_manager.current_turn_dir
-            self.turn_count = self.experiment_manager.turn_count
-
-    def _log_api_call(
-        self, call_type: str, request: Any, response: Any = None, error: Optional[Exception] = None
-    ) -> 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 in provider-specific format
-            response: Optional API response data in provider-specific format
-            error: Optional error information
-        """
-        if self.experiment_manager:
-            # Use the experiment manager to log the API call
-            provider = getattr(self, "provider", "unknown")
-            provider_str = str(provider) if provider else "unknown"
-
-            self.experiment_manager.log_api_call(
-                call_type=call_type,
-                request=request,
-                provider=provider_str,
-                model=self.model,
-                response=response,
-                error=error,
-            )
-
-    def _save_screenshot(self, img_base64: str, action_type: str = "") -> None:
-        """Save a screenshot to the experiment directory.
-
-        Args:
-            img_base64: Base64 encoded screenshot
-            action_type: Type of action that triggered the screenshot
-        """
-        if self.experiment_manager:
-            self.experiment_manager.save_screenshot(img_base64, action_type)
-            
-    ###########################################
-    # EVENT HOOKS / CALLBACKS
-    ###########################################
-    
-    async def handle_screenshot(self, screenshot_base64: str, action_type: str = "", parsed_screen: Optional[dict] = None) -> None:
-        """Process a screenshot through callback managers
-        
-        Args:
-            screenshot_base64: Base64 encoded screenshot
-            action_type: Type of action that triggered the screenshot
-        """
-        if hasattr(self, 'callback_manager'):
-            await self.callback_manager.on_screenshot(screenshot_base64, action_type, parsed_screen)

agent/core/callbacks.py

@@ -1,200 +0,0 @@
-"""Callback handlers for agent."""
-
-import json
-import logging
-from abc import ABC, abstractmethod
-from datetime import datetime
-from typing import Any, Dict, List, Optional, Protocol
-
-logger = logging.getLogger(__name__)
-
-class ContentCallback(Protocol):
-    """Protocol for content callbacks."""
-    def __call__(self, content: Dict[str, Any]) -> None: ...
-
-class ToolCallback(Protocol):
-    """Protocol for tool callbacks."""
-    def __call__(self, result: Any, tool_id: str) -> None: ...
-
-class APICallback(Protocol):
-    """Protocol for API callbacks."""
-    def __call__(self, request: Any, response: Any, error: Optional[Exception] = None) -> None: ...
-
-class ScreenshotCallback(Protocol):
-    """Protocol for screenshot callbacks."""
-    def __call__(self, screenshot_base64: str, action_type: str = "") -> Optional[str]: ...
-
-class BaseCallbackManager(ABC):
-    """Base class for callback managers."""
-    
-    def __init__(
-        self,
-        content_callback: ContentCallback,
-        tool_callback: ToolCallback,
-        api_callback: APICallback,
-    ):
-        """Initialize the callback manager.
-        
-        Args:
-            content_callback: Callback for content updates
-            tool_callback: Callback for tool execution results
-            api_callback: Callback for API interactions
-        """
-        self.content_callback = content_callback
-        self.tool_callback = tool_callback
-        self.api_callback = api_callback
-    
-    @abstractmethod
-    def on_content(self, content: Any) -> None:
-        """Handle content updates."""
-        raise NotImplementedError
-    
-    @abstractmethod
-    def on_tool_result(self, result: Any, tool_id: str) -> None:
-        """Handle tool execution results."""
-        raise NotImplementedError
-    
-    @abstractmethod
-    def on_api_interaction(
-        self,
-        request: Any,
-        response: Any,
-        error: Optional[Exception] = None
-    ) -> None:
-        """Handle API interactions."""
-        raise NotImplementedError
-
-
-class CallbackManager:
-    """Manager for callback handlers."""
-    
-    def __init__(self, handlers: Optional[List["CallbackHandler"]] = None):
-        """Initialize with optional handlers.
-        
-        Args:
-            handlers: List of callback handlers
-        """
-        self.handlers = handlers or []
-    
-    def add_handler(self, handler: "CallbackHandler") -> None:
-        """Add a callback handler.
-        
-        Args:
-            handler: Callback handler to add
-        """
-        self.handlers.append(handler)
-    
-    async def on_action_start(self, action: str, **kwargs) -> None:
-        """Called when an action starts.
-        
-        Args:
-            action: Action name
-            **kwargs: Additional data
-        """
-        for handler in self.handlers:
-            await handler.on_action_start(action, **kwargs)
-    
-    async def on_action_end(self, action: str, success: bool, **kwargs) -> None:
-        """Called when an action ends.
-        
-        Args:
-            action: Action name
-            success: Whether the action was successful
-            **kwargs: Additional data
-        """
-        for handler in self.handlers:
-            await handler.on_action_end(action, success, **kwargs)
-    
-    async def on_error(self, error: Exception, **kwargs) -> None:
-        """Called when an error occurs.
-        
-        Args:
-            error: Exception that occurred
-            **kwargs: Additional data
-        """
-        for handler in self.handlers:
-            await handler.on_error(error, **kwargs)
-            
-    async def on_screenshot(self, screenshot_base64: str, action_type: str = "", parsed_screen: Optional[dict] = None) -> None:
-        """Called when a screenshot is taken.
-        
-        Args:
-            screenshot_base64: Base64 encoded screenshot
-            action_type: Type of action that triggered the screenshot
-            parsed_screen: Optional output from parsing the screenshot
-            
-        Returns:
-            Modified screenshot or original if no modifications
-        """
-        for handler in self.handlers:
-            await handler.on_screenshot(screenshot_base64, action_type, parsed_screen)
-
-class CallbackHandler(ABC):
-    """Base class for callback handlers."""
-    
-    @abstractmethod
-    async def on_action_start(self, action: str, **kwargs) -> None:
-        """Called when an action starts.
-        
-        Args:
-            action: Action name
-            **kwargs: Additional data
-        """
-        pass
-    
-    @abstractmethod
-    async def on_action_end(self, action: str, success: bool, **kwargs) -> None:
-        """Called when an action ends.
-        
-        Args:
-            action: Action name
-            success: Whether the action was successful
-            **kwargs: Additional data
-        """
-        pass
-    
-    @abstractmethod
-    async def on_error(self, error: Exception, **kwargs) -> None:
-        """Called when an error occurs.
-        
-        Args:
-            error: Exception that occurred
-            **kwargs: Additional data
-        """
-        pass
-        
-    @abstractmethod
-    async def on_screenshot(self, screenshot_base64: str, action_type: str = "", parsed_screen: Optional[dict] = None) -> None:
-        """Called when a screenshot is taken.
-        
-        Args:
-            screenshot_base64: Base64 encoded screenshot
-            action_type: Type of action that triggered the screenshot
-            
-        Returns:
-            Optional modified screenshot
-        """
-        pass
-
-class DefaultCallbackHandler(CallbackHandler):
-    """Default implementation of CallbackHandler with no-op methods.
-    
-    This class implements all abstract methods from CallbackHandler,
-    allowing subclasses to override only the methods they need.
-    """
-    
-    async def on_action_start(self, action: str, **kwargs) -> None:
-        """Default no-op implementation."""
-        pass
-    
-    async def on_action_end(self, action: str, success: bool, **kwargs) -> None:
-        """Default no-op implementation."""
-        pass
-    
-    async def on_error(self, error: Exception, **kwargs) -> None:
-        """Default no-op implementation."""
-        pass
-        
-    async def on_screenshot(self, screenshot_base64: str, action_type: str = "") -> None:
-        """Default no-op implementation."""
-        pass

agent/core/experiment.py

@@ -1,249 +0,0 @@
-"""Core experiment management for agents."""
-
-import os
-import logging
-import base64
-from io import BytesIO
-from datetime import datetime
-from typing import Any, Dict, List, Optional
-from PIL import Image
-import json
-import re
-
-logger = logging.getLogger(__name__)
-
-
-class ExperimentManager:
-    """Manages experiment directories and logging for the agent."""
-
-    def __init__(
-        self,
-        base_dir: Optional[str] = None,
-        only_n_most_recent_images: Optional[int] = None,
-    ):
-        """Initialize the experiment manager.
-
-        Args:
-            base_dir: Base directory for saving experiment data
-            only_n_most_recent_images: Maximum number of recent screenshots to include in API requests
-        """
-        self.base_dir = base_dir
-        self.only_n_most_recent_images = only_n_most_recent_images
-        self.run_dir = None
-        self.current_turn_dir = None
-        self.turn_count = 0
-        self.screenshot_count = 0
-        # Track all screenshots for potential API request inclusion
-        self.screenshot_paths = []
-
-        # Set up experiment directories if base_dir is provided
-        if self.base_dir:
-            self.setup_experiment_dirs()
-
-    def setup_experiment_dirs(self) -> None:
-        """Setup the experiment directory structure."""
-        if not self.base_dir:
-            return
-
-        # Create base experiments directory if it doesn't exist
-        os.makedirs(self.base_dir, exist_ok=True)
-
-        # Create timestamped run directory
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        self.run_dir = os.path.join(self.base_dir, timestamp)
-        os.makedirs(self.run_dir, exist_ok=True)
-        logger.info(f"Created run directory: {self.run_dir}")
-
-        # Create first turn directory
-        self.create_turn_dir()
-
-    def create_turn_dir(self) -> None:
-        """Create a new directory for the current turn."""
-        if not self.run_dir:
-            logger.warning("Cannot create turn directory: run_dir not set")
-            return
-
-        # Increment turn counter
-        self.turn_count += 1
-
-        # Create turn directory with padded number
-        turn_name = f"turn_{self.turn_count:03d}"
-        self.current_turn_dir = os.path.join(self.run_dir, turn_name)
-        os.makedirs(self.current_turn_dir, exist_ok=True)
-        logger.info(f"Created turn directory: {self.current_turn_dir}")
-
-    def sanitize_log_data(self, data: Any) -> Any:
-        """Sanitize log data by replacing large binary data with placeholders.
-
-        Args:
-            data: Data to sanitize
-
-        Returns:
-            Sanitized copy of the data
-        """
-        if isinstance(data, dict):
-            result = {}
-            for k, v in data.items():
-                # 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]
-        elif isinstance(data, str) and len(data) > 1000 and "base64" in data.lower():
-            return f"[BASE64_DATA_LENGTH_{len(data)}]"
-        else:
-            return data
-
-    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 None
-
-        try:
-            # Increment screenshot counter
-            self.screenshot_count += 1
-
-            # Sanitize action_type to ensure valid filename
-            # Replace characters that are not safe for filenames
-            sanitized_action = ""
-            if action_type:
-                # Replace invalid filename characters with underscores
-                sanitized_action = re.sub(r'[\\/*?:"<>|]', "_", action_type)
-                # Limit the length to avoid excessively long filenames
-                sanitized_action = sanitized_action[:50]
-
-            # Create a descriptive filename
-            timestamp = int(datetime.now().timestamp() * 1000)
-            action_suffix = f"_{sanitized_action}" if sanitized_action else ""
-            filename = f"screenshot_{self.screenshot_count:03d}{action_suffix}_{timestamp}.png"
-
-            # Save directly to the turn directory
-            filepath = os.path.join(self.current_turn_dir, filename)
-
-            # Save the screenshot
-            img_data = base64.b64decode(img_base64)
-            with open(filepath, "wb") as f:
-                f.write(img_data)
-
-            # Keep track of the file path
-            self.screenshot_paths.append(filepath)
-
-            return filepath
-        except Exception as e:
-            logger.error(f"Error saving screenshot: {str(e)}")
-            return None
-
-    def save_action_visualization(
-        self, img: Image.Image, action_name: str, details: str = ""
-    ) -> str:
-        """Save a visualization of an action.
-
-        Args:
-            img: Image to save
-            action_name: Name of the action
-            details: Additional details about the action
-
-        Returns:
-            Path to the saved image
-        """
-        if not self.current_turn_dir:
-            return ""
-
-        try:
-            # Create a descriptive filename
-            timestamp = int(datetime.now().timestamp() * 1000)
-            details_suffix = f"_{details}" if details else ""
-            filename = f"vis_{action_name}{details_suffix}_{timestamp}.png"
-
-            # Save directly to the turn directory
-            filepath = os.path.join(self.current_turn_dir, filename)
-
-            # Save the image
-            img.save(filepath)
-
-            # Keep track of the file path
-            self.screenshot_paths.append(filepath)
-
-            return filepath
-        except Exception as e:
-            logger.error(f"Error saving action visualization: {str(e)}")
-            return ""
-
-    def log_api_call(
-        self,
-        call_type: str,
-        request: Any,
-        provider: str = "unknown",
-        model: str = "unknown",
-        response: Any = None,
-        error: Optional[Exception] = None,
-    ) -> None:
-        """Log API call details to file.
-
-        Args:
-            call_type: Type of API call (request, response, error)
-            request: Request data
-            provider: API provider name
-            model: Model name
-            response: Response data (for response logs)
-            error: Error information (for error logs)
-        """
-        if not self.current_turn_dir:
-            logger.warning("Cannot log API call: current_turn_dir not set")
-            return
-
-        try:
-            # Create a timestamp for the log file
-            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-
-            # Create filename based on log type
-            filename = f"api_call_{timestamp}_{call_type}.json"
-            filepath = os.path.join(self.current_turn_dir, filename)
-
-            # Sanitize data before logging
-            sanitized_request = self.sanitize_log_data(request)
-            sanitized_response = self.sanitize_log_data(response) if response is not None else None
-
-            # Prepare log data
-            log_data = {
-                "timestamp": timestamp,
-                "provider": provider,
-                "model": model,
-                "type": call_type,
-                "request": sanitized_request,
-            }
-
-            if sanitized_response is not None:
-                log_data["response"] = sanitized_response
-            if error is not None:
-                log_data["error"] = str(error)
-
-            # Write to file
-            with open(filepath, "w") as f:
-                json.dump(log_data, f, indent=2, default=str)
-
-            logger.info(f"Logged API {call_type} to {filepath}")
-
-        except Exception as e:
-            logger.error(f"Error logging API call: {str(e)}")