cua-agent 0.1.24__py3-none-any.whl → 0.1.26__py3-none-any.whl

agent/core/base.py

@@ -5,10 +5,12 @@ import asyncio
 from abc import ABC, abstractmethod
 from typing import Any, AsyncGenerator, Dict, List, Optional
 
+from agent.providers.omni.parser import ParseResult
 from computer import Computer
 from .messages import StandardMessageManager, ImageRetentionConfig
 from .types import AgentResponse
 from .experiment import ExperimentManager
+from .callbacks import CallbackManager, CallbackHandler
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +29,7 @@ class BaseLoop(ABC):
         base_dir: Optional[str] = "trajectories",
         save_trajectory: bool = True,
         only_n_most_recent_images: Optional[int] = 2,
+        callback_handlers: Optional[List[CallbackHandler]] = None,
         **kwargs,
     ):
         """Initialize base agent loop.
@@ -75,6 +78,9 @@ class BaseLoop(ABC):
 
         # Initialize basic tracking
         self.turn_count = 0
+        
+        # Initialize callback manager
+        self.callback_manager = CallbackManager(handlers=callback_handlers or [])
 
     async def initialize(self) -> None:
         """Initialize both the API client and computer interface with retries."""
@@ -187,3 +193,17 @@ class BaseLoop(ABC):
         """
         if self.experiment_manager:
             self.experiment_manager.save_screenshot(img_base64, action_type)
+            
+    ###########################################
+    # EVENT HOOKS / CALLBACKS
+    ###########################################
+    
+    async def handle_screenshot(self, screenshot_base64: str, action_type: str = "", parsed_screen: Optional[ParseResult] = None) -> None:
+        """Process a screenshot through callback managers
+        
+        Args:
+            screenshot_base64: Base64 encoded screenshot
+            action_type: Type of action that triggered the screenshot
+        """
+        if hasattr(self, 'callback_manager'):
+            await self.callback_manager.on_screenshot(screenshot_base64, action_type, parsed_screen)

agent/core/callbacks.py

@@ -6,6 +6,8 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Protocol
 
+from agent.providers.omni.parser import ParseResult
+
 logger = logging.getLogger(__name__)
 
 class ContentCallback(Protocol):
@@ -20,6 +22,10 @@ class APICallback(Protocol):
     """Protocol for API callbacks."""
     def __call__(self, request: Any, response: Any, error: Optional[Exception] = None) -> None: ...
 
+class ScreenshotCallback(Protocol):
+    """Protocol for screenshot callbacks."""
+    def __call__(self, screenshot_base64: str, action_type: str = "") -> Optional[str]: ...
+
 class BaseCallbackManager(ABC):
     """Base class for callback managers."""
     
@@ -110,7 +116,20 @@ class CallbackManager:
         """
         for handler in self.handlers:
             await handler.on_error(error, **kwargs)
-
+            
+    async def on_screenshot(self, screenshot_base64: str, action_type: str = "", parsed_screen: Optional[ParseResult] = None) -> None:
+        """Called when a screenshot is taken.
+        
+        Args:
+            screenshot_base64: Base64 encoded screenshot
+            action_type: Type of action that triggered the screenshot
+            parsed_screen: Optional output from parsing the screenshot
+            
+        Returns:
+            Modified screenshot or original if no modifications
+        """
+        for handler in self.handlers:
+            await handler.on_screenshot(screenshot_base64, action_type, parsed_screen)
 
 class CallbackHandler(ABC):
     """Base class for callback handlers."""
@@ -144,4 +163,40 @@ class CallbackHandler(ABC):
             error: Exception that occurred
             **kwargs: Additional data
         """
-        pass 
+        pass
+        
+    @abstractmethod
+    async def on_screenshot(self, screenshot_base64: str, action_type: str = "", parsed_screen: Optional[ParseResult] = None) -> None:
+        """Called when a screenshot is taken.
+        
+        Args:
+            screenshot_base64: Base64 encoded screenshot
+            action_type: Type of action that triggered the screenshot
+            
+        Returns:
+            Optional modified screenshot
+        """
+        pass
+
+class DefaultCallbackHandler(CallbackHandler):
+    """Default implementation of CallbackHandler with no-op methods.
+    
+    This class implements all abstract methods from CallbackHandler,
+    allowing subclasses to override only the methods they need.
+    """
+    
+    async def on_action_start(self, action: str, **kwargs) -> None:
+        """Default no-op implementation."""
+        pass
+    
+    async def on_action_end(self, action: str, success: bool, **kwargs) -> None:
+        """Default no-op implementation."""
+        pass
+    
+    async def on_error(self, error: Exception, **kwargs) -> None:
+        """Default no-op implementation."""
+        pass
+        
+    async def on_screenshot(self, screenshot_base64: str, action_type: str = "") -> None:
+        """Default no-op implementation."""
+        pass

agent/providers/anthropic/callbacks/manager.py

@@ -3,23 +3,33 @@ import httpx
 from anthropic.types.beta import BetaContentBlockParam
 from ..tools import ToolResult
 
+
 class APICallback(Protocol):
     """Protocol for API callbacks."""
-    def __call__(self, request: httpx.Request | None, 
-                 response: httpx.Response | object | None, 
-                 error: Exception | None) -> None: ...
+
+    def __call__(
+        self,
+        request: httpx.Request | None,
+        response: httpx.Response | object | None,
+        error: Exception | None,
+    ) -> None: ...
+
 
 class ContentCallback(Protocol):
     """Protocol for content callbacks."""
+
     def __call__(self, content: BetaContentBlockParam) -> None: ...
 
+
 class ToolCallback(Protocol):
     """Protocol for tool callbacks."""
+
     def __call__(self, result: ToolResult, tool_id: str) -> None: ...
 
+
 class CallbackManager:
     """Manages various callbacks for the agent system."""
-    
+
     def __init__(
         self,
         content_callback: ContentCallback,
@@ -27,7 +37,7 @@ class CallbackManager:
         api_callback: APICallback,
     ):
         """Initialize the callback manager.
-        
+
         Args:
             content_callback: Callback for content updates
             tool_callback: Callback for tool execution results
@@ -36,20 +46,20 @@ class CallbackManager:
         self.content_callback = content_callback
         self.tool_callback = tool_callback
         self.api_callback = api_callback
-    
+
     def on_content(self, content: BetaContentBlockParam) -> None:
         """Handle content updates."""
         self.content_callback(content)
-    
+
     def on_tool_result(self, result: ToolResult, tool_id: str) -> None:
         """Handle tool execution results."""
         self.tool_callback(result, tool_id)
-    
+
     def on_api_interaction(
         self,
         request: httpx.Request | None,
         response: httpx.Response | object | None,
-        error: Exception | None
+        error: Exception | None,
     ) -> None:
         """Handle API interactions."""
-        self.api_callback(request, response, error) 
+        self.api_callback(request, response, error)

agent/providers/omni/clients/oaicompat.py

@@ -45,8 +45,8 @@ class OAICompatClient(BaseOmniClient):
             max_tokens: Maximum tokens to generate
             temperature: Generation temperature
         """
-        super().__init__(api_key="EMPTY", model=model)
-        self.api_key = "EMPTY"  # Local endpoints typically don't require an API key
+        super().__init__(api_key=api_key or "EMPTY", model=model)
+        self.api_key = api_key or "EMPTY" # Local endpoints typically don't require an API key
         self.model = model
         self.provider_base_url = (
             provider_base_url or "http://localhost:8000/v1"
@@ -146,10 +146,18 @@ class OAICompatClient(BaseOmniClient):
                 base_url = self.provider_base_url or "http://localhost:8000/v1"
 
                 # Check if the base URL already includes the chat/completions endpoint
+                
                 endpoint_url = base_url
                 if not endpoint_url.endswith("/chat/completions"):
+                    # If URL is RunPod format, make it OpenAI compatible
+                    if endpoint_url.startswith("https://api.runpod.ai/v2/"):
+                        # Extract RunPod endpoint ID
+                        parts = endpoint_url.split("/")
+                        if len(parts) >= 5:
+                            runpod_id = parts[4]
+                            endpoint_url = f"https://api.runpod.ai/v2/{runpod_id}/openai/v1/chat/completions"
                     # If the URL ends with /v1, append /chat/completions
-                    if endpoint_url.endswith("/v1"):
+                    elif endpoint_url.endswith("/v1"):
                         endpoint_url = f"{endpoint_url}/chat/completions"
                     # If the URL doesn't end with /v1, make sure it has a proper structure
                     elif not endpoint_url.endswith("/"):

agent/providers/omni/loop.py

@@ -147,7 +147,7 @@ class OmniLoop(BaseLoop):
             )
         elif self.provider == LLMProvider.OAICOMPAT:
             self.client = OAICompatClient(
-                api_key="EMPTY",  # Local endpoints typically don't require an API key
+                api_key=self.api_key or "EMPTY",  # Local endpoints typically don't require an API key
                 model=self.model,
                 provider_base_url=self.provider_base_url,
             )
@@ -183,7 +183,7 @@ class OmniLoop(BaseLoop):
                 )
             elif self.provider == LLMProvider.OAICOMPAT:
                 self.client = OAICompatClient(
-                    api_key="EMPTY",  # Local endpoints typically don't require an API key
+                    api_key=self.api_key or "EMPTY",  # Local endpoints typically don't require an API key
                     model=self.model,
                     provider_base_url=self.provider_base_url,
                 )
@@ -443,6 +443,8 @@ class OmniLoop(BaseLoop):
                     except (json.JSONDecodeError, IndexError):
                         try:
                             # Look for JSON object pattern
+                            import re  # Local import to ensure availability
+
                             json_pattern = r"\{[^}]+\}"
                             json_match = re.search(json_pattern, raw_text)
                             if json_match:
@@ -453,8 +455,20 @@ class OmniLoop(BaseLoop):
                                 logger.error(f"No JSON found in content")
                                 return True, action_screenshot_saved
                         except json.JSONDecodeError as e:
-                            logger.error(f"Failed to parse JSON from text: {str(e)}")
-                            return True, action_screenshot_saved
+                            # Try to sanitize the JSON string and retry
+                            try:
+                                # Remove or replace invalid control characters
+                                import re  # Local import to ensure availability
+
+                                sanitized_text = re.sub(r"[\x00-\x1F\x7F]", "", raw_text)
+                                # Try parsing again with sanitized text
+                                parsed_content = json.loads(sanitized_text)
+                                logger.info(
+                                    "Successfully parsed JSON after sanitizing control characters"
+                                )
+                            except json.JSONDecodeError:
+                                logger.error(f"Failed to parse JSON from text: {str(e)}")
+                                return True, action_screenshot_saved
 
             # Step 4: Process the parsed content if available
             if parsed_content:
@@ -534,6 +548,10 @@ class OmniLoop(BaseLoop):
                     img_data = parsed_screen.annotated_image_base64
                     if "," in img_data:
                         img_data = img_data.split(",")[1]
+                    
+                    # Process screenshot through hooks and save if needed
+                    await self.handle_screenshot(img_data, action_type="state", parsed_screen=parsed_screen)
+                    
                     # Save with a generic "state" action type to indicate this is the current screen state
                     self._save_screenshot(img_data, action_type="state")
                 except Exception as e:
@@ -649,6 +667,8 @@ class OmniLoop(BaseLoop):
                     response=response,
                     messages=self.message_manager.messages,
                     model=self.model,
+                    parsed_screen=parsed_screen,
+                    parser=self.parser
                 )
 
                 # Yield the response to the caller

agent/providers/openai/loop.py

@@ -194,8 +194,13 @@ class OpenAILoop(BaseLoop):
                 # Convert to base64 if needed
                 if isinstance(screenshot, bytes):
                     screenshot_base64 = base64.b64encode(screenshot).decode("utf-8")
+                elif isinstance(screenshot, (bytearray, memoryview)):
+                    screenshot_base64 = base64.b64encode(screenshot).decode("utf-8")
                 else:
-                    screenshot_base64 = screenshot
+                    screenshot_base64 = str(screenshot)
+
+                # Emit screenshot callbacks
+                await self.handle_screenshot(screenshot_base64, action_type="initial_state")
 
                 # Save screenshot if requested
                 if self.save_trajectory:
@@ -204,8 +209,6 @@ class OpenAILoop(BaseLoop):
                         logger.warning(
                             "Converting non-string screenshot_base64 to string for _save_screenshot"
                         )
-                        if isinstance(screenshot_base64, (bytearray, memoryview)):
-                            screenshot_base64 = base64.b64encode(screenshot_base64).decode("utf-8")
                     self._save_screenshot(screenshot_base64, action_type="state")
                     logger.info("Screenshot saved to trajectory")
 
@@ -336,8 +339,14 @@ class OpenAILoop(BaseLoop):
                         screenshot = await self.computer.interface.screenshot()
                         if isinstance(screenshot, bytes):
                             screenshot_base64 = base64.b64encode(screenshot).decode("utf-8")
+                        elif isinstance(screenshot, (bytearray, memoryview)):
+                            screenshot_base64 = base64.b64encode(bytes(screenshot)).decode("utf-8")
                         else:
-                            screenshot_base64 = screenshot
+                            screenshot_base64 = str(screenshot)
+                            
+                        # Process screenshot through hooks
+                        action_type = f"after_{action.get('type', 'action')}"
+                        await self.handle_screenshot(screenshot_base64, action_type=action_type)
 
                         # Create computer_call_output
                         computer_call_output = {