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

agent/callbacks/base.py

@@ -0,0 +1,153 @@
+"""
+Base callback handler interface for ComputerAgent preprocessing and postprocessing hooks.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Optional, Union
+
+
+class AsyncCallbackHandler(ABC):
+    """
+    Base class for async callback handlers that can preprocess messages before
+    the agent loop and postprocess output after the agent loop.
+    """
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Called at the start of an agent run loop."""
+        pass
+
+    async def on_run_end(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]], new_items: List[Dict[str, Any]]) -> None:
+        """Called at the end of an agent run loop."""
+        pass
+    
+    async def on_run_continue(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]], new_items: List[Dict[str, Any]]) -> bool:
+        """Called during agent run loop to determine if execution should continue.
+        
+        Args:
+            kwargs: Run arguments
+            old_items: Original messages
+            new_items: New messages generated during run
+            
+        Returns:
+            True to continue execution, False to stop
+        """
+        return True
+    
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Called before messages are sent to the agent loop.
+        
+        Args:
+            messages: List of message dictionaries to preprocess
+            
+        Returns:
+            List of preprocessed message dictionaries
+        """
+        return messages
+    
+    async def on_llm_end(self, output: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Called after the agent loop returns output.
+        
+        Args:
+            output: List of output message dictionaries to postprocess
+            
+        Returns:
+            List of postprocessed output dictionaries
+        """
+        return output
+
+    async def on_computer_call_start(self, item: Dict[str, Any]) -> None:
+        """
+        Called when a computer call is about to start.
+        
+        Args:
+            item: The computer call item dictionary
+        """
+        pass
+    
+    async def on_computer_call_end(self, item: Dict[str, Any], result: List[Dict[str, Any]]) -> None:
+        """
+        Called when a computer call has completed.
+        
+        Args:
+            item: The computer call item dictionary
+            result: The result of the computer call
+        """
+        pass
+    
+    async def on_function_call_start(self, item: Dict[str, Any]) -> None:
+        """
+        Called when a function call is about to start.
+        
+        Args:
+            item: The function call item dictionary
+        """
+        pass
+    
+    async def on_function_call_end(self, item: Dict[str, Any], result: List[Dict[str, Any]]) -> None:
+        """
+        Called when a function call has completed.
+        
+        Args:
+            item: The function call item dictionary
+            result: The result of the function call
+        """
+        pass
+    
+    async def on_text(self, item: Dict[str, Any]) -> None:
+        """
+        Called when a text message is encountered.
+        
+        Args:
+            item: The message item dictionary
+        """
+        pass
+    
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        """
+        Called when an API call is about to start.
+        
+        Args:
+            kwargs: The kwargs being passed to the API call
+        """
+        pass
+    
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """
+        Called when an API call has completed.
+        
+        Args:
+            kwargs: The kwargs that were passed to the API call
+            result: The result of the API call
+        """
+        pass
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """
+        Called when usage information is received.
+        
+        Args:
+            usage: The usage information
+        """
+        pass
+
+    async def on_screenshot(self, screenshot: Union[str, bytes], name: str = "screenshot") -> None:
+        """
+        Called when a screenshot is taken.
+        
+        Args:
+            screenshot: The screenshot image
+            name: The name of the screenshot
+        """
+        pass
+
+    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
+        """
+        Called when responses are received.
+        
+        Args:
+            kwargs: The kwargs being passed to the agent loop
+            responses: The responses received
+        """
+        pass

agent/callbacks/budget_manager.py

@@ -0,0 +1,44 @@
+from typing import Dict, List, Any
+from .base import AsyncCallbackHandler
+
+class BudgetExceededError(Exception):
+    """Exception raised when budget is exceeded."""
+    pass
+
+class BudgetManagerCallback(AsyncCallbackHandler):
+    """Budget manager callback that tracks usage costs and can stop execution when budget is exceeded."""
+    
+    def __init__(self, max_budget: float, reset_after_each_run: bool = True, raise_error: bool = False):
+        """
+        Initialize BudgetManagerCallback.
+        
+        Args:
+            max_budget: Maximum budget allowed
+            reset_after_each_run: Whether to reset budget after each run
+            raise_error: Whether to raise an error when budget is exceeded
+        """
+        self.max_budget = max_budget
+        self.reset_after_each_run = reset_after_each_run
+        self.raise_error = raise_error
+        self.total_cost = 0.0
+    
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Reset budget if configured to do so."""
+        if self.reset_after_each_run:
+            self.total_cost = 0.0
+    
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Track usage costs."""
+        if "response_cost" in usage:
+            self.total_cost += usage["response_cost"]
+    
+    async def on_run_continue(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]], new_items: List[Dict[str, Any]]) -> bool:
+        """Check if budget allows continuation."""
+        if self.total_cost >= self.max_budget:
+            if self.raise_error:
+                raise BudgetExceededError(f"Budget exceeded: ${self.total_cost} >= ${self.max_budget}")
+            else:
+                print(f"Budget exceeded: ${self.total_cost} >= ${self.max_budget}")
+            return False
+        return True
+    

agent/callbacks/image_retention.py

@@ -0,0 +1,139 @@
+"""
+Image retention callback handler that limits the number of recent images in message history.
+"""
+
+from typing import List, Dict, Any, Optional
+from .base import AsyncCallbackHandler
+
+
+class ImageRetentionCallback(AsyncCallbackHandler):
+    """
+    Callback handler that applies image retention policy to limit the number
+    of recent images in message history to prevent context window overflow.
+    """
+    
+    def __init__(self, only_n_most_recent_images: Optional[int] = None):
+        """
+        Initialize the image retention callback.
+        
+        Args:
+            only_n_most_recent_images: If set, only keep the N most recent images in message history
+        """
+        self.only_n_most_recent_images = only_n_most_recent_images
+    
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Apply image retention policy to messages before sending to agent loop.
+        
+        Args:
+            messages: List of message dictionaries
+            
+        Returns:
+            List of messages with image retention policy applied
+        """
+        if self.only_n_most_recent_images is None:
+            return messages
+        
+        return self._apply_image_retention(messages)
+    
+    def _apply_image_retention(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Apply image retention policy to keep only the N most recent images.
+        
+        Removes computer_call_output items with image_url and their corresponding computer_call items,
+        keeping only the most recent N image pairs based on only_n_most_recent_images setting.
+        
+        Args:
+            messages: List of message dictionaries
+            
+        Returns:
+            Filtered list of messages with image retention applied
+        """
+        if self.only_n_most_recent_images is None:
+            return messages
+        
+        # First pass: Assign call_id to reasoning items based on the next computer_call
+        messages_with_call_ids = []
+        for i, msg in enumerate(messages):
+            msg_copy = msg.copy() if isinstance(msg, dict) else msg
+            
+            # If this is a reasoning item without a call_id, find the next computer_call
+            if (msg_copy.get("type") == "reasoning" and 
+                not msg_copy.get("call_id")):
+                # Look ahead for the next computer_call
+                for j in range(i + 1, len(messages)):
+                    next_msg = messages[j]
+                    if (next_msg.get("type") == "computer_call" and 
+                        next_msg.get("call_id")):
+                        msg_copy["call_id"] = next_msg.get("call_id")
+                        break
+            
+            messages_with_call_ids.append(msg_copy)
+        
+        # Find all computer_call_output items with images and their call_ids
+        image_call_ids = []
+        for msg in reversed(messages_with_call_ids):  # Process in reverse to get most recent first
+            if (msg.get("type") == "computer_call_output" and 
+                isinstance(msg.get("output"), dict) and 
+                "image_url" in msg.get("output", {})):
+                call_id = msg.get("call_id")
+                if call_id and call_id not in image_call_ids:
+                    image_call_ids.append(call_id)
+                    if len(image_call_ids) >= self.only_n_most_recent_images:
+                        break
+        
+        # Keep the most recent N image call_ids (reverse to get chronological order)
+        keep_call_ids = set(image_call_ids[:self.only_n_most_recent_images])
+        
+        # Filter messages: remove computer_call, computer_call_output, and reasoning for old images
+        filtered_messages = []
+        for msg in messages_with_call_ids:
+            msg_type = msg.get("type")
+            call_id = msg.get("call_id")
+            
+            # Remove old computer_call items
+            if msg_type == "computer_call" and call_id not in keep_call_ids:
+                # Check if this call_id corresponds to an image call
+                has_image_output = any(
+                    m.get("type") == "computer_call_output" and 
+                    m.get("call_id") == call_id and
+                    isinstance(m.get("output"), dict) and
+                    "image_url" in m.get("output", {})
+                    for m in messages_with_call_ids
+                )
+                if has_image_output:
+                    continue  # Skip this computer_call
+            
+            # Remove old computer_call_output items with images
+            if (msg_type == "computer_call_output" and 
+                call_id not in keep_call_ids and
+                isinstance(msg.get("output"), dict) and 
+                "image_url" in msg.get("output", {})):
+                continue  # Skip this computer_call_output
+            
+            # Remove old reasoning items that are paired with removed computer calls
+            if (msg_type == "reasoning" and 
+                call_id and call_id not in keep_call_ids):
+                # Check if this call_id corresponds to an image call that's being removed
+                has_image_output = any(
+                    m.get("type") == "computer_call_output" and 
+                    m.get("call_id") == call_id and
+                    isinstance(m.get("output"), dict) and
+                    "image_url" in m.get("output", {})
+                    for m in messages_with_call_ids
+                )
+                if has_image_output:
+                    continue  # Skip this reasoning item
+            
+            filtered_messages.append(msg)
+        
+        # Clean up: Remove call_id from reasoning items before returning
+        final_messages = []
+        for msg in filtered_messages:
+            if msg.get("type") == "reasoning" and "call_id" in msg:
+                # Create a copy without call_id for reasoning items
+                cleaned_msg = {k: v for k, v in msg.items() if k != "call_id"}
+                final_messages.append(cleaned_msg)
+            else:
+                final_messages.append(msg)
+        
+        return final_messages

agent/callbacks/logging.py

@@ -0,0 +1,247 @@
+"""
+Logging callback for ComputerAgent that provides configurable logging of agent lifecycle events.
+"""
+
+import json
+import logging
+from typing import Dict, List, Any, Optional, Union
+from .base import AsyncCallbackHandler
+
+
+def sanitize_image_urls(data: Any) -> Any:
+    """
+    Recursively search for 'image_url' keys and set their values to '[omitted]'.
+    
+    Args:
+        data: Any data structure (dict, list, or primitive type)
+        
+    Returns:
+        A deep copy of the data with all 'image_url' values replaced with '[omitted]'
+    """
+    if isinstance(data, dict):
+        # Create a copy of the dictionary
+        sanitized = {}
+        for key, value in data.items():
+            if key == "image_url":
+                sanitized[key] = "[omitted]"
+            else:
+                # Recursively sanitize the value
+                sanitized[key] = sanitize_image_urls(value)
+        return sanitized
+    
+    elif isinstance(data, list):
+        # Recursively sanitize each item in the list
+        return [sanitize_image_urls(item) for item in data]
+    
+    else:
+        # For primitive types (str, int, bool, None, etc.), return as-is
+        return data
+
+
+class LoggingCallback(AsyncCallbackHandler):
+    """
+    Callback handler that logs agent lifecycle events with configurable verbosity.
+    
+    Logging levels:
+    - DEBUG: All events including API calls, message preprocessing, and detailed outputs
+    - INFO: Major lifecycle events (start/end, messages, outputs)  
+    - WARNING: Only warnings and errors
+    - ERROR: Only errors
+    """
+    
+    def __init__(self, logger: Optional[logging.Logger] = None, level: int = logging.INFO):
+        """
+        Initialize the logging callback.
+        
+        Args:
+            logger: Logger instance to use. If None, creates a logger named 'agent.ComputerAgent'
+            level: Logging level (logging.DEBUG, logging.INFO, etc.)
+        """
+        self.logger = logger or logging.getLogger('agent.ComputerAgent')
+        self.level = level
+        
+        # Set up logger if it doesn't have handlers
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(
+                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+            )
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+            self.logger.setLevel(level)
+    
+    def _update_usage(self, usage: Dict[str, Any]) -> None:
+        """Update total usage statistics."""
+        def add_dicts(target: Dict[str, Any], source: Dict[str, Any]) -> None:
+            for key, value in source.items():
+                if isinstance(value, dict):
+                    if key not in target:
+                        target[key] = {}
+                    add_dicts(target[key], value)
+                else:
+                    if key not in target:
+                        target[key] = 0
+                    target[key] += value
+        add_dicts(self.total_usage, usage)
+    
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Called before the run starts."""
+        self.total_usage = {}
+    
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        self._update_usage(usage)
+
+    async def on_run_end(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]], new_items: List[Dict[str, Any]]) -> None:
+        """Called after the run ends."""
+        def format_dict(d, indent=0):
+            lines = []
+            prefix = f" - {' ' * indent}"
+            for key, value in d.items():
+                if isinstance(value, dict):
+                    lines.append(f"{prefix}{key}:")
+                    lines.extend(format_dict(value, indent + 1))
+                elif isinstance(value, float):
+                    lines.append(f"{prefix}{key}: ${value:.4f}")
+                else:
+                    lines.append(f"{prefix}{key}: {value}")
+            return lines
+        
+        formatted_output = "\n".join(format_dict(self.total_usage))
+        self.logger.info(f"Total usage:\n{formatted_output}")
+    
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Called before LLM processing starts."""
+        if self.logger.isEnabledFor(logging.INFO):
+            self.logger.info(f"LLM processing started with {len(messages)} messages")
+        if self.logger.isEnabledFor(logging.DEBUG):
+            sanitized_messages = [sanitize_image_urls(msg) for msg in messages]
+            self.logger.debug(f"LLM input messages: {json.dumps(sanitized_messages, indent=2)}")
+        return messages
+    
+    async def on_llm_end(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Called after LLM processing ends."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            sanitized_messages = [sanitize_image_urls(msg) for msg in messages]
+            self.logger.debug(f"LLM output: {json.dumps(sanitized_messages, indent=2)}")
+        return messages
+    
+    async def on_computer_call_start(self, item: Dict[str, Any]) -> None:
+        """Called when a computer call starts."""
+        action = item.get("action", {})
+        action_type = action.get("type", "unknown")
+        action_args = {k: v for k, v in action.items() if k != "type"}
+        
+        # INFO level logging for the action
+        self.logger.info(f"Computer: {action_type}({action_args})")
+        
+        # DEBUG level logging for full details
+        if self.logger.isEnabledFor(logging.DEBUG):
+            self.logger.debug(f"Computer call started: {json.dumps(action, indent=2)}")
+    
+    async def on_computer_call_end(self, item: Dict[str, Any], result: Any) -> None:
+        """Called when a computer call ends."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            action = item.get("action", "unknown")
+            self.logger.debug(f"Computer call completed: {json.dumps(action, indent=2)}")
+            if result:
+                sanitized_result = sanitize_image_urls(result)
+                self.logger.debug(f"Computer call result: {json.dumps(sanitized_result, indent=2)}")
+    
+    async def on_function_call_start(self, item: Dict[str, Any]) -> None:
+        """Called when a function call starts."""
+        name = item.get("name", "unknown")
+        arguments = item.get("arguments", "{}")
+        
+        # INFO level logging for the function call
+        self.logger.info(f"Function: {name}({arguments})")
+        
+        # DEBUG level logging for full details
+        if self.logger.isEnabledFor(logging.DEBUG):
+            self.logger.debug(f"Function call started: {name}")
+    
+    async def on_function_call_end(self, item: Dict[str, Any], result: Any) -> None:
+        """Called when a function call ends."""
+        # INFO level logging for function output (similar to function_call_output)
+        if result:
+            # Handle both list and direct result formats
+            if isinstance(result, list) and len(result) > 0:
+                output = result[0].get("output", str(result)) if isinstance(result[0], dict) else str(result[0])
+            else:
+                output = str(result)
+            
+            # Truncate long outputs
+            if len(output) > 100:
+                output = output[:100] + "..."
+            
+            self.logger.info(f"Output: {output}")
+        
+        # DEBUG level logging for full details
+        if self.logger.isEnabledFor(logging.DEBUG):
+            name = item.get("name", "unknown")
+            self.logger.debug(f"Function call completed: {name}")
+            if result:
+                self.logger.debug(f"Function call result: {json.dumps(result, indent=2)}")
+    
+    async def on_text(self, item: Dict[str, Any]) -> None:
+        """Called when a text message is encountered."""
+        # Get the role to determine if it's Agent or User
+        role = item.get("role", "unknown")
+        content_items = item.get("content", [])
+        
+        # Process content items to build display text
+        text_parts = []
+        for content_item in content_items:
+            content_type = content_item.get("type", "output_text")
+            if content_type == "output_text":
+                text_content = content_item.get("text", "")
+                if not text_content.strip():
+                    text_parts.append("[empty]")
+                else:
+                    # Truncate long text and add ellipsis
+                    if len(text_content) > 2048:
+                        text_parts.append(text_content[:2048] + "...")
+                    else:
+                        text_parts.append(text_content)
+            else:
+                # Non-text content, show as [type]
+                text_parts.append(f"[{content_type}]")
+        
+        # Join all text parts
+        display_text = ''.join(text_parts) if text_parts else "[empty]"
+        
+        # Log with appropriate level and format
+        if role == "assistant":
+            self.logger.info(f"Agent: {display_text}")
+        elif role == "user":
+            self.logger.info(f"User: {display_text}")
+        else:
+            # Fallback for unknown roles, use debug level
+            if self.logger.isEnabledFor(logging.DEBUG):
+                self.logger.debug(f"Text message ({role}): {display_text}")
+    
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        """Called when an API call is about to start."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            model = kwargs.get("model", "unknown")
+            self.logger.debug(f"API call starting for model: {model}")
+            # Log sanitized messages if present
+            if "messages" in kwargs:
+                sanitized_messages = sanitize_image_urls(kwargs["messages"])
+                self.logger.debug(f"API call messages: {json.dumps(sanitized_messages, indent=2)}")
+            elif "input" in kwargs:
+                sanitized_input = sanitize_image_urls(kwargs["input"])
+                self.logger.debug(f"API call input: {json.dumps(sanitized_input, indent=2)}")
+    
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """Called when an API call has completed."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            model = kwargs.get("model", "unknown")
+            self.logger.debug(f"API call completed for model: {model}")
+            self.logger.debug(f"API call result: {json.dumps(sanitize_image_urls(result), indent=2)}")
+
+    async def on_screenshot(self, item: Union[str, bytes], name: str = "screenshot") -> None:
+        """Called when a screenshot is taken."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            image_size = len(item) / 1024
+            self.logger.debug(f"Screenshot captured: {name} {image_size:.2f} KB")