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

agent/callbacks/telemetry.py

@@ -2,45 +2,47 @@
 Telemetry callback handler for Computer-Use Agent (cua-agent)
 """
 
+import platform
 import time
 import uuid
-from typing import List, Dict, Any, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
-from .base import AsyncCallbackHandler
-from ..telemetry import (
-    record_event,
+from core.telemetry import (
     is_telemetry_enabled,
-    set_dimension,
-    SYSTEM_INFO,
+    record_event,
 )
 
+from .base import AsyncCallbackHandler
+
+SYSTEM_INFO = {
+    "os": platform.system().lower(),
+    "os_version": platform.release(),
+    "python_version": platform.python_version(),
+}
+
 
 class TelemetryCallback(AsyncCallbackHandler):
     """
     Telemetry callback handler for Computer-Use Agent (cua-agent)
-    
+
     Tracks agent usage, performance metrics, and optionally trajectory data.
     """
-    
-    def __init__(
-        self, 
-        agent, 
-        log_trajectory: bool = False
-    ):
+
+    def __init__(self, agent, log_trajectory: bool = False):
         """
         Initialize telemetry callback.
-        
+
         Args:
             agent: The ComputerAgent instance
             log_trajectory: Whether to log full trajectory items (opt-in)
         """
         self.agent = agent
         self.log_trajectory = log_trajectory
-        
+
         # Generate session/run IDs
         self.session_id = str(uuid.uuid4())
         self.run_id = None
-        
+
         # Track timing and metrics
         self.run_start_time = None
         self.step_count = 0
@@ -49,132 +51,165 @@ class TelemetryCallback(AsyncCallbackHandler):
             "prompt_tokens": 0,
             "completion_tokens": 0,
             "total_tokens": 0,
-            "response_cost": 0.0
+            "response_cost": 0.0,
         }
-        
+
         # Record agent initialization
         if is_telemetry_enabled():
             self._record_agent_initialization()
-    
+
     def _record_agent_initialization(self) -> None:
         """Record agent type/model and session initialization."""
+        # Get the agent loop type (class name)
+        agent_type = "unknown"
+        if hasattr(self.agent, "agent_loop") and self.agent.agent_loop is not None:
+            agent_type = type(self.agent.agent_loop).__name__
+
         agent_info = {
             "session_id": self.session_id,
-            "agent_type": self.agent.agent_loop.__name__ if hasattr(self.agent, 'agent_loop') else 'unknown',
-            "model": getattr(self.agent, 'model', 'unknown'),
-            **SYSTEM_INFO
+            "agent_type": agent_type,
+            "model": getattr(self.agent, "model", "unknown"),
+            **SYSTEM_INFO,
         }
-        
-        # Set session-level dimensions
-        set_dimension("session_id", self.session_id)
-        set_dimension("agent_type", agent_info["agent_type"])
-        set_dimension("model", agent_info["model"])
-        
+
+        # Include VM name if available
+        vm_name = self._get_vm_name()
+        if vm_name:
+            agent_info["vm_name"] = vm_name
+
         record_event("agent_session_start", agent_info)
-    
+
     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 is_telemetry_enabled():
             return
-            
+
         self.run_id = str(uuid.uuid4())
         self.run_start_time = time.time()
         self.step_count = 0
-        
+
         # Calculate input context size
         input_context_size = self._calculate_context_size(old_items)
-        
+
         run_data = {
             "session_id": self.session_id,
             "run_id": self.run_id,
             "start_time": self.run_start_time,
             "input_context_size": input_context_size,
-            "num_existing_messages": len(old_items)
+            "num_existing_messages": len(old_items),
         }
-        
+
+        # Include VM name if available
+        vm_name = self._get_vm_name()
+        if vm_name:
+            run_data["vm_name"] = vm_name
+
         # Log trajectory if opted in
         if self.log_trajectory:
             trajectory = self._extract_trajectory(old_items)
             if trajectory:
                 run_data["uploaded_trajectory"] = trajectory
-        
-        set_dimension("run_id", self.run_id)
+
         record_event("agent_run_start", run_data)
-    
-    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 at the end of an agent run loop."""
         if not is_telemetry_enabled() or not self.run_start_time:
             return
-            
+
         run_duration = time.time() - self.run_start_time
-        
+
         run_data = {
             "session_id": self.session_id,
             "run_id": self.run_id,
             "end_time": time.time(),
             "duration_seconds": run_duration,
             "num_steps": self.step_count,
-            "total_usage": self.total_usage.copy()
+            "total_usage": self.total_usage.copy(),
         }
-        
+
+        # Include VM name if available
+        vm_name = self._get_vm_name()
+        if vm_name:
+            run_data["vm_name"] = vm_name
+
         # Log trajectory if opted in
         if self.log_trajectory:
             trajectory = self._extract_trajectory(new_items)
             if trajectory:
                 run_data["uploaded_trajectory"] = trajectory
-        
+
         record_event("agent_run_end", run_data)
-    
+
     async def on_usage(self, usage: Dict[str, Any]) -> None:
         """Called when usage information is received."""
         if not is_telemetry_enabled():
             return
-            
+
         # Accumulate usage stats
         self.total_usage["prompt_tokens"] += usage.get("prompt_tokens", 0)
-        self.total_usage["completion_tokens"] += usage.get("completion_tokens", 0) 
+        self.total_usage["completion_tokens"] += usage.get("completion_tokens", 0)
         self.total_usage["total_tokens"] += usage.get("total_tokens", 0)
         self.total_usage["response_cost"] += usage.get("response_cost", 0.0)
-        
+
         # Record individual usage event
         usage_data = {
             "session_id": self.session_id,
             "run_id": self.run_id,
             "step": self.step_count,
-            **usage
+            **usage,
         }
-        
+
         record_event("agent_usage", usage_data)
-    
+
     async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
         """Called when responses are received."""
         if not is_telemetry_enabled():
             return
-            
+
         self.step_count += 1
         step_duration = None
-        
+
         if self.step_start_time:
             step_duration = time.time() - self.step_start_time
-        
+
         self.step_start_time = time.time()
-        
+
         step_data = {
             "session_id": self.session_id,
             "run_id": self.run_id,
             "step": self.step_count,
-            "timestamp": self.step_start_time
+            "timestamp": self.step_start_time,
         }
-        
+
         if step_duration is not None:
             step_data["duration_seconds"] = step_duration
-        
+
         record_event("agent_step", step_data)
-    
+
+    def _get_vm_name(self) -> Optional[str]:
+        """Extract VM name from agent's computer handler if available."""
+        try:
+            if hasattr(self.agent, "computer_handler") and self.agent.computer_handler:
+                handler = self.agent.computer_handler
+                # Check if it's a cuaComputerHandler with a cua_computer
+                if hasattr(handler, "cua_computer"):
+                    computer = handler.cua_computer
+                    if hasattr(computer, "config") and hasattr(computer.config, "name"):
+                        return computer.config.name
+        except Exception:
+            pass
+        return None
+
     def _calculate_context_size(self, items: List[Dict[str, Any]]) -> int:
         """Calculate approximate context size in tokens/characters."""
         total_size = 0
-        
+
         for item in items:
             if item.get("type") == "message" and "content" in item:
                 content = item["content"]
@@ -186,25 +221,27 @@ class TelemetryCallback(AsyncCallbackHandler):
                             total_size += len(part["text"])
             elif "content" in item and isinstance(item["content"], str):
                 total_size += len(item["content"])
-                
+
         return total_size
-    
+
     def _extract_trajectory(self, items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
         """Extract trajectory items that should be logged."""
         trajectory = []
-        
+
         for item in items:
             # Include user messages, assistant messages, reasoning, computer calls, and computer outputs
             if (
-                item.get("role") == "user" or  # User inputs
-                (item.get("type") == "message" and item.get("role") == "assistant") or  # Model outputs
-                item.get("type") == "reasoning" or  # Reasoning traces
-                item.get("type") == "computer_call" or  # Computer actions
-                item.get("type") == "computer_call_output"  # Computer outputs
+                item.get("role") == "user"  # User inputs
+                or (
+                    item.get("type") == "message" and item.get("role") == "assistant"
+                )  # Model outputs
+                or item.get("type") == "reasoning"  # Reasoning traces
+                or item.get("type") == "computer_call"  # Computer actions
+                or item.get("type") == "computer_call_output"  # Computer outputs
             ):
                 # Create a copy of the item with timestamp
                 trajectory_item = item.copy()
                 trajectory_item["logged_at"] = time.time()
                 trajectory.append(trajectory_item)
-        
-        return trajectory
+
+        return trajectory