cua-agent 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

agent/providers/anthropic/loop.py

@@ -17,6 +17,7 @@ from anthropic.types.beta import (
     BetaTextBlock,
     BetaTextBlockParam,
     BetaToolUseBlockParam,
+    BetaContentBlockParam,
 )
 
 # Computer
@@ -24,12 +25,12 @@ from computer import Computer
 
 # Base imports
 from ...core.loop import BaseLoop
-from ...core.messages import ImageRetentionConfig
+from ...core.messages import ImageRetentionConfig as CoreImageRetentionConfig
 
 # Anthropic provider-specific imports
 from .api.client import AnthropicClientFactory, BaseAnthropicClient
 from .tools.manager import ToolManager
-from .messages.manager import MessageManager
+from .messages.manager import MessageManager, ImageRetentionConfig
 from .callbacks.manager import CallbackManager
 from .prompts import SYSTEM_PROMPT
 from .types import LLMProvider
@@ -48,8 +49,8 @@ class AnthropicLoop(BaseLoop):
     def __init__(
         self,
         api_key: str,
+        computer: Computer,
         model: str = "claude-3-7-sonnet-20250219",  # Fixed model
-        computer: Optional[Computer] = None,
         only_n_most_recent_images: Optional[int] = 2,
         base_dir: Optional[str] = "trajectories",
         max_retries: int = 3,
@@ -69,7 +70,7 @@ class AnthropicLoop(BaseLoop):
             retry_delay: Delay between retries in seconds
             save_trajectory: Whether to save trajectory data
         """
-        # Initialize base class
+        # Initialize base class with core config
         super().__init__(
             computer=computer,
             model=model,
@@ -93,8 +94,8 @@ class AnthropicLoop(BaseLoop):
         self.message_manager = None
         self.callback_manager = None
 
-        # Configure image retention
-        self.image_retention_config = ImageRetentionConfig(
+        # Configure image retention with core config
+        self.image_retention_config = CoreImageRetentionConfig(
             num_images_to_keep=only_n_most_recent_images
         )
 
@@ -113,7 +114,7 @@ class AnthropicLoop(BaseLoop):
 
             # Initialize message manager
             self.message_manager = MessageManager(
-                ImageRetentionConfig(
+                image_retention_config=ImageRetentionConfig(
                     num_images_to_keep=self.only_n_most_recent_images, enable_caching=True
                 )
             )
@@ -250,6 +251,10 @@ class AnthropicLoop(BaseLoop):
                 await self._process_screen(parsed_screen, self.message_history)
 
                 # Prepare messages and make API call
+                if self.message_manager is None:
+                    raise RuntimeError(
+                        "Message manager not initialized. Call initialize_client() first."
+                    )
                 prepared_messages = self.message_manager.prepare_messages(
                     cast(List[BetaMessageParam], self.message_history.copy())
                 )
@@ -257,7 +262,7 @@ class AnthropicLoop(BaseLoop):
                 # Create new turn directory for this API call
                 self._create_turn_dir()
 
-                # Make API call
+                # Use _make_api_call instead of direct client call to ensure logging
                 response = await self._make_api_call(prepared_messages)
 
                 # Handle the response
@@ -287,6 +292,11 @@ class AnthropicLoop(BaseLoop):
         Returns:
             API response
         """
+        if self.client is None:
+            raise RuntimeError("Client not initialized. Call initialize_client() first.")
+        if self.tool_manager is None:
+            raise RuntimeError("Tool manager not initialized. Call initialize_client() first.")
+
         last_error = None
 
         for attempt in range(self.max_retries):
@@ -297,6 +307,7 @@ class AnthropicLoop(BaseLoop):
                     "max_tokens": self.max_tokens,
                     "system": SYSTEM_PROMPT,
                 }
+                # Let ExperimentManager handle sanitization
                 self._log_api_call("request", request_data)
 
                 # Setup betas and system
@@ -320,7 +331,7 @@ class AnthropicLoop(BaseLoop):
                     betas=betas,
                 )
 
-                # Log success response
+                # Let ExperimentManager handle sanitization
                 self._log_api_call("response", request_data, response)
 
                 return response
@@ -365,25 +376,38 @@ class AnthropicLoop(BaseLoop):
                 }
             )
 
+            if self.callback_manager is None:
+                raise RuntimeError(
+                    "Callback manager not initialized. Call initialize_client() first."
+                )
+
             # Handle tool use blocks and collect results
             tool_result_content = []
             for content_block in response_params:
                 # Notify callback of content
-                self.callback_manager.on_content(content_block)
+                self.callback_manager.on_content(cast(BetaContentBlockParam, content_block))
 
                 # Handle tool use
                 if content_block.get("type") == "tool_use":
+                    if self.tool_manager is None:
+                        raise RuntimeError(
+                            "Tool manager not initialized. Call initialize_client() first."
+                        )
                     result = await self.tool_manager.execute_tool(
                         name=content_block["name"],
                         tool_input=cast(Dict[str, Any], content_block["input"]),
                     )
 
                     # Create tool result and add to content
-                    tool_result = self._make_tool_result(result, content_block["id"])
+                    tool_result = self._make_tool_result(
+                        cast(ToolResult, result), content_block["id"]
+                    )
                     tool_result_content.append(tool_result)
 
                     # Notify callback of tool result
-                    self.callback_manager.on_tool_result(result, content_block["id"])
+                    self.callback_manager.on_tool_result(
+                        cast(ToolResult, result), content_block["id"]
+                    )
 
             # If no tool results, we're done
             if not tool_result_content:
@@ -495,13 +519,13 @@ class AnthropicLoop(BaseLoop):
             result_text = f"<s>{result.system}</s>\n{result_text}"
         return result_text
 
-    def _handle_content(self, content: Dict[str, Any]) -> None:
+    def _handle_content(self, content: BetaContentBlockParam) -> None:
         """Handle content updates from the assistant."""
         if content.get("type") == "text":
-            text = content.get("text", "")
+            text_content = cast(BetaTextBlockParam, content)
+            text = text_content["text"]
             if text == "<DONE>":
                 return
-
             logger.info(f"Assistant: {text}")
 
     def _handle_tool_result(self, result: ToolResult, tool_id: str) -> None:
@@ -517,5 +541,10 @@ class AnthropicLoop(BaseLoop):
         """Handle API interactions."""
         if error:
             logger.error(f"API error: {error}")
+            self._log_api_call("error", request, error=error)
         else:
             logger.debug(f"API request: {request}")
+            if response:
+                self._log_api_call("response", request, response)
+            else:
+                self._log_api_call("request", request)

agent/providers/anthropic/messages/manager.py

@@ -90,7 +90,9 @@ class MessageManager:
                     blocks_with_cache_control += 1
                     # Add cache control to the last content block only
                     if content and len(content) > 0:
-                        content[-1]["cache_control"] = {"type": "ephemeral"}
+                        content[-1]["cache_control"] = BetaCacheControlEphemeralParam(
+                            type="ephemeral"
+                        )
                 else:
                     # Remove any existing cache control
                     if content and len(content) > 0:

agent/providers/anthropic/tools/base.py

@@ -6,7 +6,7 @@ from typing import Any, Dict
 
 from anthropic.types.beta import BetaToolUnionParam
 
-from ....core.tools.base import BaseTool, ToolError, ToolResult, ToolFailure, CLIResult
+from ....core.tools.base import BaseTool
 
 
 class BaseAnthropicTool(BaseTool, metaclass=ABCMeta):

agent/providers/anthropic/tools/collection.py

@@ -1,6 +1,6 @@
 """Collection classes for managing multiple tools."""
 
-from typing import Any
+from typing import Any, cast
 
 from anthropic.types.beta import BetaToolUnionParam
 
@@ -22,7 +22,7 @@ class ToolCollection:
     def to_params(
         self,
     ) -> list[BetaToolUnionParam]:
-        return [tool.to_params() for tool in self.tools]
+        return cast(list[BetaToolUnionParam], [tool.to_params() for tool in self.tools])
 
     async def run(self, *, name: str, tool_input: dict[str, Any]) -> ToolResult:
         tool = self.tool_map.get(name)

agent/providers/anthropic/tools/computer.py

@@ -61,9 +61,9 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
 
     name: Literal["computer"] = "computer"
     api_type: Literal["computer_20250124"] = "computer_20250124"
-    width: int | None
-    height: int | None
-    display_num: int | None
+    width: int | None = None
+    height: int | None = None
+    display_num: int | None = None
     computer: Computer  # The CUA Computer instance
     logger = logging.getLogger(__name__)
 
@@ -106,6 +106,7 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
         display_size = await self.computer.interface.get_screen_size()
         self.width = display_size["width"]
         self.height = display_size["height"]
+        assert isinstance(self.width, int) and isinstance(self.height, int)
         self.logger.info(f"Initialized screen dimensions to {self.width}x{self.height}")
 
     async def __call__(
@@ -120,6 +121,8 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
             # Ensure dimensions are initialized
             if self.width is None or self.height is None:
                 await self.initialize_dimensions()
+                if self.width is None or self.height is None:
+                    raise ToolError("Failed to initialize screen dimensions")
         except Exception as e:
             raise ToolError(f"Failed to initialize dimensions: {e}")
 
@@ -147,7 +150,10 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
                     self.logger.info(
                         f"Scaling image from {pre_img.size} to {self.width}x{self.height} to match screen dimensions"
                     )
-                    pre_img = pre_img.resize((self.width, self.height), Image.Resampling.LANCZOS)
+                    if not isinstance(self.width, int) or not isinstance(self.height, int):
+                        raise ToolError("Screen dimensions must be integers")
+                    size = (int(self.width), int(self.height))
+                    pre_img = pre_img.resize(size, Image.Resampling.LANCZOS)
 
                 self.logger.info(f"  Current dimensions: {pre_img.width}x{pre_img.height}")
 
@@ -160,15 +166,7 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
                     await self.computer.interface.move_cursor(x, y)
                     # Then perform drag operation - check if drag_to exists or we need to use other methods
                     try:
-                        if hasattr(self.computer.interface, "drag_to"):
-                            await self.computer.interface.drag_to(x, y)
-                        else:
-                            # Alternative approach: press mouse down, move, release
-                            await self.computer.interface.mouse_down()
-                            await asyncio.sleep(0.2)
-                            await self.computer.interface.move_cursor(x, y)
-                            await asyncio.sleep(0.2)
-                            await self.computer.interface.mouse_up()
+                        await self.computer.interface.drag_to(x, y)
                     except Exception as e:
                         self.logger.error(f"Error during drag operation: {str(e)}")
                         raise ToolError(f"Failed to perform drag: {str(e)}")
@@ -214,9 +212,10 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
                         self.logger.info(
                             f"Scaling image from {pre_img.size} to {self.width}x{self.height} to match screen dimensions"
                         )
-                        pre_img = pre_img.resize(
-                            (self.width, self.height), Image.Resampling.LANCZOS
-                        )
+                        if not isinstance(self.width, int) or not isinstance(self.height, int):
+                            raise ToolError("Screen dimensions must be integers")
+                        size = (int(self.width), int(self.height))
+                        pre_img = pre_img.resize(size, Image.Resampling.LANCZOS)
                         # Save the scaled image back to bytes
                         buffer = io.BytesIO()
                         pre_img.save(buffer, format="PNG")
@@ -275,9 +274,10 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
                         self.logger.info(
                             f"Scaling image from {pre_img.size} to {self.width}x{self.height}"
                         )
-                        pre_img = pre_img.resize(
-                            (self.width, self.height), Image.Resampling.LANCZOS
-                        )
+                        if not isinstance(self.width, int) or not isinstance(self.height, int):
+                            raise ToolError("Screen dimensions must be integers")
+                        size = (int(self.width), int(self.height))
+                        pre_img = pre_img.resize(size, Image.Resampling.LANCZOS)
 
                     # Perform the click action
                     if action == "left_click":
@@ -335,7 +335,10 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
                     self.logger.info(
                         f"Scaling image from {pre_img.size} to {self.width}x{self.height}"
                     )
-                    pre_img = pre_img.resize((self.width, self.height), Image.Resampling.LANCZOS)
+                    if not isinstance(self.width, int) or not isinstance(self.height, int):
+                        raise ToolError("Screen dimensions must be integers")
+                    size = (int(self.width), int(self.height))
+                    pre_img = pre_img.resize(size, Image.Resampling.LANCZOS)
 
                 if action == "key":
                     # Special handling for page up/down on macOS
@@ -365,7 +368,7 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
                         # Handle single key press
                         self.logger.info(f"Pressing key: {text}")
                         try:
-                            await self.computer.interface.press(text)
+                            await self.computer.interface.press_key(text)
                             output_text = text
                         except ValueError as e:
                             raise ToolError(f"Invalid key: {text}. {str(e)}")
@@ -442,7 +445,10 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
                         self.logger.info(
                             f"Scaling image from {img.size} to {self.width}x{self.height}"
                         )
-                        img = img.resize((self.width, self.height), Image.Resampling.LANCZOS)
+                        if not isinstance(self.width, int) or not isinstance(self.height, int):
+                            raise ToolError("Screen dimensions must be integers")
+                        size = (int(self.width), int(self.height))
+                        img = img.resize(size, Image.Resampling.LANCZOS)
                         buffer = io.BytesIO()
                         img.save(buffer, format="PNG")
                         screenshot = buffer.getvalue()
@@ -451,7 +457,8 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
 
                 elif action == "cursor_position":
                     pos = await self.computer.interface.get_cursor_position()
-                    return ToolResult(output=f"X={int(pos[0])},Y={int(pos[1])}")
+                    x, y = pos  # Unpack the tuple
+                    return ToolResult(output=f"X={int(x)},Y={int(y)}")
 
             except Exception as e:
                 self.logger.error(f"Error during {action} action: {str(e)}")
@@ -517,7 +524,10 @@ class ComputerTool(BaseComputerTool, BaseAnthropicTool):
             # Scale image if needed
             if img.size != (self.width, self.height):
                 self.logger.info(f"Scaling image from {img.size} to {self.width}x{self.height}")
-                img = img.resize((self.width, self.height), Image.Resampling.LANCZOS)
+                if not isinstance(self.width, int) or not isinstance(self.height, int):
+                    raise ToolError("Screen dimensions must be integers")
+                size = (int(self.width), int(self.height))
+                img = img.resize(size, Image.Resampling.LANCZOS)
                 buffer = io.BytesIO()
                 img.save(buffer, format="PNG")
                 screenshot = buffer.getvalue()

agent/providers/anthropic/tools/manager.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, cast
 from anthropic.types.beta import BetaToolUnionParam
 from computer.computer import Computer
 
@@ -37,7 +37,7 @@ class ToolManager(BaseToolManager):
         """Get tool parameters for Anthropic API calls."""
         if self.tools is None:
             raise RuntimeError("Tools not initialized. Call initialize() first.")
-        return self.tools.to_params()
+        return cast(List[BetaToolUnionParam], self.tools.to_params())
 
     async def execute_tool(self, name: str, tool_input: dict[str, Any]) -> ToolResult:
         """Execute a tool with the given input.

agent/providers/omni/experiment.py

@@ -126,15 +126,18 @@ class ExperimentManager:
         # Since we no longer want to use the images/ folder, we'll skip this functionality
         return
 
-    def save_screenshot(self, img_base64: str, action_type: str = "") -> None:
+    def save_screenshot(self, img_base64: str, action_type: str = "") -> Optional[str]:
         """Save a screenshot to the experiment directory.
 
         Args:
             img_base64: Base64 encoded screenshot
             action_type: Type of action that triggered the screenshot
+
+        Returns:
+            Optional[str]: Path to the saved screenshot, or None if saving failed
         """
         if not self.current_turn_dir:
-            return
+            return None
 
         try:
             # Increment screenshot counter

agent/providers/omni/loop.py

@@ -13,6 +13,7 @@ import asyncio
 from httpx import ConnectError, ReadTimeout
 import shutil
 import copy
+from typing import cast
 
 from .parser import OmniParser, ParseResult, ParserMetadata, UIElement
 from ...core.loop import BaseLoop
@@ -182,8 +183,6 @@ class OmniLoop(BaseLoop):
 
             if self.provider == LLMProvider.OPENAI:
                 self.client = OpenAIClient(api_key=self.api_key, model=self.model)
-            elif self.provider == LLMProvider.GROQ:
-                self.client = GroqClient(api_key=self.api_key, model=self.model)
             elif self.provider == LLMProvider.ANTHROPIC:
                 self.client = AnthropicClient(
                     api_key=self.api_key,
@@ -329,10 +328,15 @@ class OmniLoop(BaseLoop):
         raise RuntimeError(error_message)
 
     async def _handle_response(
-        self, response: Any, messages: List[Dict[str, Any]], parsed_screen: Dict[str, Any]
+        self, response: Any, messages: List[Dict[str, Any]], parsed_screen: ParseResult
     ) -> Tuple[bool, bool]:
         """Handle API response.
 
+        Args:
+            response: API response
+            messages: List of messages to update
+            parsed_screen: Current parsed screen information
+
         Returns:
             Tuple of (should_continue, action_screenshot_saved)
         """
@@ -394,7 +398,9 @@ class OmniLoop(BaseLoop):
 
                             try:
                                 # Execute action with current parsed screen info
-                                await self._execute_action(parsed_content, parsed_screen)
+                                await self._execute_action(
+                                    parsed_content, cast(ParseResult, parsed_screen)
+                                )
                                 action_screenshot_saved = True
                             except Exception as e:
                                 logger.error(f"Error executing action: {str(e)}")
@@ -463,7 +469,7 @@ class OmniLoop(BaseLoop):
 
                 try:
                     # Execute action with current parsed screen info
-                    await self._execute_action(parsed_content, parsed_screen)
+                    await self._execute_action(parsed_content, cast(ParseResult, parsed_screen))
                     action_screenshot_saved = True
                 except Exception as e:
                     logger.error(f"Error executing action: {str(e)}")
@@ -488,7 +494,7 @@ class OmniLoop(BaseLoop):
 
                 try:
                     # Execute action with current parsed screen info
-                    await self._execute_action(content, parsed_screen)
+                    await self._execute_action(content, cast(ParseResult, parsed_screen))
                     action_screenshot_saved = True
                 except Exception as e:
                     logger.error(f"Error executing action: {str(e)}")

agent/providers/omni/parser.py

@@ -122,8 +122,9 @@ class OmniParser:
             # Create a minimal valid result for error cases
             return ParseResult(
                 elements=[],
+                screen_info=None,
                 annotated_image_base64="",
-                parsed_content_list=[f"Error: {str(e)}"],
+                parsed_content_list=[{"error": str(e)}],
                 metadata=ParserMetadata(
                     image_size=(0, 0),
                     num_icons=0,

agent/providers/omni/tools/__init__.py

@@ -2,7 +2,6 @@
 
 from .bash import OmniBashTool
 from .computer import OmniComputerTool
-from .edit import OmniEditTool
 from .manager import OmniToolManager
 
 __all__ = [

agent/providers/omni/tools/computer.py

@@ -177,7 +177,7 @@ class OmniComputerTool(BaseComputerTool):
                     keys = text.split("+")
                     await self.computer.interface.hotkey(*keys)
                 else:
-                    await self.computer.interface.press(text)
+                    await self.computer.interface.press_key(text)
 
                 # Take screenshot after action
                 screenshot = await self.computer.interface.screenshot()
@@ -188,7 +188,8 @@ class OmniComputerTool(BaseComputerTool):
                 )
             elif action == "cursor_position":
                 pos = await self.computer.interface.get_cursor_position()
-                return ToolResult(output=f"X={int(pos[0])},Y={int(pos[1])}")
+                x, y = pos
+                return ToolResult(output=f"X={int(x)},Y={int(y)}")
             elif action == "scroll":
                 if direction == "down":
                     self.logger.info(f"Scrolling down, amount: {amount}")

agent/providers/omni/tools/manager.py

@@ -10,7 +10,6 @@ from ....core.tools.collection import ToolCollection
 
 from .bash import OmniBashTool
 from .computer import OmniComputerTool
-from .edit import OmniEditTool
 
 
 class ProviderType(Enum):
@@ -35,11 +34,10 @@ class OmniToolManager(BaseToolManager):
         # Initialize tools
         self.computer_tool = OmniComputerTool(self.computer)
         self.bash_tool = OmniBashTool(self.computer)
-        self.edit_tool = OmniEditTool(self.computer)
 
     def _initialize_tools(self) -> ToolCollection:
         """Initialize all available tools."""
-        return ToolCollection(self.computer_tool, self.bash_tool, self.edit_tool)
+        return ToolCollection(self.computer_tool, self.bash_tool)
 
     async def _initialize_tools_specific(self) -> None:
         """Initialize provider-specific tool requirements."""

agent/providers/omni/utils.py

@@ -96,7 +96,7 @@ def compress_image_base64(
             # Resize image
             new_width = int(img.width * scale_factor)
             new_height = int(img.height * scale_factor)
-            current_img = img.resize((new_width, new_height), Image.LANCZOS)
+            current_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
 
             # Try with reduced size and quality
             buffer = io.BytesIO()
@@ -130,7 +130,9 @@ def compress_image_base64(
 
         # Last resort: Use minimum quality and size
         buffer = io.BytesIO()
-        smallest_img = img.resize((int(img.width * 0.5), int(img.height * 0.5)), Image.LANCZOS)
+        smallest_img = img.resize(
+            (int(img.width * 0.5), int(img.height * 0.5)), Image.Resampling.LANCZOS
+        )
         # Convert to RGB if necessary
         if smallest_img.mode in ("RGBA", "LA") or (
             smallest_img.mode == "P" and "transparency" in smallest_img.info

agent/types/__init__.py

@@ -1,23 +1,20 @@
 """Type definitions for the agent package."""
 
-from .base import Provider, HostConfig, TaskResult, Annotation
+from .base import HostConfig, TaskResult, Annotation
 from .messages import Message, Request, Response, StepMessage, DisengageMessage
 from .tools import ToolInvocation, ToolInvocationState, ClientAttachment, ToolResult
 
 __all__ = [
     # Base types
-    "Provider",
     "HostConfig",
     "TaskResult",
     "Annotation",
-    
     # Message types
     "Message",
     "Request",
     "Response",
     "StepMessage",
     "DisengageMessage",
-    
     # Tool types
     "ToolInvocation",
     "ToolInvocationState",

agent/types/base.py

@@ -5,17 +5,6 @@ from typing import Dict, Any
 from pydantic import BaseModel, ConfigDict
 
 
-class Provider(str, Enum):
-    """Available AI providers."""
-
-    UNKNOWN = "unknown"  # Default provider for base class
-    ANTHROPIC = "anthropic"
-    OPENAI = "openai"
-    OLLAMA = "ollama"
-    OMNI = "omni"
-    GROQ = "groq"
-
-
 class HostConfig(BaseModel):
     """Host configuration."""
 
@@ -48,6 +37,5 @@ class AgentLoop(Enum):
     """Enumeration of available loop types."""
 
     ANTHROPIC = auto()  # Anthropic implementation
-    OPENAI = auto()  # OpenAI implementation
     OMNI = auto()  # OmniLoop implementation
     # Add more loop types as needed

{cua_agent-0.1.5.dist-info → cua_agent-0.1.6.dist-info}/METADATA

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