cua-agent 0.4.22__py3-none-any.whl → 0.7.16__py3-none-any.whl

agent/callbacks/logging.py

@@ -4,17 +4,18 @@ Logging callback for ComputerAgent that provides configurable logging of agent l
 
 import json
 import logging
-from typing import Dict, List, Any, Optional, Union
+from typing import Any, Dict, List, 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]'
     """
@@ -28,11 +29,11 @@ def sanitize_image_urls(data: Any) -> Any:
                 # 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
@@ -41,37 +42,36 @@ def sanitize_image_urls(data: Any) -> Any:
 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)  
+    - 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.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'
-            )
+            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):
@@ -82,18 +82,25 @@ class LoggingCallback(AsyncCallbackHandler):
                     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:
+    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}"
@@ -106,10 +113,10 @@ class LoggingCallback(AsyncCallbackHandler):
                 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):
@@ -118,27 +125,27 @@ class LoggingCallback(AsyncCallbackHandler):
             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):
@@ -147,48 +154,52 @@ class LoggingCallback(AsyncCallbackHandler):
             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])
+                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:
@@ -206,10 +217,10 @@ class LoggingCallback(AsyncCallbackHandler):
             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]"
-        
+        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}")
@@ -219,7 +230,7 @@ class LoggingCallback(AsyncCallbackHandler):
             # 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):
@@ -232,16 +243,18 @@ class LoggingCallback(AsyncCallbackHandler):
             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)}")
+            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")
+            self.logger.debug(f"Screenshot captured: {name} {image_size:.2f} KB")

agent/callbacks/operator_validator.py

@@ -9,6 +9,7 @@ Ensures agent output actions conform to expected schemas by fixing common issues
 This runs in on_llm_end, which receives the output array (AgentMessage[] as dicts).
 The purpose is to avoid spending another LLM call to fix broken computer call syntax when possible.
 """
+
 from __future__ import annotations
 
 from typing import Any, Dict, List
@@ -48,6 +49,7 @@ class OperatorNormalizerCallback(AsyncCallbackHandler):
                 action["type"] = "type"
 
             action_type = action.get("type")
+
             def _keep_keys(action: Dict[str, Any], keys_to_keep: List[str]):
                 """Keep only the provided keys on action; delete everything else.
                 Always ensures required 'type' is present if listed in keys_to_keep.
@@ -55,6 +57,7 @@ class OperatorNormalizerCallback(AsyncCallbackHandler):
                 for key in list(action.keys()):
                     if key not in keys_to_keep:
                         del action[key]
+
             # rename "coordinate" to "x", "y"
             if "coordinate" in action:
                 action["x"] = action["coordinate"][0]
@@ -100,39 +103,38 @@ class OperatorNormalizerCallback(AsyncCallbackHandler):
             keep = required_keys_by_type.get(action_type or "")
             if keep:
                 _keep_keys(action, keep)
-            
 
-        # Second pass: if an assistant message is immediately followed by a computer_call,
-        # replace the assistant message itself with a reasoning message with summary text.
-        if isinstance(output, list):
-            for i, item in enumerate(output):
-                # AssistantMessage shape: { type: 'message', role: 'assistant', content: OutputContent[] }
-                if item.get("type") == "message" and item.get("role") == "assistant":
-                    next_idx = i + 1
-                    if next_idx >= len(output):
-                        continue
-                    next_item = output[next_idx]
-                    if not isinstance(next_item, dict):
-                        continue
-                    if next_item.get("type") != "computer_call":
-                        continue
-                    contents = item.get("content") or []
-                    # Extract text from OutputContent[]
-                    text_parts: List[str] = []
-                    if isinstance(contents, list):
-                        for c in contents:
-                            if isinstance(c, dict) and c.get("type") == "output_text" and isinstance(c.get("text"), str):
-                                text_parts.append(c["text"])
-                    text_content = "\n".join(text_parts).strip()
-                    # Replace assistant message with reasoning message
-                    output[i] = {
-                        "type": "reasoning",
-                        "summary": [
-                            {
-                                "type": "summary_text",
-                                "text": text_content,
-                            }
-                        ],
-                    }
+        # # Second pass: if an assistant message is immediately followed by a computer_call,
+        # # replace the assistant message itself with a reasoning message with summary text.
+        # if isinstance(output, list):
+        #     for i, item in enumerate(output):
+        #         # AssistantMessage shape: { type: 'message', role: 'assistant', content: OutputContent[] }
+        #         if item.get("type") == "message" and item.get("role") == "assistant":
+        #             next_idx = i + 1
+        #             if next_idx >= len(output):
+        #                 continue
+        #             next_item = output[next_idx]
+        #             if not isinstance(next_item, dict):
+        #                 continue
+        #             if next_item.get("type") != "computer_call":
+        #                 continue
+        #             contents = item.get("content") or []
+        #             # Extract text from OutputContent[]
+        #             text_parts: List[str] = []
+        #             if isinstance(contents, list):
+        #                 for c in contents:
+        #                     if isinstance(c, dict) and c.get("type") == "output_text" and isinstance(c.get("text"), str):
+        #                         text_parts.append(c["text"])
+        #             text_content = "\n".join(text_parts).strip()
+        #             # Replace assistant message with reasoning message
+        #             output[i] = {
+        #                 "type": "reasoning",
+        #                 "summary": [
+        #                     {
+        #                         "type": "summary_text",
+        #                         "text": text_content,
+        #                     }
+        #                 ],
+        #             }
 
         return output

agent/callbacks/otel.py

@@ -0,0 +1,291 @@
+"""
+OpenTelemetry callback handler for Computer-Use Agent (cua-agent).
+
+Instruments agent operations for the Four Golden Signals:
+- Latency: Operation duration
+- Traffic: Operation counts
+- Errors: Error counts
+- Saturation: Concurrent operations
+"""
+
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import AsyncCallbackHandler
+
+# Import OTEL functions - these are available when cua-core[telemetry] is installed
+try:
+    from core.telemetry import (
+        add_breadcrumb,
+        capture_exception,
+        create_span,
+        is_otel_enabled,
+        record_error,
+        record_operation,
+        record_tokens,
+        set_context,
+        track_concurrent,
+    )
+
+    OTEL_AVAILABLE = True
+except ImportError:
+    OTEL_AVAILABLE = False
+
+    def is_otel_enabled() -> bool:
+        return False
+
+
+class OtelCallback(AsyncCallbackHandler):
+    """
+    OpenTelemetry callback handler for instrumentation.
+
+    Tracks:
+    - Agent session lifecycle (start/end)
+    - Agent run lifecycle (start/end with duration)
+    - Individual steps (with duration)
+    - Computer actions (with duration)
+    - Token usage
+    - Errors
+    """
+
+    def __init__(self, agent: Any):
+        """
+        Initialize OTEL callback.
+
+        Args:
+            agent: The ComputerAgent instance
+        """
+        self.agent = agent
+        self.model = getattr(agent, "model", "unknown")
+
+        # Timing state
+        self.run_start_time: Optional[float] = None
+        self.step_start_time: Optional[float] = None
+        self.step_count = 0
+
+        # Span management
+        self._session_span: Optional[Any] = None
+        self._run_span: Optional[Any] = None
+
+        # Track concurrent sessions
+        self._concurrent_tracker: Optional[Any] = None
+
+        if OTEL_AVAILABLE and is_otel_enabled():
+            # Set context for all events
+            set_context(
+                "agent",
+                {
+                    "model": self.model,
+                    "agent_type": self._get_agent_type(),
+                },
+            )
+
+    def _get_agent_type(self) -> str:
+        """Get the agent loop type name."""
+        if hasattr(self.agent, "agent_loop") and self.agent.agent_loop is not None:
+            return type(self.agent.agent_loop).__name__
+        return "unknown"
+
+    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."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        self.run_start_time = time.perf_counter()
+        self.step_count = 0
+
+        # Add breadcrumb for debugging
+        add_breadcrumb(
+            category="agent",
+            message=f"Agent run started with model {self.model}",
+            level="info",
+            data={
+                "model": self.model,
+                "agent_type": self._get_agent_type(),
+                "input_messages": len(old_items),
+            },
+        )
+
+    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."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        if self.run_start_time is not None:
+            duration = time.perf_counter() - self.run_start_time
+
+            # Record run metrics
+            record_operation(
+                operation="agent.run",
+                duration_seconds=duration,
+                status="success",
+                model=self.model,
+                steps=self.step_count,
+            )
+
+            add_breadcrumb(
+                category="agent",
+                message=f"Agent run completed in {duration:.2f}s",
+                level="info",
+                data={
+                    "duration_seconds": duration,
+                    "steps": self.step_count,
+                    "output_messages": len(new_items),
+                },
+            )
+
+        self.run_start_time = None
+
+    async def on_responses(
+        self, kwargs: Dict[str, Any], responses: Dict[str, Any]
+    ) -> None:
+        """Called when responses are received (each step)."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        self.step_count += 1
+        current_time = time.perf_counter()
+
+        # Calculate step duration if we have a start time
+        if self.step_start_time is not None:
+            step_duration = current_time - self.step_start_time
+            record_operation(
+                operation="agent.step",
+                duration_seconds=step_duration,
+                status="success",
+                model=self.model,
+                step_number=self.step_count,
+            )
+
+        # Start timing next step
+        self.step_start_time = current_time
+
+        add_breadcrumb(
+            category="agent",
+            message=f"Agent step {self.step_count} completed",
+            level="info",
+            data={"step": self.step_count},
+        )
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+
+        if prompt_tokens > 0 or completion_tokens > 0:
+            record_tokens(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                model=self.model,
+            )
+
+    async def on_computer_call_start(self, item: Dict[str, Any]) -> None:
+        """Called when a computer call is about to start."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        action = item.get("action", {})
+        action_type = action.get("type", "unknown")
+
+        add_breadcrumb(
+            category="computer",
+            message=f"Computer action: {action_type}",
+            level="info",
+            data={"action_type": action_type},
+        )
+
+    async def on_computer_call_end(
+        self, item: Dict[str, Any], result: List[Dict[str, Any]]
+    ) -> None:
+        """Called when a computer call has completed."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        action = item.get("action", {})
+        action_type = action.get("type", "unknown")
+
+        # Record computer action metric
+        # Note: We don't have precise timing here, so we record with 0 duration
+        # The actual timing should be done in the computer module
+        record_operation(
+            operation=f"computer.action.{action_type}",
+            duration_seconds=0,  # Timing handled elsewhere
+            status="success",
+            model=self.model,
+        )
+
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        """Called when an LLM API call is about to start."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        add_breadcrumb(
+            category="llm",
+            message="LLM API call started",
+            level="info",
+            data={"model": self.model},
+        )
+
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """Called when an LLM API call has completed."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        add_breadcrumb(
+            category="llm",
+            message="LLM API call completed",
+            level="info",
+        )
+
+
+class OtelErrorCallback(AsyncCallbackHandler):
+    """
+    Callback that captures errors and sends them to Sentry/OTEL.
+
+    Should be added early in the callback chain to catch all errors.
+    """
+
+    def __init__(self, agent: Any):
+        """
+        Initialize error callback.
+
+        Args:
+            agent: The ComputerAgent instance
+        """
+        self.agent = agent
+        self.model = getattr(agent, "model", "unknown")
+
+    async def on_error(self, error: Exception, context: Dict[str, Any]) -> None:
+        """Called when an error occurs during agent execution."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        error_type = type(error).__name__
+        operation = context.get("operation", "unknown")
+
+        # Record error metric
+        record_error(
+            error_type=error_type,
+            operation=operation,
+            model=self.model,
+        )
+
+        # Capture exception in Sentry
+        capture_exception(
+            error,
+            context={
+                "model": self.model,
+                "operation": operation,
+                **{k: v for k, v in context.items() if k != "operation"},
+            },
+        )