cua-agent 0.4.7__tar.gz → 0.4.9__tar.gz

{cua_agent-0.4.7 → cua_agent-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.4.7
+Version: 0.4.9
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: >=3.11
@@ -15,7 +15,7 @@ Requires-Dist: python-dotenv>=1.0.1
 Requires-Dist: cua-computer<0.5.0,>=0.4.0
 Requires-Dist: cua-core<0.2.0,>=0.1.8
 Requires-Dist: certifi>=2024.2.2
-Requires-Dist: litellm>=1.74.8
+Requires-Dist: litellm>=1.74.12
 Provides-Extra: openai
 Provides-Extra: anthropic
 Provides-Extra: omni

{cua_agent-0.4.7 → cua_agent-0.4.9}/agent/__init__.py

@@ -5,7 +5,7 @@ agent - Decorator-based Computer Use Agent with liteLLM integration
 import logging
 import sys
 
-from .decorators import agent_loop
+from .decorators import register_agent
 from .agent import ComputerAgent
 from .types import Messages, AgentResponse
 
@@ -13,7 +13,7 @@ from .types import Messages, AgentResponse
 from . import loops
 
 __all__ = [
-    "agent_loop",
+    "register_agent",
     "ComputerAgent",
     "Messages",
     "AgentResponse"

{cua_agent-0.4.7 → cua_agent-0.4.9}/agent/adapters/huggingfacelocal_adapter.py

@@ -48,7 +48,11 @@ class HuggingFaceLocalAdapter(CustomLLM):
             )
             
             # Load processor
-            processor = AutoProcessor.from_pretrained(model_name)
+            processor = AutoProcessor.from_pretrained(
+                model_name,
+                min_pixels=3136,
+                max_pixels=4096 * 2160
+            )
             
             # Cache them
             self.models[model_name] = model

{cua_agent-0.4.7 → cua_agent-0.4.9}/agent/agent.py

@@ -3,12 +3,12 @@ ComputerAgent - Main agent class that selects and runs agent loops
 """
 
 import asyncio
-from typing import Dict, List, Any, Optional, AsyncGenerator, Union, cast, Callable, Set
+from typing import Dict, List, Any, Optional, AsyncGenerator, Union, cast, Callable, Set, Tuple
 
 from litellm.responses.utils import Usage
 
-from .types import Messages, Computer
-from .decorators import find_agent_loop
+from .types import Messages, Computer, AgentCapability
+from .decorators import find_agent_config
 from .computer_handler import OpenAIComputerHandler, acknowledge_safety_check_callback, check_blocklisted_url
 import json
 import litellm
@@ -117,6 +117,13 @@ def sanitize_message(msg: Any) -> Any:
             return sanitized
     return msg
 
+def get_output_call_ids(messages: List[Dict[str, Any]]) -> List[str]:
+    call_ids = []
+    for message in messages:
+        if message.get("type") == "computer_call_output" or message.get("type") == "function_call_output":
+            call_ids.append(message.get("call_id"))
+    return call_ids
+
 class ComputerAgent:
     """
     Main agent class that automatically selects the appropriate agent loop
@@ -207,19 +214,21 @@ class ComputerAgent:
         litellm.custom_provider_map = [
             {"provider": "huggingface-local", "custom_handler": hf_adapter}
         ]
+        litellm.suppress_debug_info = True
 
         # == Initialize computer agent ==
 
         # Find the appropriate agent loop
         if custom_loop:
             self.agent_loop = custom_loop
-            self.agent_loop_info = None
+            self.agent_config_info = None
         else:
-            loop_info = find_agent_loop(model)
-            if not loop_info:
-                raise ValueError(f"No agent loop found for model: {model}")
-            self.agent_loop = loop_info.func
-            self.agent_loop_info = loop_info
+            config_info = find_agent_config(model)
+            if not config_info:
+                raise ValueError(f"No agent config found for model: {model}")
+            # Instantiate the agent config class
+            self.agent_loop = config_info.agent_class()
+            self.agent_config_info = config_info
         
         self.tool_schemas = []
         self.computer_handler = None
@@ -389,8 +398,10 @@ class ComputerAgent:
     # AGENT OUTPUT PROCESSING
     # ============================================================================
     
-    async def _handle_item(self, item: Any, computer: Optional[Computer] = None) -> List[Dict[str, Any]]:
+    async def _handle_item(self, item: Any, computer: Optional[Computer] = None, ignore_call_ids: Optional[List[str]] = None) -> List[Dict[str, Any]]:
         """Handle each item; may cause a computer action + screenshot."""
+        if ignore_call_ids and item.get("call_id") and item.get("call_id") in ignore_call_ids:
+            return []
         
         item_type = item.get("type", None)
         
@@ -439,7 +450,7 @@ class ComputerAgent:
             acknowledged_checks = []
             for check in pending_checks:
                 check_message = check.get("message", str(check))
-                if acknowledge_safety_check_callback(check_message):
+                if acknowledge_safety_check_callback(check_message, allow_always=True): # TODO: implement a callback for safety checks
                     acknowledged_checks.append(check)
                 else:
                     raise ValueError(f"Safety check failed: {check_message}")
@@ -514,6 +525,12 @@ class ComputerAgent:
         Returns:
             AsyncGenerator that yields response chunks
         """
+        if not self.agent_config_info:
+            raise ValueError("Agent configuration not found")
+        
+        capabilities = self.get_capabilities()
+        if "step" not in capabilities:
+            raise ValueError(f"Agent loop {self.agent_config_info.agent_class.__name__} does not support step predictions")
 
         await self._initialize_computers()
         
@@ -528,7 +545,7 @@ class ComputerAgent:
             "messages": messages,
             "stream": stream,
             "model": self.model,
-            "agent_loop": self.agent_loop.__name__,
+            "agent_loop": self.agent_config_info.agent_class.__name__,
             **merged_kwargs
         }
         await self._on_run_start(run_kwargs, old_items)
@@ -558,7 +575,7 @@ class ComputerAgent:
             }
 
             # Run agent loop iteration
-            result = await self.agent_loop(
+            result = await self.agent_loop.predict_step(
                 **loop_kwargs,
                 _on_api_start=self._on_api_start,
                 _on_api_end=self._on_api_end,
@@ -579,9 +596,12 @@ class ComputerAgent:
             # Add agent response to new_items
             new_items += result.get("output")
 
+            # Get output call ids
+            output_call_ids = get_output_call_ids(result.get("output", []))
+
             # Handle computer actions
             for item in result.get("output"):
-                partial_items = await self._handle_item(item, self.computer_handler)
+                partial_items = await self._handle_item(item, self.computer_handler, ignore_call_ids=output_call_ids)
                 new_items += partial_items
 
                 # Yield partial response
@@ -594,4 +614,51 @@ class ComputerAgent:
                     )
                 }
         
-        await self._on_run_end(loop_kwargs, old_items, new_items)
+        await self._on_run_end(loop_kwargs, old_items, new_items)
+    
+    async def predict_click(
+        self,
+        instruction: str,
+        image_b64: Optional[str] = None
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates based on image and instruction.
+        
+        Args:
+            instruction: Instruction for where to click
+            image_b64: Base64 encoded image (optional, will take screenshot if not provided)
+            
+        Returns:
+            None or tuple with (x, y) coordinates
+        """
+        if not self.agent_config_info:
+            raise ValueError("Agent configuration not found")
+        
+        capabilities = self.get_capabilities()
+        if "click" not in capabilities:
+            raise ValueError(f"Agent loop {self.agent_config_info.agent_class.__name__} does not support click predictions")
+        if hasattr(self.agent_loop, 'predict_click'):
+            if not image_b64:
+                if not self.computer_handler:
+                    raise ValueError("Computer tool or image_b64 is required for predict_click")
+                image_b64 = await self.computer_handler.screenshot()
+            return await self.agent_loop.predict_click(
+                model=self.model,
+                image_b64=image_b64,
+                instruction=instruction
+            )
+        return None
+    
+    def get_capabilities(self) -> List[AgentCapability]:
+        """
+        Get list of capabilities supported by the current agent config.
+        
+        Returns:
+            List of capability strings (e.g., ["step", "click"])
+        """
+        if not self.agent_config_info:
+            raise ValueError("Agent configuration not found")
+        
+        if hasattr(self.agent_loop, 'get_capabilities'):
+            return self.agent_loop.get_capabilities()
+        return ["step"]  # Default capability

{cua_agent-0.4.7 → cua_agent-0.4.9}/agent/cli.py

@@ -120,7 +120,7 @@ async def ainput(prompt: str = ""):
 
 async def chat_loop(agent, model: str, container_name: str, initial_prompt: str = "", show_usage: bool = True):
     """Main chat loop with the agent."""
-    print_welcome(model, agent.agent_loop.__name__, container_name)
+    print_welcome(model, agent.agent_config_info.agent_class.__name__, container_name)
     
     history = []
     
@@ -130,7 +130,7 @@ async def chat_loop(agent, model: str, container_name: str, initial_prompt: str
     total_cost = 0
 
     while True:
-        if history[-1].get("role") != "user":
+        if len(history) == 0 or history[-1].get("role") != "user":
             # Get user input with prompt
             print_colored("> ", end="")
             user_input = await ainput()
@@ -260,7 +260,12 @@ Examples:
         help="Show total cost of the agent runs"
     )
 
-
+    parser.add_argument(
+        "-r", "--max-retries",
+        type=int,
+        default=3,
+        help="Maximum number of retries for the LLM API calls"
+    )
     
     args = parser.parse_args()
     
@@ -327,6 +332,7 @@ Examples:
             "model": args.model,
             "tools": [computer],
             "verbosity": 20 if args.verbose else 30,  # DEBUG vs WARNING
+            "max_retries": args.max_retries
         }
 
         if args.images > 0:

{cua_agent-0.4.7 → cua_agent-0.4.9}/agent/computer_handler.py

@@ -3,7 +3,7 @@ Computer handler implementation for OpenAI computer-use-preview protocol.
 """
 
 import base64
-from typing import Dict, List, Any, Literal
+from typing import Dict, List, Any, Literal, Union
 from .types import Computer
 
 
@@ -61,8 +61,10 @@ class OpenAIComputerHandler:
         """Move cursor to coordinates."""
         await self.interface.move_cursor(x, y)
     
-    async def keypress(self, keys: List[str]) -> None:
+    async def keypress(self, keys: Union[List[str], str]) -> None:
         """Press key combination."""
+        if isinstance(keys, str):
+            keys = [keys]
         if len(keys) == 1:
             await self.interface.press_key(keys[0])
         else:
@@ -93,8 +95,10 @@ class OpenAIComputerHandler:
         return ""
 
 
-def acknowledge_safety_check_callback(message: str) -> bool:
+def acknowledge_safety_check_callback(message: str, allow_always: bool = False) -> bool:
     """Safety check callback for user acknowledgment."""
+    if allow_always:
+        return True
     response = input(
         f"Safety Check Warning: {message}\nDo you want to acknowledge and proceed? (y/n): "
     ).lower()

cua_agent-0.4.9/agent/decorators.py

@@ -0,0 +1,52 @@
+"""
+Decorators for agent - agent_loop decorator
+"""
+
+from typing import List, Optional
+from .types import AgentConfigInfo
+
+# Global registry
+_agent_configs: List[AgentConfigInfo] = []
+
+def register_agent(models: str, priority: int = 0):
+    """
+    Decorator to register an AsyncAgentConfig class.
+    
+    Args:
+        models: Regex pattern to match supported models
+        priority: Priority for agent selection (higher = more priority)
+    """
+    def decorator(agent_class: type):
+        # Validate that the class implements AsyncAgentConfig protocol
+        if not hasattr(agent_class, 'predict_step'):
+            raise ValueError(f"Agent class {agent_class.__name__} must implement predict_step method")
+        if not hasattr(agent_class, 'predict_click'):
+            raise ValueError(f"Agent class {agent_class.__name__} must implement predict_click method")
+        if not hasattr(agent_class, 'get_capabilities'):
+            raise ValueError(f"Agent class {agent_class.__name__} must implement get_capabilities method")
+        
+        # Register the agent config
+        config_info = AgentConfigInfo(
+            agent_class=agent_class,
+            models_regex=models,
+            priority=priority
+        )
+        _agent_configs.append(config_info)
+        
+        # Sort by priority (highest first)
+        _agent_configs.sort(key=lambda x: x.priority, reverse=True)
+        
+        return agent_class
+    
+    return decorator
+
+def get_agent_configs() -> List[AgentConfigInfo]:
+    """Get all registered agent configs"""
+    return _agent_configs.copy()
+
+def find_agent_config(model: str) -> Optional[AgentConfigInfo]:
+    """Find the best matching agent config for a model"""
+    for config_info in _agent_configs:
+        if config_info.matches_model(model):
+            return config_info
+    return None

{cua_agent-0.4.7 → cua_agent-0.4.9}/agent/loops/__init__.py

@@ -7,5 +7,7 @@ from . import anthropic
 from . import openai
 from . import uitars
 from . import omniparser
+from . import gta1
+from . import composed_grounded
 
-__all__ = ["anthropic", "openai", "uitars", "omniparser"]
+__all__ = ["anthropic", "openai", "uitars", "omniparser", "gta1", "composed_grounded"]