cua-agent 0.1.41__tar.gz → 0.1.42__tar.gz

{cua_agent-0.1.41 → cua_agent-0.1.42}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.1.41
+Version: 0.1.42
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: >=3.10

{cua_agent-0.1.41 → cua_agent-0.1.42}/agent/core/base.py

@@ -131,6 +131,15 @@ class BaseLoop(ABC):
             An async generator that yields agent responses
         """
         raise NotImplementedError
+        
+    @abstractmethod
+    async def cancel(self) -> None:
+        """Cancel the currently running agent loop task.
+        
+        This method should stop any ongoing processing in the agent loop
+        and clean up resources appropriately.
+        """
+        raise NotImplementedError
 
     ###########################################
     # EXPERIMENT AND TRAJECTORY MANAGEMENT

{cua_agent-0.1.41 → cua_agent-0.1.42}/agent/providers/anthropic/loop.py

@@ -101,6 +101,7 @@ class AnthropicLoop(BaseLoop):
         self.tool_manager = None
         self.callback_manager = None
         self.queue = asyncio.Queue()  # Initialize queue
+        self.loop_task = None  # Store the loop task for cancellation
 
         # Initialize handlers
         self.api_handler = AnthropicAPIHandler(self)
@@ -169,7 +170,7 @@ class AnthropicLoop(BaseLoop):
                 logger.info("Client initialized successfully")
 
             # Start loop in background task
-            loop_task = asyncio.create_task(self._run_loop(queue, messages))
+            self.loop_task = asyncio.create_task(self._run_loop(queue, messages))
 
             # Process and yield messages as they arrive
             while True:
@@ -184,7 +185,7 @@ class AnthropicLoop(BaseLoop):
                     continue
 
             # Wait for loop to complete
-            await loop_task
+            await self.loop_task
 
             # Send completion message
             yield {
@@ -200,6 +201,31 @@ class AnthropicLoop(BaseLoop):
                 "content": f"Error: {str(e)}",
                 "metadata": {"title": "❌ Error"},
             }
+            
+    async def cancel(self) -> None:
+        """Cancel the currently running agent loop task.
+        
+        This method stops the ongoing processing in the agent loop
+        by cancelling the loop_task if it exists and is running.
+        """
+        if self.loop_task and not self.loop_task.done():
+            logger.info("Cancelling Anthropic loop task")
+            self.loop_task.cancel()
+            try:
+                # Wait for the task to be cancelled with a timeout
+                await asyncio.wait_for(self.loop_task, timeout=2.0)
+            except asyncio.TimeoutError:
+                logger.warning("Timeout while waiting for loop task to cancel")
+            except asyncio.CancelledError:
+                logger.info("Loop task cancelled successfully")
+            except Exception as e:
+                logger.error(f"Error while cancelling loop task: {str(e)}")
+            finally:
+                # Put None in the queue to signal any waiting consumers to stop
+                await self.queue.put(None)
+                logger.info("Anthropic loop task cancelled")
+        else:
+            logger.info("No active Anthropic loop task to cancel")
 
     ###########################################
     # AGENT LOOP IMPLEMENTATION

{cua_agent-0.1.41 → cua_agent-0.1.42}/agent/providers/omni/loop.py

@@ -105,6 +105,7 @@ class OmniLoop(BaseLoop):
         # Set API client attributes
         self.client = None
         self.retry_count = 0
+        self.loop_task = None  # Store the loop task for cancellation
 
         # Initialize handlers
         self.api_handler = OmniAPIHandler(loop=self)
@@ -583,7 +584,31 @@ class OmniLoop(BaseLoop):
         # Initialize the message manager with the provided messages
         self.message_manager.messages = messages.copy()
         logger.info(f"Starting OmniLoop run with {len(self.message_manager.messages)} messages")
+        
+        # Create a task to run the loop
+        self.loop_task = asyncio.create_task(self._run_loop(messages))
 
+        # Yield from the loop task
+        try:
+            async for response in self.loop_task:
+                yield response
+        except Exception as e:
+            logger.error(f"Error in run method: {str(e)}")
+            yield {
+                "role": "assistant",
+                "content": f"Error: {str(e)}",
+                "metadata": {"title": "❌ Error"},
+            }
+            
+    async def _run_loop(self, messages: List[Dict[str, Any]]) -> AsyncGenerator[AgentResponse, None]:
+        """Internal method to run the agent loop with provided messages.
+        
+        Args:
+            messages: List of messages in standard OpenAI format
+            
+        Yields:
+            Agent response format
+        """
         # Continue running until explicitly told to stop
         running = True
         turn_created = False
@@ -688,7 +713,7 @@ class OmniLoop(BaseLoop):
 
             except Exception as e:
                 attempt += 1
-                error_msg = f"Error in run method (attempt {attempt}/{max_attempts}): {str(e)}"
+                error_msg = f"Error in _run_loop method (attempt {attempt}/{max_attempts}): {str(e)}"
                 logger.error(error_msg)
 
                 # If this is our last attempt, provide more info about the error
@@ -696,12 +721,36 @@ class OmniLoop(BaseLoop):
                     logger.error(f"Maximum retry attempts reached. Last error was: {str(e)}")
 
                 yield {
-                    "error": str(e),
+                    "role": "assistant",
+                    "content": f"Error: {str(e)}",
                     "metadata": {"title": "❌ Error"},
                 }
 
                 # Create a brief delay before retrying
                 await asyncio.sleep(1)
+                
+    async def cancel(self) -> None:
+        """Cancel the currently running agent loop task.
+        
+        This method stops the ongoing processing in the agent loop
+        by cancelling the loop_task if it exists and is running.
+        """
+        if self.loop_task and not self.loop_task.done():
+            logger.info("Cancelling Omni loop task")
+            self.loop_task.cancel()
+            try:
+                # Wait for the task to be cancelled with a timeout
+                await asyncio.wait_for(self.loop_task, timeout=2.0)
+            except asyncio.TimeoutError:
+                logger.warning("Timeout while waiting for loop task to cancel")
+            except asyncio.CancelledError:
+                logger.info("Loop task cancelled successfully")
+            except Exception as e:
+                logger.error(f"Error while cancelling loop task: {str(e)}")
+            finally:
+                logger.info("Omni loop task cancelled")
+        else:
+            logger.info("No active Omni loop task to cancel")
 
     async def process_model_response(self, response_text: str) -> Optional[Dict[str, Any]]:
         """Process model response to extract tool calls.

{cua_agent-0.1.41 → cua_agent-0.1.42}/agent/providers/openai/loop.py

@@ -87,6 +87,7 @@ class OpenAILoop(BaseLoop):
         self.acknowledge_safety_check_callback = acknowledge_safety_check_callback
         self.queue = asyncio.Queue()  # Initialize queue
         self.last_response_id = None  # Store the last response ID across runs
+        self.loop_task = None  # Store the loop task for cancellation
 
         # Initialize handlers
         self.api_handler = OpenAIAPIHandler(self)
@@ -138,7 +139,7 @@ class OpenAILoop(BaseLoop):
             await self.tool_manager.initialize()
 
             # Start loop in background task
-            loop_task = asyncio.create_task(self._run_loop(queue, messages))
+            self.loop_task = asyncio.create_task(self._run_loop(queue, messages))
 
             # Process and yield messages as they arrive
             while True:
@@ -153,7 +154,7 @@ class OpenAILoop(BaseLoop):
                     continue
 
             # Wait for loop to complete
-            await loop_task
+            await self.loop_task
 
             # Send completion message
             yield {
@@ -169,6 +170,31 @@ class OpenAILoop(BaseLoop):
                 "content": f"Error: {str(e)}",
                 "metadata": {"title": "❌ Error"},
             }
+            
+    async def cancel(self) -> None:
+        """Cancel the currently running agent loop task.
+        
+        This method stops the ongoing processing in the agent loop
+        by cancelling the loop_task if it exists and is running.
+        """
+        if self.loop_task and not self.loop_task.done():
+            logger.info("Cancelling OpenAI loop task")
+            self.loop_task.cancel()
+            try:
+                # Wait for the task to be cancelled with a timeout
+                await asyncio.wait_for(self.loop_task, timeout=2.0)
+            except asyncio.TimeoutError:
+                logger.warning("Timeout while waiting for loop task to cancel")
+            except asyncio.CancelledError:
+                logger.info("Loop task cancelled successfully")
+            except Exception as e:
+                logger.error(f"Error while cancelling loop task: {str(e)}")
+            finally:
+                # Put None in the queue to signal any waiting consumers to stop
+                await self.queue.put(None)
+                logger.info("OpenAI loop task cancelled")
+        else:
+            logger.info("No active OpenAI loop task to cancel")
 
     ###########################################
     # AGENT LOOP IMPLEMENTATION

{cua_agent-0.1.41 → cua_agent-0.1.42}/agent/providers/uitars/loop.py

@@ -93,6 +93,7 @@ class UITARSLoop(BaseLoop):
         # Set API client attributes
         self.client = None
         self.retry_count = 0
+        self.loop_task = None  # Store the loop task for cancellation
 
         # Initialize visualization helper
         self.viz_helper = VisualizationHelper(agent=self)
@@ -465,7 +466,31 @@ class UITARSLoop(BaseLoop):
         # Initialize the message manager with the provided messages
         self.message_manager.messages = messages.copy()
         logger.info(f"Starting UITARSLoop run with {len(self.message_manager.messages)} messages")
+        
+        # Create a task to run the loop
+        self.loop_task = asyncio.create_task(self._run_loop(messages))
 
+        # Yield from the loop task
+        try:
+            async for response in self.loop_task:
+                yield response
+        except Exception as e:
+            logger.error(f"Error in run method: {str(e)}")
+            yield {
+                "role": "assistant",
+                "content": f"Error: {str(e)}",
+                "metadata": {"title": "❌ Error"},
+            }
+            
+    async def _run_loop(self, messages: List[Dict[str, Any]]) -> AsyncGenerator[AgentResponse, None]:
+        """Internal method to run the agent loop with provided messages.
+        
+        Args:
+            messages: List of messages in standard OpenAI format
+            
+        Yields:
+            Agent response format
+        """
         # Continue running until explicitly told to stop
         running = True
         turn_created = False
@@ -558,6 +583,29 @@ class UITARSLoop(BaseLoop):
                 # Create a brief delay before retrying
                 await asyncio.sleep(1)
 
+    async def cancel(self) -> None:
+        """Cancel the currently running agent loop task.
+        
+        This method stops the ongoing processing in the agent loop
+        by cancelling the loop_task if it exists and is running.
+        """
+        if self.loop_task and not self.loop_task.done():
+            logger.info("Cancelling UITARS loop task")
+            self.loop_task.cancel()
+            try:
+                # Wait for the task to be cancelled with a timeout
+                await asyncio.wait_for(self.loop_task, timeout=2.0)
+            except asyncio.TimeoutError:
+                logger.warning("Timeout while waiting for loop task to cancel")
+            except asyncio.CancelledError:
+                logger.info("Loop task cancelled successfully")
+            except Exception as e:
+                logger.error(f"Error while cancelling loop task: {str(e)}")
+            finally:
+                logger.info("UITARS loop task cancelled")
+        else:
+            logger.info("No active UITARS loop task to cancel")
+
     ###########################################
     # UTILITY METHODS
     ###########################################

{cua_agent-0.1.41 → cua_agent-0.1.42}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "cua-agent"
-version = "0.1.41"
+version = "0.1.42"
 description = "CUA (Computer Use) Agent for AI-driven computer interaction"
 readme = "README.md"
 authors = [
@@ -109,7 +109,7 @@ target-version = [
 
 [tool.ruff]
 line-length = 100
-target-version = "0.1.41"
+target-version = "0.1.42"
 select = [
     "E",
     "F",
@@ -123,7 +123,7 @@ docstring-code-format = true
 
 [tool.mypy]
 strict = true
-python_version = "0.1.41"
+python_version = "0.1.42"
 ignore_missing_imports = true
 disallow_untyped_defs = true
 check_untyped_defs = true