cua-agent 0.4.0b4__tar.gz → 0.4.2__tar.gz

{cua_agent-0.4.0b4 → cua_agent-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.4.0b4
+Version: 0.4.2
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: >=3.11
@@ -13,7 +13,7 @@ Requires-Dist: pydantic>=2.6.4
 Requires-Dist: rich>=13.7.1
 Requires-Dist: python-dotenv>=1.0.1
 Requires-Dist: cua-computer<0.5.0,>=0.3.0
-Requires-Dist: cua-core<0.2.0,>=0.1.0
+Requires-Dist: cua-core<0.2.0,>=0.1.8
 Requires-Dist: certifi>=2024.2.2
 Requires-Dist: litellm>=1.74.8
 Provides-Extra: openai

cua_agent-0.4.2/agent/__init__.py

@@ -0,0 +1,64 @@
+"""
+agent - Decorator-based Computer Use Agent with liteLLM integration
+"""
+
+import logging
+import sys
+
+from .decorators import agent_loop
+from .agent import ComputerAgent
+from .types import Messages, AgentResponse
+
+# Import loops to register them
+from . import loops
+
+__all__ = [
+    "agent_loop",
+    "ComputerAgent",
+    "Messages",
+    "AgentResponse"
+]
+
+__version__ = "0.4.0"
+
+logger = logging.getLogger(__name__)
+
+# Initialize telemetry when the package is imported
+try:
+    # Import from core telemetry for basic functions
+    from core.telemetry import (
+        is_telemetry_enabled,
+        flush,
+        record_event,
+    )
+
+    # Import set_dimension from our own telemetry module
+    from .telemetry import set_dimension
+
+    # Check if telemetry is enabled
+    if is_telemetry_enabled():
+        logger.info("Telemetry is enabled")
+
+        # Record package initialization
+        record_event(
+            "module_init",
+            {
+                "module": "agent",
+                "version": __version__,
+                "python_version": sys.version,
+            },
+        )
+
+        # Set the package version as a dimension
+        set_dimension("agent_version", __version__)
+
+        # Flush events to ensure they're sent
+        flush()
+    else:
+        logger.info("Telemetry is disabled")
+except ImportError as e:
+    # Telemetry not available
+    logger.warning(f"Telemetry not available: {e}")
+except Exception as e:
+    # Other issues with telemetry
+    logger.warning(f"Error initializing telemetry: {e}")

{cua_agent-0.4.0b4 → cua_agent-0.4.2}/agent/agent.py

@@ -6,6 +6,7 @@ import asyncio
 from typing import Dict, List, Any, Optional, AsyncGenerator, Union, cast, Callable, Set
 
 from litellm.responses.utils import Usage
+
 from .types import Messages, Computer
 from .decorators import find_agent_loop
 from .computer_handler import OpenAIComputerHandler, acknowledge_safety_check_callback, check_blocklisted_url
@@ -14,7 +15,13 @@ import litellm
 import litellm.utils
 import inspect
 from .adapters import HuggingFaceLocalAdapter
-from .callbacks import ImageRetentionCallback, LoggingCallback, TrajectorySaverCallback, BudgetManagerCallback
+from .callbacks import (
+    ImageRetentionCallback, 
+    LoggingCallback, 
+    TrajectorySaverCallback, 
+    BudgetManagerCallback,
+    TelemetryCallback,
+)
 
 def get_json(obj: Any, max_depth: int = 10) -> Any:
     def custom_serializer(o: Any, depth: int = 0, seen: Set[int] = None) -> Any:
@@ -129,6 +136,7 @@ class ComputerAgent:
         screenshot_delay: Optional[float | int] = 0.5,
         use_prompt_caching: Optional[bool] = False,
         max_trajectory_budget: Optional[float | dict] = None,
+        telemetry_enabled: Optional[bool] = True,
         **kwargs
     ):
         """
@@ -146,6 +154,7 @@ class ComputerAgent:
             screenshot_delay: Delay before screenshots in seconds
             use_prompt_caching: If set, use prompt caching to avoid reprocessing the same prompt. Intended for use with anthropic providers.
             max_trajectory_budget: If set, adds BudgetManagerCallback to track usage costs and stop when budget is exceeded
+            telemetry_enabled: If set, adds TelemetryCallback to track anonymized usage data. Enabled by default.
             **kwargs: Additional arguments passed to the agent loop
         """
         self.model = model
@@ -158,10 +167,18 @@ class ComputerAgent:
         self.max_retries = max_retries
         self.screenshot_delay = screenshot_delay
         self.use_prompt_caching = use_prompt_caching
+        self.telemetry_enabled = telemetry_enabled
         self.kwargs = kwargs
 
         # == Add built-in callbacks ==
 
+        # Add telemetry callback if telemetry_enabled is set
+        if self.telemetry_enabled:
+            if isinstance(self.telemetry_enabled, bool):
+                self.callbacks.append(TelemetryCallback(self))
+            else:
+                self.callbacks.append(TelemetryCallback(self, **self.telemetry_enabled))
+
         # Add logging callback if verbosity is set
         if self.verbosity is not None:
             self.callbacks.append(LoggingCallback(level=self.verbosity))

{cua_agent-0.4.0b4 → cua_agent-0.4.2}/agent/callbacks/__init__.py

@@ -7,6 +7,7 @@ from .image_retention import ImageRetentionCallback
 from .logging import LoggingCallback
 from .trajectory_saver import TrajectorySaverCallback
 from .budget_manager import BudgetManagerCallback
+from .telemetry import TelemetryCallback
 
 __all__ = [
     "AsyncCallbackHandler",
@@ -14,4 +15,5 @@ __all__ = [
     "LoggingCallback",
     "TrajectorySaverCallback",
     "BudgetManagerCallback",
+    "TelemetryCallback",
 ]

{cua_agent-0.4.0b4 → cua_agent-0.4.2}/agent/callbacks/budget_manager.py

@@ -1,44 +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
+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
     

cua_agent-0.4.2/agent/callbacks/telemetry.py

@@ -0,0 +1,210 @@
+"""
+Telemetry callback handler for Computer-Use Agent (cua-agent)
+"""
+
+import time
+import uuid
+from typing import List, Dict, Any, Optional, Union
+
+from .base import AsyncCallbackHandler
+from ..telemetry import (
+    record_event,
+    is_telemetry_enabled,
+    set_dimension,
+    SYSTEM_INFO,
+)
+
+
+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
+    ):
+        """
+        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
+        self.step_start_time = None
+        self.total_usage = {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 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."""
+        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
+        }
+        
+        # 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"])
+        
+        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)
+        }
+        
+        # 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:
+        """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()
+        }
+        
+        # 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["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
+        }
+        
+        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
+        }
+        
+        if step_duration is not None:
+            step_data["duration_seconds"] = step_duration
+        
+        record_event("agent_step", step_data)
+    
+    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"]
+                if isinstance(content, str):
+                    total_size += len(content)
+                elif isinstance(content, list):
+                    for part in content:
+                        if isinstance(part, dict) and "text" in part:
+                            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
+            ):
+                # Create a copy of the item with timestamp
+                trajectory_item = item.copy()
+                trajectory_item["logged_at"] = time.time()
+                trajectory.append(trajectory_item)
+        
+        return trajectory

{cua_agent-0.4.0b4 → cua_agent-0.4.2}/agent/cli.py

@@ -92,26 +92,30 @@ def print_welcome(model: str, agent_loop: str, container_name: str):
 async def ainput(prompt: str = ""):
     return await asyncio.to_thread(input, prompt)
 
-async def chat_loop(agent, model: str, container_name: str):
+async def chat_loop(agent, model: str, container_name: str, initial_prompt: str = ""):
     """Main chat loop with the agent."""
     print_welcome(model, agent.agent_loop.__name__, container_name)
     
     history = []
     
+    if initial_prompt:
+        history.append({"role": "user", "content": initial_prompt})
+    
     while True:
-        # Get user input with prompt
-        print_colored("> ", end="")
-        user_input = await ainput()
-        
-        if user_input.lower() in ['exit', 'quit', 'q']:
-            print_colored("\n👋 Goodbye!")
-            break
+        if history[-1].get("role") != "user":
+            # Get user input with prompt
+            print_colored("> ", end="")
+            user_input = await ainput()
             
-        if not user_input:
-            continue
-            
-        # Add user message to history
-        history.append({"role": "user", "content": user_input})
+            if user_input.lower() in ['exit', 'quit', 'q']:
+                print_colored("\n👋 Goodbye!")
+                break
+                
+            if not user_input:
+                continue
+                
+            # Add user message to history
+            history.append({"role": "user", "content": user_input})
         
         # Stream responses from the agent with spinner
         with yaspin(text="Thinking...", spinner="line", attrs=["dark"]) as spinner:
@@ -204,6 +208,12 @@ Examples:
         action="store_true",
         help="Enable verbose logging"
     )
+
+    parser.add_argument(
+        "-p", "--prompt",
+        type=str,
+        help="Initial prompt to send to the agent. Leave blank for interactive mode."
+    )
     
     args = parser.parse_args()
     
@@ -269,9 +279,11 @@ Examples:
         agent_kwargs = {
             "model": args.model,
             "tools": [computer],
-            "only_n_most_recent_images": args.images,
             "verbosity": 20 if args.verbose else 30,  # DEBUG vs WARNING
         }
+
+        if args.images > 0:
+            agent_kwargs["only_n_most_recent_images"] = args.images
         
         if args.trajectory:
             agent_kwargs["trajectory_dir"] = "trajectories"
@@ -286,7 +298,7 @@ Examples:
         agent = ComputerAgent(**agent_kwargs)
         
         # Start chat loop
-        await chat_loop(agent, args.model, container_name)
+        await chat_loop(agent, args.model, container_name, args.prompt)