baqr/computer_use_ootb

1import asyncio2import os3from typing import ClassVar, Literal4 5from anthropic.types.beta import BetaToolBash20241022Param6 7from .base import BaseAnthropicTool, CLIResult, ToolError, ToolResult8 9 10class _BashSession:11    """A session of a bash shell."""12 13    _started: bool14    _process: asyncio.subprocess.Process15 16    command: str = "/bin/bash"17    _output_delay: float = 0.2  # seconds18    _timeout: float = 120.0  # seconds19    _sentinel: str = "<<exit>>"20 21    def __init__(self):22        self._started = False23        self._timed_out = False24 25    async def start(self):26        if self._started:27            return28 29        self._process = await asyncio.create_subprocess_shell(30            self.command,31            shell=False,32            stdin=asyncio.subprocess.PIPE,33            stdout=asyncio.subprocess.PIPE,34            stderr=asyncio.subprocess.PIPE,35        )36 37        self._started = True38 39    def stop(self):40        """Terminate the bash shell."""41        if not self._started:42            raise ToolError("Session has not started.")43        if self._process.returncode is not None:44            return45        self._process.terminate()46 47    async def run(self, command: str):48        """Execute a command in the bash shell."""49        if not self._started:50            raise ToolError("Session has not started.")51        if self._process.returncode is not None:52            return ToolResult(53                system="tool must be restarted",54                error=f"bash has exited with returncode {self._process.returncode}",55            )56        if self._timed_out:57            raise ToolError(58                f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",59            )60 61        # we know these are not None because we created the process with PIPEs62        assert self._process.stdin63        assert self._process.stdout64        assert self._process.stderr65 66        # send command to the process67        self._process.stdin.write(68            command.encode() + f"; echo '{self._sentinel}'\n".encode()69        )70        await self._process.stdin.drain()71 72        # read output from the process, until the sentinel is found73        output = ""74        try:75            async with asyncio.timeout(self._timeout):76                while True:77                    await asyncio.sleep(self._output_delay)78                    data = await self._process.stdout.readline()79                    if not data:80                        break81                    line = data.decode()82                    output += line83                    if self._sentinel in line:84                        output = output.replace(self._sentinel, "")85                        break86        except asyncio.TimeoutError:87            self._timed_out = True88            raise ToolError(89                f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",90            ) from None91 92        error = await self._process.stderr.read()93        error = error.decode()94 95        return CLIResult(output=output.strip(), error=error.strip())96 97 98class BashTool(BaseAnthropicTool):99    """100    A tool that allows the agent to run bash commands.101    The tool parameters are defined by Anthropic and are not editable.102    """103 104    _session: _BashSession | None105    name: ClassVar[Literal["bash"]] = "bash"106    api_type: ClassVar[Literal["bash_20241022"]] = "bash_20241022"107 108    def __init__(self):109        self._session = None110        super().__init__()111 112    async def __call__(113        self, command: str | None = None, restart: bool = False, **kwargs114    ):115        if restart:116            if self._session:117                self._session.stop()118            self._session = _BashSession()119            await self._session.start()120 121            return ToolResult(system="tool has been restarted.")122 123        if self._session is None:124            self._session = _BashSession()125            await self._session.start()126 127        if command is not None:128            return await self._session.run(command)129 130        raise ToolError("no command provided.")131 132    def to_params(self) -> BetaToolBash20241022Param:133        return {134            "type": self.api_type,135            "name": self.name,136        }