Blame for tiny_computer_agent/observe.py

"""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:  # pragma: no cover - pillow is a desktop extra
            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)