hai-drivers 0.1.0__py3-none-any.whl

hai_drivers/common/exceptions.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+
+class DriverException(Exception):
+    """Base exception for all driver exceptions."""
+
+    error_type: str = "driver_error"
+    description: str = "A driver error occurred."
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "error_type": self.error_type,
+            "description": self.description,
+            "message": self.args,
+        }

hai_drivers/common/types.py

@@ -0,0 +1,205 @@
+from typing import Literal
+
+MouseButton = Literal["left", "right", "middle"]
+
+ScrollDirection = Literal["up", "down", "left", "right"]
+
+SecretField = Literal["username", "password", "email", "totp", "api_key"]
+
+Key = Literal[
+    "\t",
+    "\n",
+    "\r",
+    " ",
+    "!",
+    '"',
+    "#",
+    "$",
+    "%",
+    "&",
+    "'",
+    "(",
+    ")",
+    "*",
+    "+",
+    ",",
+    "-",
+    ".",
+    "/",
+    "0",
+    "1",
+    "2",
+    "3",
+    "4",
+    "5",
+    "6",
+    "7",
+    "8",
+    "9",
+    ":",
+    ";",
+    "<",
+    "=",
+    ">",
+    "?",
+    "@",
+    "[",
+    "\\",
+    "]",
+    "^",
+    "_",
+    "`",
+    "a",
+    "b",
+    "c",
+    "d",
+    "e",
+    "f",
+    "g",
+    "h",
+    "i",
+    "j",
+    "k",
+    "l",
+    "m",
+    "n",
+    "o",
+    "p",
+    "q",
+    "r",
+    "s",
+    "t",
+    "u",
+    "v",
+    "w",
+    "x",
+    "y",
+    "z",
+    "{",
+    "|",
+    "}",
+    "~",
+    "accept",
+    "add",
+    "alt",
+    "altleft",
+    "altright",
+    "apps",
+    "backspace",
+    "browserback",
+    "browserfavorites",
+    "browserforward",
+    "browserhome",
+    "browserrefresh",
+    "browsersearch",
+    "browserstop",
+    "capslock",
+    "clear",
+    "convert",
+    "ctrl",
+    "ctrlleft",
+    "ctrlright",
+    "decimal",
+    "del",
+    "delete",
+    "divide",
+    "down",
+    "end",
+    "enter",
+    "esc",
+    "escape",
+    "execute",
+    "f1",
+    "f10",
+    "f11",
+    "f12",
+    "f13",
+    "f14",
+    "f15",
+    "f16",
+    "f17",
+    "f18",
+    "f19",
+    "f2",
+    "f20",
+    "f21",
+    "f22",
+    "f23",
+    "f24",
+    "f3",
+    "f4",
+    "f5",
+    "f6",
+    "f7",
+    "f8",
+    "f9",
+    "final",
+    "fn",
+    "hanguel",
+    "hangul",
+    "hanja",
+    "help",
+    "home",
+    "insert",
+    "junja",
+    "kana",
+    "kanji",
+    "launchapp1",
+    "launchapp2",
+    "launchmail",
+    "launchmediaselect",
+    "left",
+    "modechange",
+    "multiply",
+    "nexttrack",
+    "nonconvert",
+    "num0",
+    "num1",
+    "num2",
+    "num3",
+    "num4",
+    "num5",
+    "num6",
+    "num7",
+    "num8",
+    "num9",
+    "numlock",
+    "pagedown",
+    "pageup",
+    "pause",
+    "pgdn",
+    "pgup",
+    "playpause",
+    "prevtrack",
+    "print",
+    "printscreen",
+    "prntscrn",
+    "prtsc",
+    "prtscr",
+    "return",
+    "right",
+    "scrolllock",
+    "select",
+    "separator",
+    "shift",
+    "shiftleft",
+    "shiftright",
+    "sleep",
+    "space",
+    "stop",
+    "subtract",
+    "tab",
+    "up",
+    "volumedown",
+    "volumemute",
+    "volumeup",
+    "win",
+    "winleft",
+    "winright",
+    "yen",
+    "command",
+    "option",
+    "optionleft",
+    "optionright",
+    "universal_command",
+]

hai_drivers/desktop/interface.py

@@ -0,0 +1,143 @@
+from abc import ABC, abstractmethod
+from functools import cached_property
+from pathlib import Path
+from typing import Sequence
+
+from hai_drivers.common.types import Key, MouseButton, ScrollDirection
+from hai_drivers.desktop.models import RunCommandResponse
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class DesktopDriverInterface(ABC):
+    @cached_property
+    @abstractmethod
+    def platform(self) -> str:
+        """Get the platform of the desktop."""
+
+    @abstractmethod
+    def screenshot_b64(self) -> str:
+        """Take a screenshot and return it as a base64 string."""
+
+    @abstractmethod
+    def get_screen_size(self) -> tuple[int, int]:
+        """Get the screen dimensions as (width, height)."""
+
+    @abstractmethod
+    def get_mouse_position(self) -> tuple[int, int]:
+        """Get the current mouse position as (x, y)."""
+
+    @abstractmethod
+    def get_accessibility_tree(self) -> str:
+        """Get the accessibility tree of the visible screen as an xml string."""
+
+    @abstractmethod
+    def get_number_of_pixels_scrolled_per_click(self) -> int:
+        """Return the number of px scrolled by a scrolling click.
+
+        This value should either be a constant or cached by the implementation.
+        """
+
+    @abstractmethod
+    def mouse_move_to(self, x: int, y: int):
+        """Move the mouse to the coordinates."""
+
+    @abstractmethod
+    def mouse_press(self, button: MouseButton):
+        """Press a mouse button down without releasing it."""
+
+    @abstractmethod
+    def mouse_release(self, button: MouseButton):
+        """Release a previously pressed mouse button."""
+
+    @abstractmethod
+    def click(self, x: int, y: int, button: MouseButton = "left"):
+        """Click at the specified screen coordinates."""
+
+    @abstractmethod
+    def double_click(self, x: int, y: int, button: MouseButton = "left", delay_between_clicks: float = 0.05):
+        """Double click at the specified screen coordinates."""
+
+    @abstractmethod
+    def write(self, text: str, delay_between_keys: float = 0.05) -> None:
+        """Type a string of text.
+        Capital letters and other special characters will be translated into the correct keycode sequence.
+
+        delay_between_keys: secs between each key press.
+        """
+
+    @abstractmethod
+    def hotkey(self, keys: list[Key]):
+        """Performs key down presses on the arguments passed in order, then performs key releases in reverse order."""
+
+    @abstractmethod
+    def tap_key(self, key: Key):
+        """Tap a keyboard key. Performs both press and release of the key."""
+
+    @abstractmethod
+    def press_key(self, key: Key):
+        """Press a keyboard key down without releasing it.
+
+        Behaviour of pressing a Key might be OS specific. e.g Ctrl vs Command
+        """
+
+    @abstractmethod
+    def release_key(self, key: Key):
+        """Release a previously pressed keyboard key."""
+
+    @abstractmethod
+    def scroll_by_n_clicks(self, direction: ScrollDirection, clicks: int):
+        """Scroll in a direction by a number of atomic scrolling click.
+
+        The size of a click differs per platform, you can get it with get_number_of_pixels_scrolled_per_click().
+        """
+
+    def scroll_by_n_pxs(self, direction: ScrollDirection, pixels: int):
+        """Scroll in a direction by at least `pixels`, 1 px ~= 0.26 mm.
+
+        If pixels is not a multiple of get_number_of_pixels_scrolled_per_click(), the n pixels scrolled will be higher than pixels.
+        """
+        npixels_per_click = self.get_number_of_pixels_scrolled_per_click()
+        self.scroll_by_n_clicks(
+            direction=direction,
+            clicks=(pixels // npixels_per_click) + (1 if pixels % npixels_per_click > 0 else 0),
+        )
+
+    @abstractmethod
+    def run_command(
+        self,
+        command: Sequence[str],
+        timeout: int | None = 60,
+        env: dict[str, str] | None = None,
+        cwd: Path | None = None,
+        detach: bool = False,
+        ignore_errors: bool = False,
+    ) -> RunCommandResponse:
+        """Run a shell command on the desktop."""
+
+    @abstractmethod
+    def read_file(self, path: str) -> bytes:
+        """Read a file from the desktop filesystem. Returns raw bytes."""
+
+    @abstractmethod
+    def write_file(self, path: str, content: bytes) -> None:
+        """Write raw bytes to a file on the desktop filesystem."""
+
+    def get_observation_snapshot(self) -> "ObservationSnapshot":
+        """Bundle the data needed to build an observation in one round-trip.
+
+        Drivers with cheap AX trees or window metadata override this so
+        observations don't pay for multiple sequential round-trips.
+        """
+        return ObservationSnapshot(
+            screenshot_b64=self.screenshot_b64(),
+            cursor_position=self.get_mouse_position(),
+        )
+
+
+class ObservationSnapshot(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    screenshot_b64: str
+    cursor_position: tuple[int, int] | None = None
+    accessibility_tree: str | None = None
+    metadata: dict[str, str] = Field(default_factory=dict)

hai_drivers/desktop/local/__init__.py

@@ -0,0 +1,5 @@
+"""Local desktop driver."""
+
+from hai_drivers.desktop.local.driver import LocalDesktopDriver
+
+__all__ = ["LocalDesktopDriver"]

hai_drivers/desktop/local/driver.py

@@ -0,0 +1,306 @@
+"""Local desktop driver using pyautogui and pynput."""
+
+import base64
+import io
+import logging
+import platform
+import time
+from pathlib import Path
+from typing import Sequence
+
+import pyautogui
+from hai_drivers.common.types import Key, MouseButton, ScrollDirection
+from hai_drivers.desktop.interface import DesktopDriverInterface, ObservationSnapshot
+from hai_drivers.desktop.models import RunCommandResponse
+from hai_drivers.desktop.utils import run_command_impl
+from PIL import Image as PILImage
+from pynput.keyboard import Controller
+from pynput.keyboard import Key as PynputKey
+
+logger = logging.getLogger(__name__)
+MACOS_CLICK_TO_KEYBOARD_SETTLE_S = 0.2
+
+
+class LocalDesktopDriver(DesktopDriverInterface):
+    """Local desktop driver with pyautogui and pynput."""
+
+    def __init__(self, *, screenshot_max_width: int | None = None):
+        """Create a pyautogui desktop driver.
+
+        Args:
+            screenshot_max_width: Downscale model-facing screenshots wider than
+                this before upload (pyautogui captures native HiDPI). ``None``
+                keeps native resolution.
+        """
+        if screenshot_max_width is not None and screenshot_max_width <= 0:
+            raise ValueError(f"screenshot_max_width must be a positive int or None; got {screenshot_max_width}")
+
+        pyautogui.FAILSAFE = False
+        pyautogui.PAUSE = 0
+        _patch_pyautogui_macos_multiclick()
+        self._keyboard = Controller()
+        self._screenshot_max_width: int | None = screenshot_max_width
+        # pyautogui itself assumes a fixed primary display for the process
+        # lifetime, so cache the logical screen size under the same assumption.
+        size = pyautogui.size()
+        self._screen_size: tuple[int, int] = (size.width, size.height)
+
+    @property
+    def platform(self) -> str:
+        """Get the platform of the desktop."""
+        return platform.system()
+
+    def screenshot_b64(self) -> str:
+        """Take a native-resolution screenshot and return it as a base64-encoded PNG string."""
+        return base64.b64encode(_encode_png(pyautogui.screenshot())).decode("ascii")
+
+    def get_observation_snapshot(self) -> ObservationSnapshot:
+        """Capture one model-facing screenshot and cursor position.
+
+        pyautogui can report cursor/screen coordinates in logical pixels while
+        screenshots arrive in physical HiDPI pixels. Cursor coordinates are
+        first mapped into the screenshot image frame, then scaled again if the
+        image is downsampled. The capture stays a live PIL image until the
+        single encode at output size.
+        """
+        captured = pyautogui.screenshot()
+        cursor = self._cursor_position_in_screenshot(
+            cursor_xy=self.get_mouse_position(),
+            screenshot_size=(captured.width, captured.height),
+        )
+        if self._screenshot_max_width is not None and captured.width > self._screenshot_max_width:
+            image, out_cursor = _resize_image_to_width(captured, self._screenshot_max_width, cursor)
+        else:
+            image, out_cursor = captured, cursor
+        return ObservationSnapshot(
+            screenshot_b64=base64.b64encode(_encode_png(image)).decode("ascii"),
+            cursor_position=out_cursor,
+        )
+
+    def _cursor_position_in_screenshot(
+        self, *, cursor_xy: tuple[int, int], screenshot_size: tuple[int, int]
+    ) -> tuple[int, int]:
+        """Map pyautogui cursor coordinates into the screenshot pixel frame."""
+        screen_w, screen_h = self.get_screen_size()
+        image_w, image_h = screenshot_size
+        scale_x = image_w / screen_w
+        scale_y = image_h / screen_h
+        return (round(cursor_xy[0] * scale_x), round(cursor_xy[1] * scale_y))
+
+    def get_screen_size(self) -> tuple[int, int]:
+        """Get the screen dimensions as (width, height)."""
+        return self._screen_size
+
+    def get_mouse_position(self) -> tuple[int, int]:
+        """Get the current mouse position as (x, y)."""
+        pos = pyautogui.position()
+        return (pos.x, pos.y)
+
+    def get_accessibility_tree(self) -> str:
+        """Get the accessibility tree of the visible screen as an xml string.
+
+        Note: This requires platform-specific implementations.
+        Not currently supported by pyautogui/pynput.
+        """
+        raise NotImplementedError(
+            "Accessibility tree not available with pyautogui/pynput. "
+            "Use platform-specific drivers for accessibility tree access."
+        )
+
+    def get_number_of_pixels_scrolled_per_click(self) -> int:
+        """Return the number of px scrolled by a scrolling click.
+
+        - Windows: 120 (WHEEL_DELTA: https://learn.microsoft.com/en-us/windows/win32/inputdev/wm-mousewheel)
+        - Linux: 53 (empirical GTK/X11 default, varies by DE)
+        - macOS: 20 (empirical, varies due to smooth scrolling)
+        """
+        return {"Darwin": 20, "Windows": 120}.get(platform.system(), 53)
+
+    def mouse_move_to(self, x: int, y: int):
+        """Move the mouse to the coordinates."""
+        pyautogui.moveTo(x, y)
+
+    def mouse_press(self, button: MouseButton):
+        """Press a mouse button down without releasing it."""
+        pyautogui.mouseDown(button=button)
+
+    def mouse_release(self, button: MouseButton):
+        """Release a previously pressed mouse button."""
+        pyautogui.mouseUp(button=button)
+
+    def click(self, x: int, y: int, button: MouseButton = "left"):
+        """Click at the specified screen coordinates."""
+        pyautogui.click(x=x, y=y, button=button)
+        _settle_after_click_for_focus()
+
+    def double_click(self, x: int, y: int, button: MouseButton = "left", delay_between_clicks: float = 0.05):
+        """Double click at the specified screen coordinates."""
+        pyautogui.doubleClick(x=x, y=y, button=button, interval=delay_between_clicks)
+        _settle_after_click_for_focus()
+
+    def write(self, text: str, delay_between_keys: float = 0.05):
+        """Type a string of text.
+
+        Capital letters and other special characters will be translated into the correct keycode sequence.
+        delay_between_keys: secs between each key press.
+        """
+        if delay_between_keys > 0:
+            for char in text:
+                self._keyboard.type(char)
+                time.sleep(delay_between_keys)
+        else:
+            self._keyboard.type(text)
+
+    def press_key(self, key: Key):
+        """Press a keyboard key down without releasing it.
+
+        Behaviour of pressing a Key might be OS specific. e.g Ctrl vs Command
+        """
+        key_parsed = self._parse_pynput_key(key)
+        self._keyboard.press(key_parsed)
+
+    def release_key(self, key: Key):
+        """Release a previously pressed keyboard key."""
+        key_parsed = self._parse_pynput_key(key)
+        self._keyboard.release(key_parsed)
+
+    def tap_key(self, key: Key):
+        """Tap a keyboard key. Performs both press and release of the key."""
+        self.press_key(key)
+        self.release_key(key)
+
+    def hotkey(self, keys: list[Key]):
+        """Performs key down presses on the arguments passed in order, then performs key releases in reverse order."""
+        for key in keys:
+            self.press_key(key)
+        for key in reversed(keys):
+            self.release_key(key)
+
+    def scroll_by_n_clicks(self, direction: ScrollDirection, clicks: int):
+        """Scroll in a direction by a number of atomic scrolling click.
+
+        The size of a click differs per platform, you can get it with get_number_of_pixels_scrolled_per_click().
+        """
+        if direction == "up":
+            pyautogui.scroll(clicks)
+        elif direction == "down":
+            pyautogui.scroll(-clicks)
+        elif direction == "left":
+            pyautogui.hscroll(-clicks)
+        elif direction == "right":
+            pyautogui.hscroll(clicks)
+        else:
+            raise ValueError(f"Invalid scroll direction: {direction}")
+
+    def _parse_pynput_key(self, key: str) -> PynputKey | str:
+        """Parse a key string to a pynput.keyboard.Key enum or a string."""
+        if key in PynputKey.__members__:
+            return PynputKey[key]
+        if key == "command":
+            return PynputKey.cmd
+        return key
+
+    def run_command(
+        self,
+        command: Sequence[str],
+        timeout: int | None = 60,
+        env: dict[str, str] | None = None,
+        cwd: Path | None = None,
+        detach: bool = False,
+        ignore_errors: bool = False,
+    ) -> RunCommandResponse:
+        """Run a shell command on the local desktop using subprocess."""
+        return run_command_impl(
+            command=command,
+            timeout=timeout,
+            env=env,
+            cwd=cwd,
+            detach=detach,
+        )
+
+    def read_file(self, path: str) -> bytes:
+        """Read a file from the local filesystem."""
+        p = Path(path).expanduser()
+        if not p.exists():
+            raise FileNotFoundError(f"File not found: {path}")
+        return p.read_bytes()
+
+    def write_file(self, path: str, content: bytes) -> None:
+        """Write a file to the local filesystem."""
+        p = Path(path).expanduser()
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.write_bytes(content)
+
+
+def _resize_image_to_width(
+    image: PILImage.Image, width: int, cursor_xy: tuple[int, int] | None
+) -> tuple[PILImage.Image, tuple[int, int] | None]:
+    """Resize the image to ``width`` and scale the cursor to match.
+
+    Operates on the live PIL image so callers encode exactly once, at output size.
+    """
+    scale = width / image.width
+    new_size = (width, max(1, round(image.height * scale)))
+    resized = image.resize(new_size, PILImage.Resampling.LANCZOS)
+    out_cursor = (round(cursor_xy[0] * scale), round(cursor_xy[1] * scale)) if cursor_xy is not None else None
+    return resized, out_cursor
+
+
+def _encode_png(image: PILImage.Image) -> bytes:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return buffer.getvalue()
+
+
+def _settle_after_click_for_focus() -> None:
+    if platform.system() == "Darwin":
+        time.sleep(MACOS_CLICK_TO_KEYBOARD_SETTLE_S)
+
+
+def _patch_pyautogui_macos_multiclick() -> None:
+    """Apply PyAutoGUI PR #949's macOS click-state fix."""
+    if platform.system() != "Darwin":
+        return
+
+    platform_module = getattr(pyautogui, "platformModule", None)
+    if platform_module is None:
+        logger.info("Skipping macOS PyAutoGUI multi-click patch: platformModule is unavailable")
+        return
+
+    if getattr(platform_module, "_hai_drivers_macos_multiclick_patch", False):
+        logger.info("Skipping macOS PyAutoGUI multi-click patch: patch is already applied")
+        return
+
+    quartz = getattr(platform_module, "Quartz", None)
+    if quartz is None:
+        logger.info("Skipping macOS PyAutoGUI multi-click patch: Quartz is unavailable")
+        return
+
+    def _button_events(button: MouseButton):
+        if button == platform_module.LEFT:
+            return quartz.kCGMouseButtonLeft, quartz.kCGEventLeftMouseDown, quartz.kCGEventLeftMouseUp
+        if button == platform_module.MIDDLE:
+            return quartz.kCGMouseButtonCenter, quartz.kCGEventOtherMouseDown, quartz.kCGEventOtherMouseUp
+        if button == platform_module.RIGHT:
+            return quartz.kCGMouseButtonRight, quartz.kCGEventRightMouseDown, quartz.kCGEventRightMouseUp
+        raise AssertionError("button argument not in ('left', 'middle', 'right')")
+
+    def _send_mouse_event(event_type, x: int, y: int, button, click_state: int | None = None) -> None:
+        mouse_event = quartz.CGEventCreateMouseEvent(None, event_type, (x, y), button)
+        if click_state is not None:
+            quartz.CGEventSetIntegerValueField(mouse_event, quartz.kCGMouseEventClickState, click_state)
+        quartz.CGEventPost(quartz.kCGHIDEventTap, mouse_event)
+
+    def _multi_click(x: int, y: int, button: MouseButton, num: int, interval: float = 0.0) -> None:
+        native_button, down, up = _button_events(button)
+        for index in range(num):
+            pyautogui.failSafeCheck()
+            click_state = index + 1
+            platform_module._sendMouseEvent(down, x, y, native_button, click_state)
+            platform_module._sendMouseEvent(up, x, y, native_button, click_state)
+            time.sleep(interval)
+
+    platform_module._sendMouseEvent = _send_mouse_event
+    platform_module._multiClick = _multi_click
+    platform_module._hai_drivers_macos_multiclick_patch = True
+    logger.info("Applied macOS PyAutoGUI multi-click patch")

hai_drivers/desktop/models.py

@@ -0,0 +1,12 @@
+"""Models shared across desktop drivers."""
+
+from pydantic import BaseModel
+
+
+class RunCommandResponse(BaseModel):
+    """Response from running a shell command."""
+
+    returncode: int
+    stdout: str
+    stderr: str
+    exception: str | None = None