| """The Observation contract: screen frames and the element index built from them. |
| |
| The harness never asks the model to *find* anything. OCR or the accessibility tree |
| produces a labelled element index (see :mod:`tiny_computer_agent.perception`); the |
| model only chooses among those labels. |
| |
| Every :class:`Element` carries **absolute screen coordinates**, whoever built it. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import sys |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Protocol, Sequence |
|
|
|
|
| def normalise_label(text: str) -> str: |
| """Lower-case and keep only letters and digits. |
| |
| OCR routinely loses or invents spaces ("Untitleddocument-unsavedchanges") and |
| models add punctuation ("Save."), so labels are compared on their alphanumeric |
| core - otherwise nothing a model says ever matches what was seen on screen. |
| """ |
| return "".join(character for character in text.lower() if character.isalnum()) |
|
|
|
|
| @dataclass(frozen=True) |
| class Element: |
| """A clickable thing, located by the platform — not by the model.""" |
|
|
| label: str |
| kind: str = "button" |
| x: int = 0 |
| y: int = 0 |
| width: int = 0 |
| height: int = 0 |
|
|
| def __post_init__(self) -> None: |
| if not self.label.strip(): |
| raise ValueError("Element.label must be non-empty") |
| if min(self.x, self.y, self.width, self.height) < 0: |
| raise ValueError("Element geometry must be non-negative") |
|
|
| @property |
| def center(self) -> tuple[int, int]: |
| if self.width and self.height: |
| return (self.x + self.width // 2, self.y + self.height // 2) |
| return (self.x, self.y) |
|
|
| def matches(self, needle: str) -> bool: |
| wanted = normalise_label(needle) |
| return bool(wanted) and wanted in normalise_label(self.label) |
|
|
|
|
| @dataclass(frozen=True) |
| class Observation: |
| """One frame of the world: pixels plus whatever index we could build.""" |
|
|
| index: int = 0 |
| frame: Any = None |
| elements: tuple[Element, ...] = () |
| text: str = "" |
|
|
| def find(self, label: str) -> Element | None: |
| """Exact label first, then substring; None if nothing matches. |
| |
| Comparison is on normalised text (see :func:`normalise_label`), so |
| ``find("Save")`` still matches a label the OCR rendered as ``"Save ."``. |
| """ |
| wanted = normalise_label(label) |
| if not wanted: |
| return None |
| exact = [e for e in self.elements if normalise_label(e.label) == wanted] |
| if exact: |
| return exact[0] |
| loose = [e for e in self.elements if e.matches(label)] |
| return loose[0] if loose else None |
|
|
| def summary(self, max_elements: int = 40) -> str: |
| """Compact text form — this is what a model would be shown.""" |
| lines = [f"step {self.index}"] |
| if self.text: |
| lines.append("visible text: " + " ".join(self.text.split())[:400]) |
| if self.elements: |
| lines.append("clickable elements:") |
| for e in self.elements[:max_elements]: |
| x, y = e.center |
| lines.append(f" - {e.label!r} ({e.kind}) at {x},{y}") |
| if len(self.elements) > max_elements: |
| lines.append(f" ... and {len(self.elements) - max_elements} more") |
| if not self.elements and not self.text: |
| lines.append("no element index available (pixels only)") |
| return "\n".join(lines) |
|
|
|
|
| class ScreenSource(Protocol): |
| def capture(self) -> Observation: ... |
|
|
|
|
| class FakeScreen: |
| """Scripted screen: each capture() advances one frame. |
| |
| Used to exercise the loop with no display at all. The final frame repeats |
| forever, so a runaway agent cannot exhaust the script and crash. |
| """ |
|
|
| def __init__(self, frames: Sequence[Sequence[Element]], *, text: str = "") -> None: |
| if not frames: |
| raise ValueError("FakeScreen needs at least one frame") |
| self._frames: tuple[tuple[Element, ...], ...] = tuple(tuple(f) for f in frames) |
| self._text = text |
| self.captures = 0 |
|
|
| def capture(self) -> Observation: |
| i = min(self.captures, len(self._frames) - 1) |
| observation = Observation( |
| index=self.captures, |
| frame=f"frame-{i}", |
| elements=self._frames[i], |
| text=self._text, |
| ) |
| self.captures += 1 |
| return observation |
|
|
|
|
| class ImageScreen: |
| """A static screen source: replay one screenshot (or a sequence of them). |
| |
| Lets the full pipeline - OCR, index, menu, model, action - be exercised with no |
| display at all, and lets you replay a captured failure offline. Accepts a path, |
| a PIL image, or a sequence of either. |
| """ |
|
|
| def __init__(self, source: Any) -> None: |
| from PIL import Image |
|
|
| frames = [source] if isinstance(source, (str, Path)) or not isinstance(source, (list, tuple)) else list(source) |
| self._images: list[Any] = [] |
| for frame in frames: |
| if isinstance(frame, (str, Path)): |
| self._images.append(Image.open(str(frame)).convert("RGB")) |
| else: |
| self._images.append(frame) |
| if not self._images: |
| raise ValueError("ImageScreen needs at least one image") |
| self.captures = 0 |
|
|
| def capture(self) -> Observation: |
| index = min(self.captures, len(self._images) - 1) |
| self.captures += 1 |
| return Observation(index=self.captures - 1, frame=self._images[index], elements=(), text="") |
|
|
|
|
| class ScreenUnavailable(RuntimeError): |
| """Raised when real capture is requested but the machine cannot provide it.""" |
|
|
|
|
| class GrabScreen: |
| """Capture the real desktop as an image. |
| |
| Prefers ``mss`` (fast, multi-monitor) and falls back to Pillow's ``ImageGrab``. |
| Caps are deliberately absent: a frame pixel must be a *screen* pixel, which is |
| what lets every detector produce Elements with absolute screen coordinates. |
| |
| Two options break that identity, so both are explicit: ``monitor`` may select a |
| display whose origin is not (0, 0), and ``max_width`` downscales a wide capture |
| (faster OCR, and a smaller image for a model to look at). |
| :meth:`map_to_screen` converts a frame coordinate back to a real one, and |
| :class:`~tiny_computer_agent.perception.IndexedScreen` applies it to every |
| detected element, so a click never lands on the wrong pixel. |
| """ |
|
|
| def __init__(self, *, monitor: int = 1, max_width: int | None = None, |
| scale: float | None = None) -> None: |
| if monitor < 0: |
| raise ValueError("monitor must be 0 (all displays) or a positive index") |
| if max_width is not None and max_width < 1: |
| raise ValueError("max_width must be positive") |
| if scale is not None and scale <= 0: |
| raise ValueError("scale must be positive") |
| self.monitor = monitor |
| self.max_width = max_width |
| self.scale = scale |
| self.origin: tuple[int, int] = (0, 0) |
| self.frame_scale: float = 1.0 |
| self._index = 0 |
|
|
| def capture(self) -> Observation: |
| image, origin = self._grab() |
| native_width = image.width |
| if self.max_width is not None and image.width > self.max_width: |
| ratio = self.max_width / image.width |
| image = image.resize((self.max_width, max(1, round(image.height * ratio)))) |
| if self.scale is not None: |
| self.frame_scale = float(self.scale) |
| else: |
| self.frame_scale = native_width / image.width |
| self.origin = origin |
| observation = Observation(index=self._index, frame=image, elements=(), text="") |
| self._index += 1 |
| return observation |
|
|
| def map_to_screen(self, element: Element) -> Element: |
| """Frame pixel -> absolute screen pixel, honouring the origin and any scale.""" |
| if self.frame_scale == 1.0 and self.origin == (0, 0): |
| return element |
| factor = self.frame_scale |
| return Element( |
| label=element.label, kind=element.kind, |
| x=round(element.x * factor) + self.origin[0], |
| y=round(element.y * factor) + self.origin[1], |
| width=max(1, round(element.width * factor)), |
| height=max(1, round(element.height * factor)), |
| ) |
|
|
| def _grab(self) -> tuple[Any, tuple[int, int]]: |
| """(image, screen position of the image's top-left pixel).""" |
| try: |
| import mss |
| except ImportError: |
| return self._grab_with_pillow() |
| try: |
| from PIL import Image |
|
|
| with mss.mss() as sct: |
| monitors = sct.monitors |
| monitor = monitors[self.monitor] if self.monitor < len(monitors) else monitors[-1] |
| shot = sct.grab(monitor) |
| image = Image.frombytes("RGB", shot.size, shot.bgra, "raw", "BGRX") |
| return image, (int(monitor.get("left", 0)), int(monitor.get("top", 0))) |
| except Exception as exc: |
| raise ScreenUnavailable( |
| f"screen capture failed (no display?): {type(exc).__name__}: {exc}" |
| ) from exc |
|
|
| def _grab_with_pillow(self) -> tuple[Any, tuple[int, int]]: |
| try: |
| from PIL import ImageGrab |
| except ImportError as exc: |
| raise ScreenUnavailable( |
| "screen capture needs mss or Pillow: pip install \'tiny-computer-agent[desktop]\'" |
| ) from exc |
| try: |
| if sys.platform.startswith("win"): |
| image = ImageGrab.grab(all_screens=self.monitor == 0) |
| else: |
| image = ImageGrab.grab() |
| except Exception as exc: |
| raise ScreenUnavailable( |
| f"screen capture failed (no display?): {type(exc).__name__}: {exc}" |
| ) from exc |
| return image.convert("RGB"), (0, 0) |
|
|
