inspect-ai 0.3.71__py3-none-any.whl → 0.3.73__py3-none-any.whl

inspect_ai/tool/_tools/_computer/_resources/tool/_x11_client.py

@@ -1,4 +1,4 @@
-"""Based on https://github.com/anthropics/anthropic-quickstarts/blob/main/computer-use-demo/computer_use_demo/tools/computer.py"""
+"""Inspired by https://github.com/anthropics/anthropic-quickstarts/blob/main/computer-use-demo/computer_use_demo/tools/computer.py"""
 
 import asyncio
 import base64
@@ -19,21 +19,8 @@ TYPING_GROUP_SIZE = 50
 
 ColorCount = Literal[4096, 2048, 1024, 512, 256, 128, 64, 32, 16, 8, 4]
 
-Action = Literal[
-    "key",
-    "type",
-    "mouse_move",
-    "left_click",
-    "left_click_drag",
-    "right_click",
-    "middle_click",
-    "double_click",
-    "screenshot",
-    "cursor_position",
-]
-
-
-class ToolError(Exception):
+
+class X11ClientError(Exception):
     def __init__(self, message):
         self.message = message
 
@@ -83,7 +70,7 @@ class X11Client:
 
     @property
     def options(self) -> ComputerToolOptions:
-        width, height = self.scale_coordinates("computer", self.width, self.height)
+        width, height = self._scale_coordinates("computer", self.width, self.height)
         return {
             "display_width_px": width,
             "display_height_px": height,
@@ -105,120 +92,191 @@ class X11Client:
 
         self.xdotool = f"{self._display_prefix}xdotool"
 
-    async def __call__(
+    async def key(self, text: str) -> ToolResult:
+        return await self._shell(f"{self.xdotool} key -- {_key_arg_for_text(text)}")
+
+    async def hold_key(self, text: str, duration: int) -> ToolResult:
+        key_arg = _key_arg_for_text(text)
+        await self._shell(f"{self.xdotool} keydown -- {key_arg}", False)
+        await asyncio.sleep(duration)
+        return await self._shell(f"{self.xdotool} keyup -- {key_arg}")
+
+    async def type(self, text: str) -> ToolResult:
+        results: list[ToolResult] = []
+        for chunk in chunks(text, TYPING_GROUP_SIZE):
+            cmd = (
+                f"{self.xdotool} type --delay {TYPING_DELAY_MS} -- {shlex.quote(chunk)}"
+            )
+            results.append(await self._shell(cmd, take_screenshot=False))
+
+        screenshot_base64 = await self._take_screenshot_after_delay()
+        return ToolResult(
+            output="".join(result.output or "" for result in results),
+            error="".join(result.error or "" for result in results),
+            base64_image=screenshot_base64,
+        )
+
+    async def cursor_position(self) -> ToolResult:
+        result = await self._shell(
+            f"{self.xdotool} getmouselocation --shell",
+            take_screenshot=False,
+        )
+        output = result.output or ""
+        x, y = self._scale_coordinates(
+            "computer",
+            int(output.split("X=")[1].split("\n")[0]),
+            int(output.split("Y=")[1].split("\n")[0]),
+        )
+        return result.replace(output=f"X={x},Y={y}")
+
+    async def left_mouse_down(self) -> ToolResult:
+        return await self._shell(f"{self.xdotool} mousedown 1")
+
+    async def left_mouse_up(self) -> ToolResult:
+        return await self._shell(f"{self.xdotool} mouseup 1")
+
+    async def mouse_move(self, coordinate: tuple[int, int]) -> ToolResult:
+        return await self._mouse_move_and("mouse_move", coordinate, None)
+
+    async def left_click(
+        self, coordinate: tuple[int, int] | None, text: str | None
+    ) -> ToolResult:
+        return await self._mouse_move_and("left_click", coordinate, text)
+
+    async def right_click(
+        self, coordinate: tuple[int, int] | None, text: str | None
+    ) -> ToolResult:
+        return await self._mouse_move_and("right_click", coordinate, text)
+
+    async def middle_click(
+        self, coordinate: tuple[int, int] | None, text: str | None
+    ) -> ToolResult:
+        return await self._mouse_move_and("middle_click", coordinate, text)
+
+    async def double_click(
+        self, coordinate: tuple[int, int] | None, text: str | None
+    ) -> ToolResult:
+        return await self._mouse_move_and("double_click", coordinate, text)
+
+    async def triple_click(
+        self, coordinate: tuple[int, int] | None, text: str | None
+    ) -> ToolResult:
+        return await self._mouse_move_and("triple_click", coordinate, text)
+
+    async def left_click_drag(
+        self, start_coordinate: tuple[int, int], coordinate: tuple[int, int]
+    ) -> ToolResult:
+        await self._move_mouse_to_coordinate(start_coordinate, False)
+        x, y = self._scale_coordinates("api", *coordinate)
+        return await self._shell(
+            f"{self.xdotool} mousedown 1 mousemove --sync {x} {y} mouseup 1"
+        )
+
+    async def scroll(
         self,
-        *,
-        action: Action,
-        text: str | None = None,
-        coordinate: tuple[int, int] | None = None,
-        **kwargs,
-    ):
-        if action in ("mouse_move", "left_click_drag"):
-            if coordinate is None:
-                raise ToolError(f"coordinate is required for {action}")
-            if text is not None:
-                raise ToolError(f"text is not accepted for {action}")
-            if not isinstance(coordinate, list) or len(coordinate) != 2:
-                raise ToolError(f"{coordinate} must be a tuple of length 2")
-            if not all(isinstance(i, int) and i >= 0 for i in coordinate):
-                raise ToolError(f"{coordinate} must be a tuple of non-negative ints")
-
-            x, y = self.scale_coordinates("api", coordinate[0], coordinate[1])
+        scroll_direction: Literal["up", "down", "left", "right"],
+        scroll_amount: int,
+        coordinate: tuple[int, int] | None,
+        text: str | None,
+    ) -> ToolResult:
+        if coordinate:
+            await self._move_mouse_to_coordinate(coordinate, False)
+        scroll_button = {
+            "up": 4,
+            "down": 5,
+            "left": 6,
+            "right": 7,
+        }[scroll_direction]
+
+        if text:
+            key_arg = _key_arg_for_text(text)
+            await self._shell(f"{self.xdotool} keydown -- {key_arg}", False)
+            await self._shell(
+                f"{self.xdotool} click --repeat {scroll_amount} {scroll_button}",
+                False,
+            )
+            return await self._shell(f"{self.xdotool} keyup -- {key_arg}")
+        else:
+            return await self._shell(
+                f"{self.xdotool} click --repeat {scroll_amount} {scroll_button}"
+            )
 
-            if action == "mouse_move":
-                return await self.shell(f"{self.xdotool} mousemove --sync {x} {y}")
-            elif action == "left_click_drag":
-                return await self.shell(
-                    f"{self.xdotool} mousedown 1 mousemove --sync {x} {y} mouseup 1"
-                )
-
-        if action in ("key", "type"):
-            if text is None:
-                raise ToolError(f"text is required for {action}")
-            if coordinate is not None:
-                raise ToolError(f"coordinate is not accepted for {action}")
-            if not isinstance(text, str):
-                raise ToolError(f"{text} must be a string")
-
-            if action == "key":
-                return await self.shell(
-                    f"{self.xdotool} key -- {' '.join(shlex.quote(part) for part in text.split())}"
-                )
-            elif action == "type":
-                results: list[ToolResult] = []
-                for chunk in chunks(text, TYPING_GROUP_SIZE):
-                    cmd = f"{self.xdotool} type --delay {TYPING_DELAY_MS} -- {shlex.quote(chunk)}"
-                    results.append(await self.shell(cmd, take_screenshot=False))
-
-                screenshot_base64 = await self.take_screenshot_after_delay()
-                return ToolResult(
-                    output="".join(result.output or "" for result in results),
-                    error="".join(result.error or "" for result in results),
-                    base64_image=screenshot_base64,
-                )
-
-        if action in (
+    async def wait(self, duration: int) -> ToolResult:
+        await asyncio.sleep(duration)
+        return await self.screenshot()
+
+    async def screenshot(self) -> ToolResult:
+        return await self._screenshot()
+
+    async def _mouse_move_and(
+        self,
+        action: Literal[
+            "mouse_move",
             "left_click",
             "right_click",
-            "double_click",
             "middle_click",
-            "screenshot",
-            "cursor_position",
-        ):
-            if text is not None:
-                raise ToolError(f"text is not accepted for {action}")
-            if coordinate is not None:
-                raise ToolError(f"coordinate is not accepted for {action}")
-
-            if action == "screenshot":
-                return await self.screenshot()
-            elif action == "cursor_position":
-                result = await self.shell(
-                    f"{self.xdotool} getmouselocation --shell",
-                    take_screenshot=False,
-                )
-                output = result.output or ""
-                x, y = self.scale_coordinates(
-                    "computer",
-                    int(output.split("X=")[1].split("\n")[0]),
-                    int(output.split("Y=")[1].split("\n")[0]),
-                )
-                return result.replace(output=f"X={x},Y={y}")
-            else:
-                click_arg = {
-                    "left_click": "1",
-                    "right_click": "3",
-                    "middle_click": "2",
-                    "double_click": "--repeat 2 --delay 300 1",
-                }[action]
-                return await self.shell(f"{self.xdotool} click {click_arg}")
-
-        raise ToolError(f"Invalid action: {action}")
-
-    async def screenshot(self):
+            "double_click",
+            "triple_click",
+        ],
+        coordinate: tuple[int, int] | None,
+        text: str | None,
+    ):
+        should_move = action == "mouse_move" or coordinate
+        if should_move:
+            assert coordinate  # coding/type safety error
+            move_result = await self._move_mouse_to_coordinate(
+                coordinate, action == "mouse_move"
+            )
+            if action == "mouse_move":
+                return move_result
+        click_arg = {
+            "left_click": "1",
+            "right_click": "3",
+            "middle_click": "2",
+            "double_click": "--repeat 2 --delay 300 1",
+            "triple_click": "--repeat 3 --delay 300 1",
+        }[action]
+
+        if text:
+            key_arg = _key_arg_for_text(text)
+            await self._shell(f"{self.xdotool} keydown -- {key_arg}", False)
+            await self._shell(f"{self.xdotool} click {click_arg}", False)
+            return await self._shell(f"{self.xdotool} keyup -- {key_arg}")
+        else:
+            return await self._shell(f"{self.xdotool} click {click_arg}")
+
+    async def _move_mouse_to_coordinate(
+        self, coordinate: tuple[int, int], take_screenshot: bool
+    ):
+        x, y = self._scale_coordinates("api", *coordinate)
+        return await self._shell(
+            f"{self.xdotool} mousemove --sync {x} {y}", take_screenshot=take_screenshot
+        )
+
+    async def _screenshot(self):
         """Take a screenshot of the current screen and return the base64 encoded image."""
         output_dir = Path(OUTPUT_DIR)
         output_dir.mkdir(parents=True, exist_ok=True)
         path = output_dir / f"screenshot_{uuid4().hex}.png"
 
-        result = await self.shell(
+        result = await self._shell(
             f"{self._display_prefix}scrot --silent -p {path}", take_screenshot=False
         )
         if self._scaling_enabled:
-            x, y = self.scale_coordinates("computer", self.width, self.height)
+            x, y = self._scale_coordinates("computer", self.width, self.height)
             convert_cmd = f"convert {path} -resize {x}x{y}!"
             if self.color_count is not None:
                 convert_cmd += f" -colors {self.color_count}"
             convert_cmd += f" {path}"
-            await self.shell(convert_cmd, take_screenshot=False)
+            await self._shell(convert_cmd, take_screenshot=False)
 
         if path.exists():
             return result.replace(
                 base64_image=base64.b64encode(path.read_bytes()).decode()
             )
-        raise ToolError(f"Failed to take screenshot: {result.error}")
+        raise X11ClientError(f"Failed to take screenshot: {result.error}")
 
-    async def shell(self, command: str, take_screenshot=True) -> ToolResult:
+    async def _shell(self, command: str, take_screenshot=True) -> ToolResult:
         """Run a shell command and return the output, error, and optionally a screenshot."""
         logging.debug(f"running shell command {command}")
         _, stdout, stderr = await run(command)
@@ -226,17 +284,17 @@ class X11Client:
         return ToolResult(
             output=stdout,
             error=stderr,
-            base64_image=(await self.take_screenshot_after_delay())
+            base64_image=(await self._take_screenshot_after_delay())
             if take_screenshot
             else None,
         )
 
-    async def take_screenshot_after_delay(self) -> str:
+    async def _take_screenshot_after_delay(self) -> str:
         # delay to let things settle before taking a screenshot
         await asyncio.sleep(self._screenshot_delay)
-        return (await self.screenshot()).base64_image
+        return (await self._screenshot()).base64_image
 
-    def scale_coordinates(self, source: ScalingSource, x: int, y: int):
+    def _scale_coordinates(self, source: ScalingSource, x: int, y: int):
         """Scale coordinates to a target maximum resolution."""
         if not self._scaling_enabled:
             return x, y
@@ -255,8 +313,12 @@ class X11Client:
         y_scaling_factor = target_dimension["height"] / self.height
         if source == "api":
             if x > self.width or y > self.height:
-                raise ToolError(f"Coordinates {x}, {y} are out of bounds")
+                raise X11ClientError(f"Coordinates {x}, {y} are out of bounds")
             # scale up
             return round(x / x_scaling_factor), round(y / y_scaling_factor)
         # scale down
         return round(x * x_scaling_factor), round(y * y_scaling_factor)
+
+
+def _key_arg_for_text(text: str) -> str:
+    return " ".join(shlex.quote(part) for part in text.split())

inspect_ai/tool/_tools/_computer/_resources/tool/computer_tool.py

@@ -1,15 +1,24 @@
-import argparse
 import asyncio
 import json
 import logging
 import os
 import sys
 import time
+from argparse import Namespace
+from typing import TypeVar
 
+from _args import parse_arguments
+from _constants import Action
 from _logger import setup_logger
 from _tool_result import ToolResult
 from _x11_client import X11Client
 
+
+class ComputerToolError(Exception):
+    def __init__(self, message):
+        self.message = message
+
+
 # This is a bit sketchy. We really want to use relative imports here. Using absolute imports
 # works at runtime, but it prevents intellisense from working. However, when this folder is
 # copied to the container, by default relative imports won't work if this file is launched
@@ -44,29 +53,67 @@ def main():
         sys.exit(1)
 
 
-def parse_arguments():
-    parser = argparse.ArgumentParser(description="Execute computer tool action")
-    parser.add_argument("--action", type=str, required=True, help="Action to perform")
-    parser.add_argument("--text", type=str, help="Optional text parameter")
-    parser.add_argument(
-        "--coordinate",
-        type=int,
-        nargs=2,
-        help="Optional coordinate parameter as a list of two integers",
-    )
-    return parser.parse_args()
-
-
-async def execute_action(args) -> ToolResult:
+async def execute_action(args: Namespace) -> ToolResult:
     # we can't do anything until X11 is ready to go.
     await wait_for_file("/tmp/xfce_started")
 
     computer = X11Client()
-    return await computer(
-        action=args.action,
-        text=args.text,
-        coordinate=args.coordinate if args.coordinate else None,
-    )
+    action: Action = args.action
+    match action:
+        case "key":
+            return await computer.key(not_none(args.text, "text"))
+        case "hold_key":
+            return await computer.hold_key(
+                not_none(args.text, "text"), not_none(args.duration, "duration")
+            )
+        case "type":
+            return await computer.type(not_none(args.text, "text"))
+        case "cursor_position":
+            return await computer.cursor_position()
+        case "left_mouse_down":
+            return await computer.left_mouse_down()
+        case "left_mouse_up":
+            return await computer.left_mouse_up()
+        case "mouse_move":
+            return await computer.mouse_move(not_none(args.coordinate, "coordinate"))
+        case "left_click":
+            return await computer.left_click(
+                getattr(args, "coordinate", None), getattr(args, "text", None)
+            )
+        case "right_click":
+            return await computer.right_click(
+                getattr(args, "coordinate", None), getattr(args, "text", None)
+            )
+        case "middle_click":
+            return await computer.middle_click(
+                getattr(args, "coordinate", None), getattr(args, "text", None)
+            )
+        case "double_click":
+            return await computer.double_click(
+                getattr(args, "coordinate", None), getattr(args, "text", None)
+            )
+        case "triple_click":
+            return await computer.triple_click(
+                getattr(args, "coordinate", None), getattr(args, "text", None)
+            )
+        case "left_click_drag":
+            return await computer.left_click_drag(
+                not_none(args.start_coordinate, "start_coordinate"),
+                not_none(args.coordinate, "coordinate"),
+            )
+        case "scroll":
+            return await computer.scroll(
+                not_none(args.scroll_direction, "scroll_direction"),
+                not_none(args.scroll_amount, "scroll_amount"),
+                getattr(args, "coordinate", None),
+                getattr(args, "text", None),
+            )
+        case "wait":
+            return await computer.wait(not_none(args.duration, "duration"))
+        case "screenshot":
+            return await computer.screenshot()
+
+    raise ComputerToolError(f"Invalid action: {action}")
 
 
 async def wait_for_file(file_path, check_interval=1):
@@ -81,5 +128,14 @@ async def wait_for_file(file_path, check_interval=1):
     )
 
 
+T = TypeVar("T")
+
+
+def not_none(value: T | None, name: str) -> T:
+    if value is None:
+        raise ComputerToolError(f"{name} must be provided")
+    return value
+
+
 if __name__ == "__main__":
     main()

inspect_ai/tool/_tools/_computer/_resources/tool/pyproject.toml

@@ -0,0 +1,65 @@
+[build-system]
+requires = ["setuptools>=64", "setuptools_scm[toml]>=8"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools_scm]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["inspect_ai*"]
+
+[tool.ruff]
+src = ["."]
+
+[tool.ruff.lint]
+select = [
+    "E",      # pycodestyle errors
+    "W",      # pycodestyle warnings
+    "F",      # flake8
+    "D",      # pydocstyle
+    "I",      # isort
+    "SIM101", # duplicate isinstance
+    "UP038",  # non-pep604-isinstance
+    # "RET", # flake8-return
+    # "RUF", # ruff rules
+]
+ignore = ["E203", "E501", "D10", "D212", "D415"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.pytest.ini_options]
+minversion = "7.0"
+addopts = "-rA --doctest-modules --color=yes"
+doctest_optionflags = ["NORMALIZE_WHITESPACE", "IGNORE_EXCEPTION_DETAIL"]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+log_level = "warning"
+
+[tool.mypy]
+warn_unused_ignores = true
+no_implicit_reexport = true
+strict_equality = true
+warn_redundant_casts = true
+warn_unused_configs = true
+disallow_any_explicit = true
+disallow_any_generics = true
+disallow_subclassing_any = true
+plugins=["pydantic.mypy"]
+
+
+[tool.pydantic-mypy]
+init_forbid_extra = true
+init_typed = true
+
+[tool.check-wheel-contents]
+ignore = ["W002", "W009"]
+
+[project]
+name = "web_browser_tool_container"
+requires-python = ">=3.10"
+dynamic = ["version", "dependencies"]
+
+
+[project.optional-dependencies]
+dev = ["pytest"]