inspect-ai 0.3.58__py3-none-any.whl → 0.3.60__py3-none-any.whl

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

@@ -0,0 +1,262 @@
+"""Based on https://github.com/anthropics/anthropic-quickstarts/blob/main/computer-use-demo/computer_use_demo/tools/computer.py"""
+
+import asyncio
+import base64
+import logging
+import os
+import shlex
+from pathlib import Path
+from typing import Literal, TypedDict
+from uuid import uuid4
+
+from _run import run
+from _tool_result import ToolResult
+
+OUTPUT_DIR = "/tmp/outputs"
+
+TYPING_DELAY_MS = 12
+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):
+    def __init__(self, message):
+        self.message = message
+
+
+class Resolution(TypedDict):
+    width: int
+    height: int
+
+
+# sizes above XGA/WXGA are not recommended (see README.md)
+# scale down to one of these targets if ComputerTool._scaling_enabled is set
+MAX_SCALING_TARGETS: dict[str, Resolution] = {
+    "XGA": Resolution(width=1024, height=768),  # 4:3
+    "WXGA": Resolution(width=1280, height=800),  # 16:10
+    "FWXGA": Resolution(width=1366, height=768),  # ~16:9
+}
+
+
+ScalingSource = Literal["computer", "api"]
+
+
+class ComputerToolOptions(TypedDict):
+    display_height_px: int
+    display_width_px: int
+    display_number: int | None
+
+
+def chunks(s: str, chunk_size: int) -> list[str]:
+    return [s[i : i + chunk_size] for i in range(0, len(s), chunk_size)]
+
+
+class X11Client:
+    """
+    A tool that allows the agent to interact with the screen, keyboard, and mouse of the current computer.
+
+    The tool parameters are defined by Anthropic and are not editable.
+    """
+
+    width: int
+    height: int
+    display_num: int | None
+    # TODO: Complete plumbing this or remove it
+    color_count: ColorCount | None = 256
+
+    _screenshot_delay = 2.0
+    _scaling_enabled = True
+
+    @property
+    def options(self) -> ComputerToolOptions:
+        width, height = self.scale_coordinates("computer", self.width, self.height)
+        return {
+            "display_width_px": width,
+            "display_height_px": height,
+            "display_number": self.display_num,
+        }
+
+    def __init__(self):
+        super().__init__()
+
+        self.width = int(os.getenv("WIDTH") or 0)
+        self.height = int(os.getenv("HEIGHT") or 0)
+        assert self.width and self.height, "WIDTH, HEIGHT must be set"
+        if (display_num := os.getenv("DISPLAY_NUM")) is not None:
+            self.display_num = int(display_num)
+            self._display_prefix = f"DISPLAY=:{self.display_num} "
+        else:
+            self.display_num = None
+            self._display_prefix = ""
+
+        self.xdotool = f"{self._display_prefix}xdotool"
+
+    async def __call__(
+        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])
+
+            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(output=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 (
+            "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):
+        """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(
+            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)
+            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)
+
+        if path.exists():
+            return result.replace(
+                base64_image=base64.b64encode(path.read_bytes()).decode()
+            )
+        raise ToolError(f"Failed to take screenshot: {result.error}")
+
+    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)
+        logging.debug(f"shell command returned stdout: {stdout}, stderr: {stderr}")
+        return ToolResult(
+            output=stdout,
+            error=stderr,
+            base64_image=(await self.take_screenshot_after_delay())
+            if take_screenshot
+            else None,
+        )
+
+    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
+
+    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
+        ratio = self.width / self.height
+        target_dimension = None
+        for dimension in MAX_SCALING_TARGETS.values():
+            # allow some error in the aspect ratio - not ratios are exactly 16:9
+            if abs(dimension["width"] / dimension["height"] - ratio) < 0.02:
+                if dimension["width"] < self.width:
+                    target_dimension = dimension
+                break
+        if target_dimension is None:
+            return x, y
+        # should be less than 1
+        x_scaling_factor = target_dimension["width"] / self.width
+        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")
+            # 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)

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

@@ -0,0 +1,85 @@
+import argparse
+import asyncio
+import json
+import logging
+import os
+import sys
+import time
+
+from _logger import setup_logger
+from _tool_result import ToolResult
+from _x11_client import X11Client
+
+# 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
+# normally. To overcome this, two things need to happen:
+# 1. PYTHONPATH must be set to the parent of the container folder. `PYTHONPATH=/opt`
+# 2. The program must be launched with the -m flag. `python3 -m computer_tool.computer_tool`
+#
+# TODO: There's got to be a cleaner way.
+
+my_logger = setup_logger(logging.INFO)
+
+
+def main():
+    try:
+        args = parse_arguments()
+        my_logger.info(f"({args})")
+        result = asyncio.run(execute_action(args))
+
+        print(
+            json.dumps(
+                {
+                    "output": result.output,
+                    "error": result.error,
+                    "base64_image": result.base64_image,
+                }
+            )
+        )
+        my_logger.debug("SUCCESS")
+    except Exception as e:
+        my_logger.warning(f"An error occurred: {e}")
+        print(f"An error occurred: {e}", file=sys.stderr)
+        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:
+    # 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,
+    )
+
+
+async def wait_for_file(file_path, check_interval=1):
+    if os.path.exists(file_path):
+        return
+    my_logger.info(f"Waiting for {file_path}")
+    start_time = time.time()
+    while not os.path.exists(file_path):
+        await asyncio.sleep(check_interval)
+    my_logger.info(
+        f"Done waiting for {file_path} after {time.time() - start_time:.1f} seconds"
+    )
+
+
+if __name__ == "__main__":
+    main()

inspect_ai/util/__init__.py

@@ -3,6 +3,7 @@ from inspect_ai._util.trace import trace_action, trace_message
 from ._concurrency import concurrency
 from ._console import input_screen
 from ._display import DisplayType, display_type
+from ._limit import SampleLimitExceededError
 from ._panel import InputPanel, input_panel
 from ._resource import resource
 from ._sandbox import (
@@ -36,6 +37,7 @@ __all__ = [
     "input_panel",
     "input_screen",
     "OutputLimitExceededError",
+    "SampleLimitExceededError",
     "resource",
     "subprocess",
     "SandboxEnvironment",

inspect_ai/util/_display.py

@@ -49,3 +49,8 @@ def display_type() -> DisplayType:
         return _display_type
     else:
         return init_display_type()
+
+
+def display_type_initialized() -> bool:
+    global _display_type
+    return _display_type is not None

inspect_ai/util/_limit.py

@@ -0,0 +1,26 @@
+from typing import Literal
+
+
+class SampleLimitExceededError(Exception):
+    """Exception raised when a sample limit is exceeded.
+
+    Args:
+       type (Literal["message", "time", "token", "operator"]): Type of limit exceeded.
+       value (int): Value compared to
+       limit (int): Limit applied.
+       message (str | None): Optional. Human readable message.
+    """
+
+    def __init__(
+        self,
+        type: Literal["message", "time", "token", "operator", "custom"],
+        *,
+        value: int,
+        limit: int,
+        message: str | None = None,
+    ) -> None:
+        self.type = type
+        self.value = value
+        self.limit = limit
+        self.message = f"Exceeded {type} limit: {limit:,}"
+        super().__init__(message)

inspect_ai/util/_sandbox/docker/docker.py

@@ -1,4 +1,5 @@
 import errno
+import json
 import os
 import tempfile
 from logging import getLogger
@@ -7,9 +8,11 @@ from typing import Literal, Union, cast, overload
 
 from typing_extensions import override
 
-from inspect_ai.util._subprocess import ExecResult
+from inspect_ai.util._subprocess import ExecResult, subprocess
 
 from ..environment import (
+    HostMapping,
+    PortMapping,
     SandboxConnection,
     SandboxEnvironment,
     SandboxEnvironmentConfigType,
@@ -439,6 +442,7 @@ class DockerSandboxEnvironment(SandboxEnvironment):
                     "remote-containers.attachToRunningContainer",
                     container,
                 ],
+                ports=await get_ports_info(container),
                 container=container,
             )
         # error (not currently running)
@@ -468,3 +472,62 @@ async def container_working_dir(
             + f"{result.stderr}"
         )
         return default
+
+
+async def get_ports_info(container: str) -> list[PortMapping] | None:
+    try:
+        result = await subprocess(
+            [
+                "docker",
+                "inspect",
+                container,
+                "--format",
+                "{{json .NetworkSettings.Ports}}",
+            ],
+            timeout=60,
+        )
+
+        if not result.success:
+            raise RuntimeError(result.stderr)
+
+        return parse_docker_inspect_ports(result.stdout)
+
+    # It's currently a policy decision to let docker timeouts to be silent.
+    except TimeoutError:
+        return None
+
+
+def parse_docker_inspect_ports(json_str: str) -> list[PortMapping] | None:
+    """
+    Parses the JSON output from `docker inspect {container_name} --format='{{json .NetworkSettings.Ports}}'` to extract port mappings.
+
+    Args:
+        json_str (str): A JSON string representing the `NetworkSettings.Ports` output of `docker inspect`. e.g.
+          ```
+          {
+              "5900/tcp": [{"HostIp": "0.0.0.0", "HostPort": "54023"}],
+              "8080/tcp": [{"HostIp": "0.0.0.0", "HostPort": "54024"}]
+          }
+          ```
+
+    Returns:
+        list[PortMapping] | None: A list of PortMapping objects if any port mappings are found,
+                                   otherwise None.
+    """
+    data = json.loads(json_str)
+    port_mappings = []
+    for port_protocol, mappings in data.items():
+        if mappings is None:
+            continue
+        container_port, protocol = port_protocol.split("/")
+        host_mappings = [
+            HostMapping(host_ip=mapping["HostIp"], host_port=int(mapping["HostPort"]))
+            for mapping in mappings
+        ]
+        port_mapping = PortMapping(
+            container_port=int(container_port),
+            protocol=protocol,
+            mappings=host_mappings,
+        )
+        port_mappings.append(port_mapping)
+    return port_mappings if port_mappings else None

inspect_ai/util/_sandbox/docker/internal.py

@@ -6,13 +6,15 @@ from inspect_ai.util._subprocess import subprocess
 INSPECT_WEB_BROWSER_IMAGE_DOCKERHUB = "aisiuk/inspect-web-browser-tool"
 
 INSPECT_WEB_BROWSER_IMAGE = "inspect_web_browser"
+INSPECT_COMPUTER_IMAGE = "inspect-computer-tool"
 
 INTERNAL_IMAGES = {
     INSPECT_WEB_BROWSER_IMAGE: PKG_PATH
     / "tool"
     / "_tools"
     / "_web_browser"
-    / "_resources"
+    / "_resources",
+    INSPECT_COMPUTER_IMAGE: PKG_PATH / "tool" / "beta" / "_computer" / "_resources",
 }
 
 

inspect_ai/util/_sandbox/docker/prereqs.py

@@ -57,7 +57,7 @@ async def validate_docker_compose(
     version: str = DOCKER_COMPOSE_REQUIRED_VERSION,
 ) -> None:
     def parse_version(stdout: str) -> semver.Version:
-        version = json.loads(stdout)["version"].removeprefix("v")
+        version = json.loads(stdout)["version"].removeprefix("v").split("+")[0]
         return semver.Version.parse(version)
 
     await validate_version(

inspect_ai/util/_sandbox/environment.py

@@ -28,6 +28,17 @@ SampleCleanup = Callable[
 ]
 
 
+class HostMapping(BaseModel):
+    host_ip: str
+    host_port: int
+
+
+class PortMapping(BaseModel):
+    container_port: int
+    protocol: Literal["tcp", "udp"]
+    mappings: list[HostMapping]
+
+
 class SandboxConnection(BaseModel):
     """Information required to connect to sandbox."""
 
@@ -40,6 +51,9 @@ class SandboxConnection(BaseModel):
     vscode_command: list[Any] | None = Field(default=None)
     """Optional vscode command (+args) to connect to sandbox."""
 
+    ports: list[PortMapping] | None = Field(default=None)
+    """Optional list of port mappings into container"""
+
     container: str | None = Field(default=None)
     """Optional container name (does not apply to all sandboxes)."""
 

{inspect_ai-0.3.58.dist-info → inspect_ai-0.3.60.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: inspect_ai
-Version: 0.3.58
+Version: 0.3.60
 Summary: Framework for large language model evaluations
 Author: UK AI Safety Institute
 License: MIT License
@@ -54,6 +54,7 @@ Requires-Dist: aioboto3; extra == "dev"
 Requires-Dist: azure-ai-inference; extra == "dev"
 Requires-Dist: google-cloud-aiplatform; extra == "dev"
 Requires-Dist: google-generativeai; extra == "dev"
+Requires-Dist: goodfire; extra == "dev"
 Requires-Dist: groq; extra == "dev"
 Requires-Dist: ipython; extra == "dev"
 Requires-Dist: mistralai; extra == "dev"
@@ -67,7 +68,7 @@ Requires-Dist: pytest-asyncio; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"
 Requires-Dist: pytest-dotenv; extra == "dev"
 Requires-Dist: pytest-xdist; extra == "dev"
-Requires-Dist: ruff==0.9.1; extra == "dev"
+Requires-Dist: ruff==0.9.3; extra == "dev"
 Requires-Dist: textual-dev>=0.86.2; extra == "dev"
 Requires-Dist: types-PyYAML; extra == "dev"
 Requires-Dist: types-beautifulsoup4; extra == "dev"