cua-agent 0.1.0__py3-none-any.whl

agent/providers/anthropic/tools/edit.py

@@ -0,0 +1,326 @@
+from collections import defaultdict
+from pathlib import Path
+from typing import Literal, get_args, Dict, Any
+from computer.computer import Computer
+
+from .base import BaseAnthropicTool, CLIResult, ToolError, ToolResult
+from ....core.tools.edit import BaseEditTool
+from .run import maybe_truncate
+
+Command = Literal[
+    "view",
+    "create",
+    "str_replace",
+    "insert",
+    "undo_edit",
+]
+SNIPPET_LINES: int = 4
+
+
+class EditTool(BaseEditTool, BaseAnthropicTool):
+    """
+    An filesystem editor tool that allows the agent to view, create, and edit files.
+    The tool parameters are defined by Anthropic and are not editable.
+    """
+
+    api_type: Literal["text_editor_20250124"] = "text_editor_20250124"
+    name: Literal["str_replace_editor"] = "str_replace_editor"
+    _timeout: float = 30.0  # seconds
+
+    def __init__(self, computer: Computer):
+        """Initialize the edit tool.
+
+        Args:
+            computer: Computer instance for file operations
+        """
+        # Initialize the base edit tool first
+        BaseEditTool.__init__(self, computer)
+        # Then initialize the Anthropic tool
+        BaseAnthropicTool.__init__(self)
+
+        # Edit history for the current session
+        self.edit_history = defaultdict(list)
+
+    async def __call__(
+        self,
+        *,
+        command: Command,
+        path: str,
+        file_text: str | None = None,
+        view_range: list[int] | None = None,
+        old_str: str | None = None,
+        new_str: str | None = None,
+        insert_line: int | None = None,
+        **kwargs,
+    ):
+        _path = Path(path)
+        await self.validate_path(command, _path)
+
+        if command == "view":
+            return await self.view(_path, view_range)
+        elif command == "create":
+            if file_text is None:
+                raise ToolError("Parameter `file_text` is required for command: create")
+            await self.write_file(_path, file_text)
+            self.edit_history[_path].append(file_text)
+            return ToolResult(output=f"File created successfully at: {_path}")
+        elif command == "str_replace":
+            if old_str is None:
+                raise ToolError("Parameter `old_str` is required for command: str_replace")
+            return await self.str_replace(_path, old_str, new_str)
+        elif command == "insert":
+            if insert_line is None:
+                raise ToolError("Parameter `insert_line` is required for command: insert")
+            if new_str is None:
+                raise ToolError("Parameter `new_str` is required for command: insert")
+            return await self.insert(_path, insert_line, new_str)
+        elif command == "undo_edit":
+            return await self.undo_edit(_path)
+
+        raise ToolError(
+            f'Unrecognized command {command}. The allowed commands for the {self.name} tool are: {", ".join(get_args(Command))}'
+        )
+
+    async def validate_path(self, command: str, path: Path):
+        """Check that the path/command combination is valid."""
+        # Check if its an absolute path
+        if not path.is_absolute():
+            suggested_path = Path("") / path
+            raise ToolError(
+                f"The path {path} is not an absolute path, it should start with `/`. Maybe you meant {suggested_path}?"
+            )
+
+        # Check if path exists using bash commands
+        try:
+            result = await self.computer.interface.run_command(
+                f'[ -e "{str(path)}" ] && echo "exists" || echo "not exists"'
+            )
+            exists = result[0].strip() == "exists"
+
+            if exists:
+                result = await self.computer.interface.run_command(
+                    f'[ -d "{str(path)}" ] && echo "dir" || echo "file"'
+                )
+                is_dir = result[0].strip() == "dir"
+            else:
+                is_dir = False
+
+            # Check path validity
+            if not exists and command != "create":
+                raise ToolError(f"The path {path} does not exist. Please provide a valid path.")
+            if exists and command == "create":
+                raise ToolError(
+                    f"File already exists at: {path}. Cannot overwrite files using command `create`."
+                )
+            if is_dir and command != "view":
+                raise ToolError(
+                    f"The path {path} is a directory and only the `view` command can be used on directories"
+                )
+        except Exception as e:
+            raise ToolError(f"Failed to validate path: {str(e)}")
+
+    async def view(self, path: Path, view_range: list[int] | None = None):
+        """Implement the view command"""
+        try:
+            # Check if path is a directory
+            result = await self.computer.interface.run_command(
+                f'[ -d "{str(path)}" ] && echo "dir" || echo "file"'
+            )
+            is_dir = result[0].strip() == "dir"
+
+            if is_dir:
+                if view_range:
+                    raise ToolError(
+                        "The `view_range` parameter is not allowed when `path` points to a directory."
+                    )
+
+                # List directory contents using ls
+                result = await self.computer.interface.run_command(f'ls -la "{str(path)}"')
+                contents = result[0]
+                if contents:
+                    stdout = f"Here's the files and directories in {path}:\n{contents}\n"
+                else:
+                    stdout = f"Directory {path} is empty\n"
+                return CLIResult(output=stdout)
+
+            # Read file content using cat
+            file_content = await self.read_file(path)
+            init_line = 1
+
+            if view_range:
+                if len(view_range) != 2 or not all(isinstance(i, int) for i in view_range):
+                    raise ToolError("Invalid `view_range`. It should be a list of two integers.")
+
+                file_lines = file_content.split("\n")
+                n_lines_file = len(file_lines)
+                init_line, final_line = view_range
+
+                if init_line < 1 or init_line > n_lines_file:
+                    raise ToolError(
+                        f"Invalid `view_range`: {view_range}. Its first element `{init_line}` should be within the range of lines of the file: {[1, n_lines_file]}"
+                    )
+                if final_line > n_lines_file:
+                    raise ToolError(
+                        f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be smaller than the number of lines in the file: `{n_lines_file}`"
+                    )
+                if final_line != -1 and final_line < init_line:
+                    raise ToolError(
+                        f"Invalid `view_range`: {view_range}. Its second element `{final_line}` should be larger or equal than its first `{init_line}`"
+                    )
+
+                if final_line == -1:
+                    file_content = "\n".join(file_lines[init_line - 1 :])
+                else:
+                    file_content = "\n".join(file_lines[init_line - 1 : final_line])
+
+            return CLIResult(output=self._make_output(file_content, str(path), init_line=init_line))
+        except Exception as e:
+            raise ToolError(f"Failed to view path: {str(e)}")
+
+    async def str_replace(self, path: Path, old_str: str, new_str: str | None):
+        """Implement the str_replace command"""
+        # Read the file content
+        file_content = await self.read_file(path)
+        file_content = file_content.expandtabs()
+        old_str = old_str.expandtabs()
+        new_str = new_str.expandtabs() if new_str is not None else ""
+
+        # Check if old_str is unique in the file
+        occurrences = file_content.count(old_str)
+        if occurrences == 0:
+            raise ToolError(
+                f"No replacement was performed, old_str `{old_str}` did not appear verbatim in {path}."
+            )
+        elif occurrences > 1:
+            file_content_lines = file_content.split("\n")
+            lines = [idx + 1 for idx, line in enumerate(file_content_lines) if old_str in line]
+            raise ToolError(
+                f"No replacement was performed. Multiple occurrences of old_str `{old_str}` in lines {lines}. Please ensure it is unique"
+            )
+
+        # Replace old_str with new_str
+        new_file_content = file_content.replace(old_str, new_str)
+
+        # Write the new content to the file
+        await self.write_file(path, new_file_content)
+
+        # Save the content to history
+        self.edit_history[path].append(file_content)
+
+        # Create a snippet of the edited section
+        replacement_line = file_content.split(old_str)[0].count("\n")
+        start_line = max(0, replacement_line - SNIPPET_LINES)
+        end_line = replacement_line + SNIPPET_LINES + new_str.count("\n")
+        snippet = "\n".join(new_file_content.split("\n")[start_line : end_line + 1])
+
+        # Prepare the success message
+        success_msg = f"The file {path} has been edited. "
+        success_msg += self._make_output(snippet, f"a snippet of {path}", start_line + 1)
+        success_msg += "Review the changes and make sure they are as expected. Edit the file again if necessary."
+
+        return CLIResult(output=success_msg)
+
+    async def insert(self, path: Path, insert_line: int, new_str: str):
+        """Implement the insert command"""
+        file_text = await self.read_file(path)
+        file_text = file_text.expandtabs()
+        new_str = new_str.expandtabs()
+        file_text_lines = file_text.split("\n")
+        n_lines_file = len(file_text_lines)
+
+        if insert_line < 0 or insert_line > n_lines_file:
+            raise ToolError(
+                f"Invalid `insert_line` parameter: {insert_line}. It should be within the range of lines of the file: {[0, n_lines_file]}"
+            )
+
+        new_str_lines = new_str.split("\n")
+        new_file_text_lines = (
+            file_text_lines[:insert_line] + new_str_lines + file_text_lines[insert_line:]
+        )
+        snippet_lines = (
+            file_text_lines[max(0, insert_line - SNIPPET_LINES) : insert_line]
+            + new_str_lines
+            + file_text_lines[insert_line : insert_line + SNIPPET_LINES]
+        )
+
+        new_file_text = "\n".join(new_file_text_lines)
+        snippet = "\n".join(snippet_lines)
+
+        await self.write_file(path, new_file_text)
+        self.edit_history[path].append(file_text)
+
+        success_msg = f"The file {path} has been edited. "
+        success_msg += self._make_output(
+            snippet, "a snippet of the edited file", max(1, insert_line - SNIPPET_LINES + 1)
+        )
+        success_msg += "Review the changes and make sure they are as expected (correct indentation, no duplicate lines, etc). Edit the file again if necessary."
+        return CLIResult(output=success_msg)
+
+    async def undo_edit(self, path: Path):
+        """Implement the undo_edit command"""
+        if not self.edit_history[path]:
+            raise ToolError(f"No edit history found for {path}.")
+
+        old_text = self.edit_history[path].pop()
+        await self.write_file(path, old_text)
+
+        return CLIResult(
+            output=f"Last edit to {path} undone successfully. {self._make_output(old_text, str(path))}"
+        )
+
+    async def read_file(self, path: Path) -> str:
+        """Read the content of a file using cat command."""
+        try:
+            result = await self.computer.interface.run_command(f'cat "{str(path)}"')
+            if result[1]:  # If there's stderr output
+                raise ToolError(f"Error reading file: {result[1]}")
+            return result[0]
+        except Exception as e:
+            raise ToolError(f"Failed to read {path}: {str(e)}")
+
+    async def write_file(self, path: Path, content: str):
+        """Write content to a file using echo and redirection."""
+        try:
+            # Create parent directories if they don't exist
+            parent = path.parent
+            if parent != Path("/"):
+                await self.computer.interface.run_command(f'mkdir -p "{str(parent)}"')
+
+            # Write content to file using echo and heredoc to preserve formatting
+            cmd = f"""cat > "{str(path)}" << 'EOFCUA'
+{content}
+EOFCUA"""
+            result = await self.computer.interface.run_command(cmd)
+            if result[1]:  # If there's stderr output
+                raise ToolError(f"Error writing file: {result[1]}")
+        except Exception as e:
+            raise ToolError(f"Failed to write to {path}: {str(e)}")
+
+    def _make_output(
+        self,
+        file_content: str,
+        file_descriptor: str,
+        init_line: int = 1,
+        expand_tabs: bool = True,
+    ) -> str:
+        """Generate output for the CLI based on the content of a file."""
+        file_content = maybe_truncate(file_content)
+        if expand_tabs:
+            file_content = file_content.expandtabs()
+        file_content = "\n".join(
+            [f"{i + init_line:6}\t{line}" for i, line in enumerate(file_content.split("\n"))]
+        )
+        return (
+            f"Here's the result of running `cat -n` on {file_descriptor}:\n" + file_content + "\n"
+        )
+
+    def to_params(self) -> Dict[str, Any]:
+        """Convert tool to API parameters.
+
+        Returns:
+            Dictionary with tool parameters
+        """
+        return {
+            "name": self.name,
+            "type": self.api_type,
+        }

agent/providers/anthropic/tools/manager.py

@@ -0,0 +1,54 @@
+from typing import Any, Dict, List
+from anthropic.types.beta import BetaToolUnionParam
+from computer.computer import Computer
+
+from ....core.tools import BaseToolManager, ToolResult
+from ....core.tools.collection import ToolCollection
+
+from .bash import BashTool
+from .computer import ComputerTool
+from .edit import EditTool
+
+
+class ToolManager(BaseToolManager):
+    """Manages Anthropic-specific tool initialization and execution."""
+
+    def __init__(self, computer: Computer):
+        """Initialize the tool manager.
+
+        Args:
+            computer: Computer instance for computer-related tools
+        """
+        super().__init__(computer)
+        # Initialize Anthropic-specific tools
+        self.computer_tool = ComputerTool(self.computer)
+        self.bash_tool = BashTool(self.computer)
+        self.edit_tool = EditTool(self.computer)
+
+    def _initialize_tools(self) -> ToolCollection:
+        """Initialize all available tools."""
+        return ToolCollection(self.computer_tool, self.bash_tool, self.edit_tool)
+
+    async def _initialize_tools_specific(self) -> None:
+        """Initialize Anthropic-specific tool requirements."""
+        await self.computer_tool.initialize_dimensions()
+
+    def get_tool_params(self) -> List[BetaToolUnionParam]:
+        """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()
+
+    async def execute_tool(self, name: str, tool_input: dict[str, Any]) -> ToolResult:
+        """Execute a tool with the given input.
+
+        Args:
+            name: Name of the tool to execute
+            tool_input: Input parameters for the tool
+
+        Returns:
+            Result of the tool execution
+        """
+        if self.tools is None:
+            raise RuntimeError("Tools not initialized. Call initialize() first.")
+        return await self.tools.run(name=name, tool_input=tool_input)

agent/providers/anthropic/tools/run.py

@@ -0,0 +1,42 @@
+"""Utility to run shell commands asynchronously with a timeout."""
+
+import asyncio
+
+TRUNCATED_MESSAGE: str = "<response clipped><NOTE>To save on context only part of this file has been shown to you. You should retry this tool after you have searched inside the file with `grep -n` in order to find the line numbers of what you are looking for.</NOTE>"
+MAX_RESPONSE_LEN: int = 16000
+
+
+def maybe_truncate(content: str, truncate_after: int | None = MAX_RESPONSE_LEN):
+    """Truncate content and append a notice if content exceeds the specified length."""
+    return (
+        content
+        if not truncate_after or len(content) <= truncate_after
+        else content[:truncate_after] + TRUNCATED_MESSAGE
+    )
+
+
+async def run(
+    cmd: str,
+    timeout: float | None = 120.0,  # seconds
+    truncate_after: int | None = MAX_RESPONSE_LEN,
+):
+    """Run a shell command asynchronously with a timeout."""
+    process = await asyncio.create_subprocess_shell(
+        cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+    )
+
+    try:
+        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
+        return (
+            process.returncode or 0,
+            maybe_truncate(stdout.decode(), truncate_after=truncate_after),
+            maybe_truncate(stderr.decode(), truncate_after=truncate_after),
+        )
+    except asyncio.TimeoutError as exc:
+        try:
+            process.kill()
+        except ProcessLookupError:
+            pass
+        raise TimeoutError(
+            f"Command '{cmd}' timed out after {timeout} seconds"
+        ) from exc

agent/providers/anthropic/types.py

@@ -0,0 +1,16 @@
+from enum import StrEnum
+
+
+class APIProvider(StrEnum):
+    """Enum for supported API providers."""
+
+    ANTHROPIC = "anthropic"
+    BEDROCK = "bedrock"
+    VERTEX = "vertex"
+
+
+PROVIDER_TO_DEFAULT_MODEL_NAME: dict[APIProvider, str] = {
+    APIProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
+    APIProvider.BEDROCK: "anthropic.claude-3-7-sonnet-20250219-v2:0",
+    APIProvider.VERTEX: "claude-3-5-sonnet-v2@20241022",
+}

agent/providers/omni/__init__.py

@@ -0,0 +1,27 @@
+"""Omni provider implementation."""
+
+# The OmniComputerAgent has been replaced by the unified ComputerAgent
+# which can be found in agent.core.agent
+from .types import APIProvider
+from .experiment import ExperimentManager
+from .visualization import visualize_click, visualize_scroll, calculate_element_center
+from .image_utils import (
+    decode_base64_image,
+    encode_image_base64,
+    clean_base64_data,
+    extract_base64_from_text,
+    get_image_dimensions,
+)
+
+__all__ = [
+    "APIProvider",
+    "ExperimentManager",
+    "visualize_click",
+    "visualize_scroll",
+    "calculate_element_center",
+    "decode_base64_image",
+    "encode_image_base64",
+    "clean_base64_data",
+    "extract_base64_from_text",
+    "get_image_dimensions",
+]

agent/providers/omni/callbacks.py

@@ -0,0 +1,78 @@
+"""Omni callback manager implementation."""
+
+import logging
+from typing import Any, Dict, Optional, Set
+
+from ...core.callbacks import BaseCallbackManager, ContentCallback, ToolCallback, APICallback
+from ...types.tools import ToolResult
+
+logger = logging.getLogger(__name__)
+
+class OmniCallbackManager(BaseCallbackManager):
+    """Callback manager for multi-provider support."""
+    
+    def __init__(
+        self,
+        content_callback: ContentCallback,
+        tool_callback: ToolCallback,
+        api_callback: APICallback,
+    ):
+        """Initialize Omni callback manager.
+        
+        Args:
+            content_callback: Callback for content updates
+            tool_callback: Callback for tool execution results
+            api_callback: Callback for API interactions
+        """
+        super().__init__(
+            content_callback=content_callback,
+            tool_callback=tool_callback,
+            api_callback=api_callback
+        )
+        self._active_tools: Set[str] = set()
+    
+    def on_content(self, content: Any) -> None:
+        """Handle content updates.
+        
+        Args:
+            content: Content update data
+        """
+        logger.debug(f"Content update: {content}")
+        self.content_callback(content)
+    
+    def on_tool_result(self, result: ToolResult, tool_id: str) -> None:
+        """Handle tool execution results.
+        
+        Args:
+            result: Tool execution result
+            tool_id: ID of the tool
+        """
+        logger.debug(f"Tool result for {tool_id}: {result}")
+        self.tool_callback(result, tool_id)
+    
+    def on_api_interaction(
+        self,
+        request: Any,
+        response: Any,
+        error: Optional[Exception] = None
+    ) -> None:
+        """Handle API interactions.
+        
+        Args:
+            request: API request data
+            response: API response data
+            error: Optional error that occurred
+        """
+        if error:
+            logger.error(f"API error: {str(error)}")
+        else:
+            logger.debug(f"API interaction - Request: {request}, Response: {response}")
+        self.api_callback(request, response, error)
+    
+    def get_active_tools(self) -> Set[str]:
+        """Get currently active tools.
+        
+        Returns:
+            Set of active tool names
+        """
+        return self._active_tools.copy()

agent/providers/omni/clients/anthropic.py

@@ -0,0 +1,99 @@
+"""Anthropic API client implementation."""
+
+import logging
+from typing import Any, Dict, List, Optional, Tuple, cast
+import asyncio
+from httpx import ConnectError, ReadTimeout
+
+from anthropic import AsyncAnthropic, Anthropic
+from anthropic.types import MessageParam
+from .base import BaseOmniClient
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicClient(BaseOmniClient):
+    """Client for making calls to Anthropic API."""
+
+    def __init__(self, api_key: str, model: str, max_retries: int = 3, retry_delay: float = 1.0):
+        """Initialize the Anthropic client.
+
+        Args:
+            api_key: Anthropic API key
+            model: Anthropic model name (e.g. "claude-3-opus-20240229")
+            max_retries: Maximum number of retries for API calls
+            retry_delay: Base delay between retries in seconds
+        """
+        if not model:
+            raise ValueError("Model name must be provided")
+
+        self.client = AsyncAnthropic(api_key=api_key)
+        self.model: str = model  # Add explicit type annotation
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+
+    def _convert_message_format(self, messages: List[Dict[str, Any]]) -> List[MessageParam]:
+        """Convert messages from standard format to Anthropic format.
+
+        Args:
+            messages: Messages in standard format
+
+        Returns:
+            Messages in Anthropic format
+        """
+        anthropic_messages = []
+
+        for message in messages:
+            if message["role"] == "user":
+                anthropic_messages.append({"role": "user", "content": message["content"]})
+            elif message["role"] == "assistant":
+                anthropic_messages.append({"role": "assistant", "content": message["content"]})
+
+        # Cast the list to the correct type expected by Anthropic
+        return cast(List[MessageParam], anthropic_messages)
+
+    async def run_interleaved(
+        self, messages: List[Dict[str, Any]], system: str, max_tokens: int
+    ) -> Any:
+        """Run model with interleaved conversation format.
+
+        Args:
+            messages: List of messages to process
+            system: System prompt
+            max_tokens: Maximum tokens to generate
+
+        Returns:
+            Model response
+        """
+        last_error = None
+
+        for attempt in range(self.max_retries):
+            try:
+                # Convert messages to Anthropic format
+                anthropic_messages = self._convert_message_format(messages)
+
+                response = await self.client.messages.create(
+                    model=self.model,
+                    max_tokens=max_tokens,
+                    temperature=0,
+                    system=system,
+                    messages=anthropic_messages,
+                )
+
+                return response
+
+            except (ConnectError, ReadTimeout) as e:
+                last_error = e
+                logger.warning(
+                    f"Connection error on attempt {attempt + 1}/{self.max_retries}: {str(e)}"
+                )
+                if attempt < self.max_retries - 1:
+                    await asyncio.sleep(self.retry_delay * (attempt + 1))  # Exponential backoff
+                continue
+
+            except Exception as e:
+                logger.error(f"Unexpected error in Anthropic API call: {str(e)}")
+                raise RuntimeError(f"Anthropic API call failed: {str(e)}")
+
+        # If we get here, all retries failed
+        raise RuntimeError(f"Connection error after {self.max_retries} retries: {str(last_error)}")

agent/providers/omni/clients/base.py

@@ -0,0 +1,44 @@
+"""Base client implementation for Omni providers."""
+
+import os
+import logging
+from typing import Dict, List, Optional, Any, Tuple
+import aiohttp
+import json
+
+logger = logging.getLogger(__name__)
+
+class BaseOmniClient:
+    """Base class for provider-specific clients."""
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None
+    ):
+        """Initialize base client.
+        
+        Args:
+            api_key: Optional API key
+            model: Optional model name
+        """
+        self.api_key = api_key
+        self.model = model
+        
+    async def run_interleaved(
+        self,
+        messages: List[Dict[str, Any]],
+        system: str,
+        max_tokens: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Run interleaved chat completion.
+        
+        Args:
+            messages: List of message dicts
+            system: System prompt
+            max_tokens: Optional max tokens override
+            
+        Returns:
+            Response dict
+        """
+        raise NotImplementedError