cua-agent 0.1.40__py3-none-any.whl → 0.1.42__py3-none-any.whl

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

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

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.

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

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
     ###########################################

agent/ui/gradio/app.py

@@ -1018,12 +1018,18 @@ if __name__ == "__main__":
                         model_string_to_analyze = model_choice_value  # Use the full UI string initially
 
                     try:
-                        # Special case for UITARS - use MLXVLM provider
+                        # Special case for UITARS - use MLXVLM provider or OAICOMPAT for custom
                         if agent_loop_choice == "UITARS":
-                            provider = LLMProvider.MLXVLM
-                            cleaned_model_name_from_func = model_string_to_analyze
-                            agent_loop_type = AgentLoop.UITARS
-                            print(f"Using MLXVLM provider for UITARS model: {model_string_to_analyze}")
+                            if is_custom_openai_api:
+                                provider = LLMProvider.OAICOMPAT
+                                cleaned_model_name_from_func = custom_model_value
+                                agent_loop_type = AgentLoop.UITARS
+                                print(f"Using OAICOMPAT provider for custom UITARS model: {custom_model_value}")
+                            else:
+                                provider = LLMProvider.MLXVLM
+                                cleaned_model_name_from_func = model_string_to_analyze
+                                agent_loop_type = AgentLoop.UITARS
+                                print(f"Using MLXVLM provider for UITARS model: {model_string_to_analyze}")
                         # Special case for Ollama custom model
                         elif is_custom_ollama and agent_loop_choice == "OMNI":
                             provider = LLMProvider.OLLAMA
@@ -1046,8 +1052,8 @@ if __name__ == "__main__":
                             else cleaned_model_name_from_func
                         )
 
-                        # Determine if OAICOMPAT should be used (only for OpenAI compatible API custom model)
-                        is_oaicompat = is_custom_openai_api and agent_loop_choice != "UITARS"
+                        # Determine if OAICOMPAT should be used (for OpenAI compatible API custom model)
+                        is_oaicompat = is_custom_openai_api
 
                         # Get API key based on provider determined by get_provider_and_model
                         if is_oaicompat and custom_api_key:

{cua_agent-0.1.40.dist-info → cua_agent-0.1.42.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.1.40
+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.40.dist-info → cua_agent-0.1.42.dist-info}/RECORD

@@ -1,7 +1,7 @@
 agent/__init__.py,sha256=guFGtorDBF6R5hVep0Bvci3_sUJfBlcsq9ss5Kwrej8,1484
 agent/core/__init__.py,sha256=7DhJ_6KKooM6uTmDIlumCnd7OFcU67BYIIR1dpIYUB0,506
 agent/core/agent.py,sha256=HUfBe7Uam3TObAmf6KH0GDKuNCNunNmmMcuxS7aZg0Q,8332
-agent/core/base.py,sha256=2sg8B2VqUKImRlkLTNj5lx-Oarlu7_GoMR6MbNzSY9Q,8078
+agent/core/base.py,sha256=AiSjnBAcHhZIca4KWBP1vQRE3HyikAPkr4Ij9WDevZQ,8374
 agent/core/callbacks.py,sha256=FKAxyajJ-ZJ5SxNXoupNcrm0GYBgjOjJEsStqst0EAk,6453
 agent/core/experiment.py,sha256=Ywj6q3JZFDKicfPuQsDl0vSN55HS7-Cnk3u3EcUCKe8,8866
 agent/core/factory.py,sha256=zzlCdibctqhf8Uta-SrvE-G7h59wAw-7SGhHiGvS9GY,4608
@@ -25,7 +25,7 @@ agent/providers/anthropic/api/logging.py,sha256=vHpwkIyOZdkSTVIH4ycbBPd4a_rzhP7O
 agent/providers/anthropic/api_handler.py,sha256=pWXcqDs0ruviDhRNRrz5Ac9ZH4yDv6ZlwpeG3a42cDg,5206
 agent/providers/anthropic/callbacks/__init__.py,sha256=PciBb6Z6MKSwfXqDjU3pV_0FS4MOn_Np_A7_skD-6dA,104
 agent/providers/anthropic/callbacks/manager.py,sha256=euIah5yiM8nhisN-RWXewo6v0WQr0c-FbMBO04r6dJk,1865
-agent/providers/anthropic/loop.py,sha256=A9ce3q5a4SQfOKVUk0En4x6fuc8a8s9wXFlQg4ypIkg,20394
+agent/providers/anthropic/loop.py,sha256=Sepfo0b0oQT98xd3Sv2S7Xc81bfU7L4_Zv3VTiapKkg,21661
 agent/providers/anthropic/prompts.py,sha256=nHFfgPrfvnWrEdVP7EUBGUHAI85D2X9HeZirk9EwncU,1941
 agent/providers/anthropic/response_handler.py,sha256=ZTprV4NTP9Eb9jQ7QgEKZBX0L6rMj5nqBRiE3Zfws8I,8008
 agent/providers/anthropic/tools/__init__.py,sha256=JyZwuVtPUnZwRSZBSCdQv9yxbLCsygm3l8Ywjjt9qTQ,661
@@ -47,7 +47,7 @@ agent/providers/omni/clients/ollama.py,sha256=PmR5EhU9Mi43_o5mZN36XcpiGKp5HbQwlX
 agent/providers/omni/clients/openai.py,sha256=iTSYWEJEM8INFPGJMiUVs8rFn0781XF_ofRkd7NT3gk,5920
 agent/providers/omni/clients/utils.py,sha256=Ani9CVVBm_J2Dl51WG6p1GVuoI6cq8scISrG0pmQ37o,688
 agent/providers/omni/image_utils.py,sha256=wejhWb36yqedsPnLFTFwk2wth8a6txfVWSg4EaNrRdA,908
-agent/providers/omni/loop.py,sha256=3eL80w2btw7Gt9FpvJjZeHf97fnDSzYpTi1hBjzRjIk,40929
+agent/providers/omni/loop.py,sha256=ecsDEUm9KPeBJVznXfWVyw4_YhbbslEjDkMMKDzUDwI,42973
 agent/providers/omni/parser.py,sha256=REpQwlwvY1z_N8wbMj6GhOeTiiWVWHhVja_LOxgzbks,11734
 agent/providers/omni/prompts.py,sha256=Mupjy0bUwBjcAeLXpE1r1jisYPSlhwsp-IXJKEKrEtw,3779
 agent/providers/omni/tools/__init__.py,sha256=IC1cMEDoR2ljGcNNthzBRF_VtnDbRL5qvHJWErtNp98,774
@@ -58,7 +58,7 @@ agent/providers/omni/tools/manager.py,sha256=UhtasaxGcmkxtz-bP1UJ1a4xdYnD3Cv8Pbt
 agent/providers/omni/utils.py,sha256=Ikp6ONL1HO637o3KDtv5yv6q-4uIWAzMSQDvGetWXC8,8724
 agent/providers/openai/__init__.py,sha256=8DS6YNZp42NLCacwXsfRaghyczaOCVovX8TgzXUZf_o,165
 agent/providers/openai/api_handler.py,sha256=L1K56dR1j4JsX1sX4OFYeKoCUMM25Fwj2y9nqv8oOhw,17736
-agent/providers/openai/loop.py,sha256=yJrFBy6uTlYf-MMcq_Q-EKQxbdBrEueJIeMhVo5OV1E,19344
+agent/providers/openai/loop.py,sha256=y0oLPYMjdjm0g5sCg3zF28FnMLaO2JdN6FiDdVK_77Q,20602
 agent/providers/openai/response_handler.py,sha256=K8v_92uSr9R74Y5INY4naeEZZZm35CLIl4h74MBZhsw,7953
 agent/providers/openai/tools/__init__.py,sha256=-KbHMWcd2OVTk5RYQ3ACBEMygwbH-VW6n_98p0lwM4A,344
 agent/providers/openai/tools/base.py,sha256=Np_BC9Cm6TslK99etE9hVTtsBlcEaGhoNCK3NXdB_Lw,2474
@@ -70,7 +70,7 @@ agent/providers/uitars/__init__.py,sha256=sq5OMVJP9E_sok9tIiKJreGkjmNWXPMObjPTCl
 agent/providers/uitars/clients/base.py,sha256=5w8Ajmq1JiPyUQJUAq1lSkfpA8_Ts80NQiDxPMTtQrI,948
 agent/providers/uitars/clients/mlxvlm.py,sha256=lMnN6ecMmWHf_l7khJ2iJHHvT7PE4XagUjrWhB0zEhc,10893
 agent/providers/uitars/clients/oaicompat.py,sha256=uYjwrGCVpFi8wj4kcaJ905ABiY6ksJZXaLlM61B2DUA,8907
-agent/providers/uitars/loop.py,sha256=4-cgQteixPy03vp7xWezd6jWpuPkBmlLS3tizaOmd0U,23494
+agent/providers/uitars/loop.py,sha256=cyzLGg5BEPOLHJk_gnoSLEy2BFnBXVgGry3TUvMX0_A,25467
 agent/providers/uitars/prompts.py,sha256=_pQNd438mFpZKZT0aMl6Bd0_GgQxuy9y08kQAMPi9UM,2536
 agent/providers/uitars/tools/__init__.py,sha256=0hc3W6u5TvcXYztYKIyve_C2G3XMfwt_y7grmH0ZHC0,29
 agent/providers/uitars/tools/computer.py,sha256=TeIg_aCtMroxWOBJEiYY_YI4krW_C3pYu51tgGsVUYU,11808
@@ -79,8 +79,8 @@ agent/providers/uitars/utils.py,sha256=493STTEEJcVhVbQgR0e8rNTI1DjkxUx8IgIv3wkJ1
 agent/telemetry.py,sha256=pVGxbj0ewnvq4EGj28CydN4a1iOfvZR_XKL3vIOqhOM,390
 agent/ui/__init__.py,sha256=ohhxJLBin6k1hl5sKcmBST8mgh23WXgAXz3pN4f470E,45
 agent/ui/gradio/__init__.py,sha256=ANKZhv1HqsLheWbLVBlyRQ7Q5qGeXuPi5jDs8vu-ZMo,579
-agent/ui/gradio/app.py,sha256=M_pqSiN40F1u-8luBdqvTJxQFGzqd0WPsSz8APLbPus,67826
-cua_agent-0.1.40.dist-info/METADATA,sha256=7ZNaoMyIfm5l_vnWUYNhcNxplhPj9_ocnfWF9sgN9yk,12689
-cua_agent-0.1.40.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-cua_agent-0.1.40.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-cua_agent-0.1.40.dist-info/RECORD,,
+agent/ui/gradio/app.py,sha256=-ccsE6LrXFfxnPeMlEqm49QGvdjCgm-l6TudZZEM9r0,68241
+cua_agent-0.1.42.dist-info/METADATA,sha256=rW2ijMboeUh6BKDZopZOuR01Fo2946tT0ftyDlwXL0k,12689
+cua_agent-0.1.42.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
+cua_agent-0.1.42.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+cua_agent-0.1.42.dist-info/RECORD,,