Blame for tiny_computer_agent/brain.py

"""Decision making: menu-mediated, model-agnostic.

The design lesson from testing tiny models on real screenshots: a 256M VLM can
*read* a button label from pixels but cannot reliably reason, and a 0.5B text model
writes perfect JSON but picks the wrong widget when handed noisy free-form text.
Both do much better when the harness hands them a **short numbered menu** and asks
for a number. That is what this module builds, and it keeps the model's output space
finite no matter how small the model is.

Brains shipped here:

* :class:`MockBrain`    - replays a scripted action list (tests, demos, replays)
* :class:`PlannerBrain` - rule-based, no model: clicks the element matching the goal
* :class:`ModelBrain`   - base for real models: menu, prompt, lenient parsing
  (see :mod:`tiny_computer_agent.models` for the text and vision backends)
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any, Iterator, Mapping, Protocol, Sequence

from tiny_computer_agent.actions import ACTION_KINDS, Action, ActionError, parse_action
from tiny_computer_agent.decoding import (
    brevity_ladder,
    choice_token_ids,
    digit_token_ids,
    majority_choice,
    pair_token_ids,
    read_choice,
    read_verdict,
)
from tiny_computer_agent.observe import Element, Observation, normalise_label


class BrainError(RuntimeError):
    """Raised when a brain cannot produce a valid action."""


class GoalUnreachable(BrainError):
    """Raised by a brain that can decide nothing further, without claiming success.

    Deliberately distinct from a crash: the harness records the episode as a
    *failure*, because a brain admitting "I give up" must never be logged as
    "goal met" - and the harness, not the brain, owns that judgement.
    """


def extract_json_objects(raw: str) -> list[dict[str, Any]]:
    """Every flat JSON object in a reply, in the order they appear.

    Regex-based and non-nesting: replies from tiny models are one-level objects.
    Callers should prefer the **last** object, because a model that echoes its own
    prompt before answering puts the prompt's example first and its answer last.
    """
    objects: list[dict[str, Any]] = []
    for match in re.finditer(r"\{[^{}]*\}", raw or "", re.DOTALL):
        try:
            payload = json.loads(match.group(0))
        except json.JSONDecodeError:
            continue
        if isinstance(payload, dict):
            objects.append(payload)
    return objects


def extract_json_object(raw: str) -> dict[str, Any]:
    """The decision out of a reply: the *last* JSON object, not the first."""
    objects = extract_json_objects(raw)
    if not objects:
        raise BrainError(f"no JSON object in the model reply: {(raw or '')[:200]!r}")
    return objects[-1]


@dataclass(frozen=True)
class MenuItem:
    """One numbered choice shown to the model."""

    index: int
    element: Element

    @property
    def label(self) -> str:
        return self.element.label


class DecisionMenu:
    """A short, numbered, closed choice set built from the element index.

    Elements are presented in reading order (top to bottom, then left to right),
    and capped: a 20-item menu is as good as free-form text to a small model.
    """

    def __init__(self, observation: Observation, *, max_items: int = 12) -> None:
        ordered = sorted(observation.elements, key=lambda e: (e.y, e.x))[:max_items]
        self.items: tuple[MenuItem, ...] = tuple(
            MenuItem(index=i + 1, element=e) for i, e in enumerate(ordered)
        )

    def __len__(self) -> int:
        return len(self.items)

    def __bool__(self) -> bool:
        return bool(self.items)

    def __iter__(self) -> Iterator[MenuItem]:
        return iter(self.items)

    def item(self, value: Any) -> MenuItem | None:
        """Look up a 1-based index, tolerating the strings models actually send."""
        try:
            number = int(float(str(value).strip()))
        except (TypeError, ValueError):
            return None
        if 1 <= number <= len(self.items):
            return self.items[number - 1]
        return None

    def find(self, label: str) -> Element | None:
        """Resolve a label against the menu (normalised, like Observation.find)."""
        wanted = normalise_label(label)
        if not wanted:
            return None
        for item in self.items:
            if normalise_label(item.label) == wanted:
                return item.element
        for item in self.items:
            if wanted in normalise_label(item.label):
                return item.element
        return None

    def mentioned_in(self, text: str) -> tuple[MenuItem, ...]:
        """Menu entries whose label appears in the text, in order of first mention."""
        haystack = normalise_label(text)
        found: list[tuple[int, MenuItem]] = []
        for item in self.items:
            needle = normalise_label(item.label)
            if needle and (position := haystack.find(needle)) != -1:
                found.append((position, item))
        return tuple(item for _, item in sorted(found, key=lambda pair: pair[0]))

    def render(self) -> str:
        return "\n".join(
            f"  {item.index}. {item.label}  (at {item.element.center[0]},{item.element.center[1]})"
            for item in self.items
        )


def _recover_from_text(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None:
    """Salvage a decision from a payload that named the right thing in the wrong field.

    Observed for real: a 0.5B planner replied ``{"kind": "click", "index": 0,
    "reason": "Save"}`` - the correct answer, in the wrong place. If the index is
    unusable but exactly one menu label is mentioned anywhere in the payload, take
    the label. It is still only ever a label the detector saw in this frame, so the
    deny list, the dry run and the step budget all continue to apply.
    """
    text = json.dumps(payload, ensure_ascii=False)
    mentioned = menu.mentioned_in(text)
    if len({item.label for item in mentioned}) == 1:
        item = mentioned[0]
        return Action(
            "click", label=item.label,
            reason=f"recovered {item.label!r} from the model's reply: {text[:80]}",
        )
    return None


def _action_from_payload(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None:
    """Interpret one JSON object, tolerating the ways small models mangle the shape."""
    kind = payload.get("kind")
    if isinstance(kind, str) and kind.strip().lower() not in ACTION_KINDS:
        kind = None  # e.g. {"kind": "file"} - a garbled click; fall back to index/label

    if payload.get("index") is not None:
        item = menu.item(payload["index"])
        if item is not None:
            return Action("click", label=item.label, reason=f"menu index {payload['index']}")
        recovered = _recover_from_text(payload, menu)
        if recovered is not None:
            return recovered
        raise BrainError(
            f"the model chose index {payload['index']!r}, which is not in the "
            f"menu of {len(menu)} items"
        )

    if kind in (None, "click") and "label" in payload:
        element = menu.find(str(payload["label"]))
        if element is None:
            raise BrainError(f"the model named {payload['label']!r}, which is not on screen")
        return Action("click", label=element.label, reason=f"the model named {element.label!r}")

    if kind is not None:
        try:
            return parse_action(dict(payload, kind=kind))
        except ActionError as exc:
            # a valid-looking verb with no usable target: maybe the label is in the prose
            return _recover_from_text(payload, menu) or (_raise_brain_error(exc))

    return _recover_from_text(payload, menu)


def _raise_brain_error(exc: Exception) -> Action:
    raise BrainError(f"the model proposed an invalid action: {exc}") from exc


_DONE = re.compile(r"\b(done|finished|complete[d]?)\b", re.IGNORECASE)


def parse_decision(raw: str, observation: Observation, *, max_menu_items: int = 12) -> Action:
    """Turn a small model's reply into a validated Action - leniently, but safely.

    Tried in order:

    1. the **last** JSON object in the reply (``index``, ``label`` or a full action)
    2. a label the reply names (most specific evidence of intent)
    3. a bare menu number
    4. a "done" word

    A reply that mentions three or more menu entries is rejected as an echo of the
    prompt rather than a decision - the failure mode of the weakest models, and the
    one most likely to make an agent click something arbitrary.
    """
    text = " ".join((raw or "").split())
    if not text:
        raise BrainError("the model returned an empty reply")
    menu = DecisionMenu(observation, max_items=max_menu_items)

    for payload in reversed(extract_json_objects(text)):
        action = _action_from_payload(payload, menu)
        if action is not None:
            return action

    mentioned = menu.mentioned_in(text)
    if len(mentioned) >= 3:
        raise BrainError(
            f"the reply mentions {len(mentioned)} menu entries, so it looks like an echo "
            f"of the prompt rather than a decision: {text[:120]!r}"
        )
    if mentioned:
        item = mentioned[0]
        note = f"the reply names {item.label!r}"
        if len(mentioned) > 1:
            note += f" (ambiguous: {len(mentioned)} labels mentioned; took the first)"
        return Action("click", label=item.label, reason=note)

    number = re.search(r"\b(\d{1,2})\b", text)
    if number is not None:
        item = menu.item(number.group(1))
        if item is not None:
            return Action(
                "click", label=item.label,
                reason=f"menu index {number.group(1)} from reply {text[:40]!r}",
            )

    if _DONE.search(text):
        return Action("done", reason=_DONE.search(text).group(0).lower())

    raise BrainError(f"could not read a decision out of the reply: {text[:120]!r}")


class Brain(Protocol):
    """Anything that can turn (goal, observation) into one validated action."""

    name: str

    def decide(self, goal: str, observation: Observation) -> Action: ...


class MockBrain:
    """Replays a fixed script of actions, exactly once. Deterministic and offline."""

    name = "mock"

    def __init__(self, script: Sequence[Action | Mapping[str, Any]]) -> None:
        self._script: list[Action] = [
            a if isinstance(a, Action) else parse_action(a) for a in script
        ]
        self._cursor = 0

    @property
    def remaining(self) -> int:
        return max(0, len(self._script) - self._cursor)

    def decide(self, goal: str, observation: Observation) -> Action:
        if self._cursor >= len(self._script):
            return Action("done", reason="script exhausted")
        action = self._script[self._cursor]
        self._cursor += 1
        return action


class PlannerBrain:
    """Rule-based baseline: click the element matching the goal, then stop.

    Its purpose is threefold - it proves the harness works end to end with no model
    at all, it is a fallback when a model download is undesirable, and it shows how
    much of a "computer use" task is really an indexing problem. It waits when the
    target is not on screen yet, refuses to click the same label twice, and gives up
    rather than flailing. It will not generalise the way a model can.
    """

    name = "planner"

    def __init__(self, target: str | None = None, *,
                 done_when_text_contains: str | None = None,
                 max_waits: int = 3, wait_seconds: float = 0.5) -> None:
        self.target = target
        self.done_when_text_contains = done_when_text_contains
        self.max_waits = max_waits
        self.wait_seconds = wait_seconds
        self._clicked: set[str] = set()
        self._waits = 0

    def decide(self, goal: str, observation: Observation) -> Action:
        if self.done_when_text_contains:
            needle = self.done_when_text_contains.lower()
            if needle in observation.text.lower():
                return Action("done", reason=f"page text contains {self.done_when_text_contains!r}")

        target = self.target or self._target_from(goal)
        if target:
            element = observation.find(target)
            if element is None:
                if self._waits < self.max_waits:
                    self._waits += 1
                    return Action("wait", seconds=self.wait_seconds,
                                  reason=f"{target!r} is not on screen yet")
                raise GoalUnreachable(f"{target!r} never appeared on screen")
            key = element.label.strip().lower()
            if key in self._clicked:
                return Action("done", reason=f"already clicked {element.label!r}; goal assumed met")
            self._clicked.add(key)
            return Action("click", label=element.label,
                          reason=f"matched {target!r} in the element index")
        raise GoalUnreachable(f"no element name found in the goal {goal!r}")

    @staticmethod
    def _target_from(goal: str) -> str | None:
        """Best effort: quoted text wins, else the last capitalised word."""
        quoted = re.search(r"[\"\']([^\"\']{2,40})[\"\']", goal)
        if quoted:
            return quoted.group(1)
        words = [w for w in goal.split() if w[:1].isupper()]
        return words[-1].strip(".,!?") if words else None


class ModelBrain:
    """Shared plumbing for model-backed brains: menu, prompt, lenient parsing.

    Subclasses implement :meth:`_generate`. ``requires_index`` marks brains that
    cannot decide without an element index (a text-only planner): for those, an
    empty screen is a hard stop rather than a prompt.
    """

    name = "model"
    requires_index = False
    # a backend that can mask its own logits opts in, and then decision_mode
    # "choice" is available: ask for a menu number instead of asking for JSON
    supports_constrained = False
    # "json"   - free text, parsed leniently (the original behaviour)
    # "choice" - one masked generation, a menu number
    # "scan"   - one yes/no question per item, masked, first yes wins
    # "rank"   - the same questions, but ranked by the yes/no log-odds margin,
    #            which does not reward whichever item happened to be asked first
    decision_mode = "json"

    def __init__(self, model_id: str, *, max_menu_items: int = 12,
                 max_new_tokens: int = 64, decision_mode: str = "json",
                 retries: int = 3, votes: int = 1) -> None:
        self.model_id = model_id
        self.max_menu_items = max_menu_items
        self.max_new_tokens = max_new_tokens
        self.decision_mode = decision_mode
        self.retries = max(1, retries)
        self.votes = max(1, votes)
        self.last_choice: int | None = None
        self.echoes_refused = 0
        self.last_prompt: str | None = None
        self.last_reply: str | None = None
        # "what I already tried", newest last; each entry's outcome is filled in at
        # the next decision, when the new observation can be compared to the old one
        self.history: list[str] = []
        self._last_fingerprint: tuple | None = None

    @staticmethod
    def _fingerprint(observation: Observation) -> tuple:
        return (tuple(element.label for element in observation.elements),
                observation.text[:200])

    def build_prompt(self, goal: str, menu: DecisionMenu, observation: Observation,
                     *, history: Sequence[str] = (), screen_changed: bool = True,
                     brevity: int = 0) -> str:
        """A short prompt with a strict output shape.

        Deliberately contains placeholders (``<number>``) rather than worked
        examples: the weakest models echo the prompt back, and an echoed
        placeholder fails validation instead of being mistaken for a decision.

        ``brevity`` is the retry ladder: echoing is driven by how much text the
        model has to copy, so each step upwards drops a section - first the screen
        text and the history, then everything but the menu.
        """
        if self.decision_mode == "choice":
            return self._choice_prompt(goal, menu, observation,
                                       history=history, brevity=brevity)
        lines = [
            "You are controlling a computer. Choose exactly one action.",
            "",
            "Reply with one JSON object and nothing else, in this shape:",
            '  {"kind": "click", "index": <number from the list>, "reason": "<brief>"}',
            '  {"kind": "done", "reason": "<brief>"}',
            "Use done only when the goal is already met.",
            "",
            f"Goal: {goal}",
        ]
        if observation.text:
            lines += ["", f"Text on screen: {observation.text[:300]}"]
        lines += ["", f"Clickable items ({len(menu)}):"]
        lines.append(menu.render() if menu else "  (none detected)")
        if history:
            lines += ["", "Already attempted - do NOT repeat these, they did not work:"]
            lines += [f"  - {entry}" for entry in history]
        if not screen_changed:
            lines += ["", "The screen did not change after your last action. "
                          "Choose a DIFFERENT action."]
        return "\n".join(lines)

    def _choice_prompt(self, goal: str, menu: DecisionMenu, observation: Observation,
                       *, history: Sequence[str] = (), brevity: int = 0) -> str:
        """The constrained-mode prompt: one number, and as little text as possible.

        There is no JSON shape to copy and no worked example - the whole point of
        the masked path is that the model has nothing to echo.
        """
        lines = [
            "You are controlling a computer. Reply with ONE NUMBER and nothing else.",
            "0 means the goal is already met. Any other number is the item to click.",
            "",
            f"Goal: {goal}",
        ]
        if brevity == 0 and observation.text:
            lines += ["", f"Text on screen: {observation.text[:200]}"]
        lines += ["", f"Items ({len(menu)}):", menu.render() if menu else "  (none detected)"]
        if brevity == 0 and history:
            lines += ["", "Already tried, and it did not work:"]
            lines += [f"  - {entry}" for entry in history]
        return "\n".join(lines)

    def decide(self, goal: str, observation: Observation) -> Action:
        menu = DecisionMenu(observation, max_items=self.max_menu_items)
        if self.requires_index and not menu:
            raise GoalUnreachable(
                "no clickable elements on screen - is screen capture and OCR working?"
            )
        fingerprint = self._fingerprint(observation)
        screen_changed = (self._last_fingerprint is None
                          or fingerprint != self._last_fingerprint)

        # Finalise the previous attempt first: only now do we know whether it changed
        # anything, and "you already tried this and nothing happened" is the single most
        # useful thing to tell a small model.
        if self.history:
            previous = self.history[-1].split(" -> ")[0]
            self.history[-1] = (
                f"{previous} -> "
                + ("the screen changed" if screen_changed else "the screen did NOT change")
            )

        self.last_prompt = self.build_prompt(
            goal, menu, observation, history=tuple(self.history),
            screen_changed=screen_changed,
        )
        if self.decision_mode == "rank" and self.supports_constrained:
            action = self._decide_by_ranking(goal, menu, observation)
        elif self.decision_mode == "scan" and self.supports_constrained:
            action = self._decide_by_scanning(goal, menu, observation)
        elif self.decision_mode == "choice" and self.supports_constrained:
            action = self._decide_by_choice(goal, menu, observation)
        else:
            self.last_reply = self._generate(self.last_prompt, observation)
            action = parse_decision(self.last_reply, observation,
                                    max_menu_items=self.max_menu_items)
        self.history.append(action.describe())
        del self.history[:-3]
        self._last_fingerprint = fingerprint
        return action

    def _decide_by_choice(self, goal: str, menu: DecisionMenu,
                          observation: Observation) -> Action:
        """Ask for a menu number with every other token masked out.

        The echo failure cannot happen here: the model is unable to emit the prompt
        back, only one of the numbers we allow. When even that fails, the prompt
        gets shorter (echoing is driven by prompt length) and, if ``votes`` is set,
        the answer is taken by majority over several samples.
        """
        if not menu:
            raise GoalUnreachable("no clickable items to choose from")
        tokenizer = self._tokenizer_for_choice()
        allowed = choice_token_ids(tokenizer, len(menu))
        if not allowed:
            allowed = digit_token_ids(tokenizer)
        if not allowed:
            raise BrainError(
                "this tokenizer cannot express a menu number as a single token, so "
                "the choice cannot be constrained; use decision_mode='json'"
            )

        last_error: BrainError | None = None
        for brevity in brevity_ladder(self.retries):
            self.last_prompt = self.build_prompt(
                goal, menu, observation, history=tuple(self.history), brevity=brevity)
            replies = [
                self._generate_constrained(self.last_prompt, observation, allowed, 4,
                                           sample=index > 0)
                for index in range(self.votes)
            ]
            self.last_reply = replies[0] if replies else ""
            choices = [read_choice(reply, len(menu)) for reply in replies]
            number = majority_choice(choices)
            if number is None:
                self.echoes_refused += sum(1 for reply in replies
                                           if reply and reply.strip() == self.last_prompt.strip())
                last_error = BrainError(
                    f"the constrained reply was not a menu number: "
                    f"{self.last_reply[:80]!r} (brevity {brevity}, {len(replies)} vote(s))"
                )
                continue
            self.last_choice = number
            if number == 0:
                return Action("done", reason="the model chose 0: the goal is already met")
            item = menu.item(number)
            if item is None:                                   # pragma: no cover
                last_error = BrainError(f"the model chose {number}, outside the menu")
                continue
            return Action("click", label=item.label,
                          reason=f"constrained choice {number} of {len(menu)}")

        raise last_error or BrainError("the model produced no constrained choice")

    def _scan_prompt(self, goal: str, label: str) -> str:
        """The prompt for one binary question: does *this* item do the job?

        Short on purpose. Every word a tiny model is given is a word it can copy
        instead of answering, and there is nothing here worth copying back.
        """
        return "\n".join([
            "You are controlling a computer. Answer with one word: yes or no.",
            "",
            f"Goal: {goal}",
            f"Candidate item: {label}",
            "Does clicking this item achieve the goal?",
        ])

    def _decide_by_scanning(self, goal: str, menu: DecisionMenu,
                            observation: Observation) -> Action:
        """Ask one masked yes/no question per menu item, and take the first yes.

        Why this rather than one N-way choice: a 256M model cannot hold twelve
        options in mind, but "does this button do what the goal says?" is a
        question it can answer. Each answer is one token out of two, so the mask
        makes the reply space as small as it can possibly be - and the scan stops
        at the first yes, so it is usually cheaper than it looks.
        """
        if not menu:
            raise GoalUnreachable("no clickable items to choose from")
        tokenizer = self._tokenizer_for_choice()
        yes_ids, no_ids = pair_token_ids(tokenizer)
        if not yes_ids or not no_ids:
            raise BrainError(
                "this tokenizer cannot express yes/no as single tokens, so the scan "
                "cannot be constrained; use decision_mode='json'"
            )
        allowed = sorted(set(yes_ids) | set(no_ids))
        scanned: list[tuple[str, str | None]] = []
        for item in menu:
            self.last_prompt = self._scan_prompt(goal, item.label)
            self.last_reply = self._generate_constrained(
                self.last_prompt, observation, allowed, 1, sample=False)
            verdict = read_verdict(self.last_reply)
            scanned.append((item.label, verdict))
            if verdict is None:
                # a reply that is neither yes nor no means the mask is not being
                # applied; refuse rather than fall through to a guess
                raise BrainError(
                    f"the masked reply to a yes/no question was {self.last_reply[:40]!r}"
                )
            if verdict == "yes":
                return Action("click", label=item.label,
                              reason=f"the model said yes to {item.label!r} "
                                     f"({len(scanned)} of {len(menu)} scanned)")
        self.last_reply = ""
        raise BrainError(
            f"the model said no to all {len(menu)} item(s): "
            f"{[label for label, _ in scanned][:6]}"
        )

    def _decide_by_ranking(self, goal: str, menu: DecisionMenu,
                           observation: Observation) -> Action:
        """Ask every item, then click the one the model liked most.

        ``scan`` stops at the first yes, which quietly rewards position. Here the
        model is asked about every candidate and the winner is the largest
        yes-minus-no log-odds margin - the model's own opinion, read off the
        distribution instead of off a sampled token. Costs N short generations,
        and it is the cheapest way to find out whether the weights can do it at all.
        """
        if not menu:
            raise GoalUnreachable("no clickable items to choose from")
        tokenizer = self._tokenizer_for_choice()
        yes_ids, no_ids = pair_token_ids(tokenizer)
        if not yes_ids or not no_ids:
            raise BrainError(
                "this tokenizer cannot express yes/no as single tokens, so the ranking "
                "cannot be constrained; use decision_mode='json'"
            )
        scored: list[tuple[Any, float]] = []
        for item in menu:
            self.last_prompt = self._scan_prompt(goal, item.label)
            margin = self._score_pair(self.last_prompt, observation, yes_ids, no_ids)
            scored.append((item, margin))
        scored.sort(key=lambda pair: pair[1], reverse=True)
        best, margin = scored[0]
        self.last_reply = f"margin {margin:+.2f}"
        if margin <= 0:
            raise BrainError(
                f"the model preferred no item: best was {best.label!r} at {margin:+.2f}; "
                f"runner-up {scored[1][0].label!r} at {scored[1][1]:+.2f}"
                if len(scored) > 1 else
                f"the model preferred no item (only {best.label!r}, {margin:+.2f})"
            )
        return Action("click", label=best.label,
                      reason=f"highest yes/no margin {margin:+.2f} of {len(scored)} item(s)")

    def _score_pair(self, prompt: str, observation: Observation,
                    yes_ids: Sequence[int], no_ids: Sequence[int]) -> float:
        """How much more likely "yes" is than "no". Backends implement this."""
        raise NotImplementedError("a constrained backend implements _score_pair")

    def _tokenizer_for_choice(self) -> Any:
        """The tokenizer used to build the mask. Subclasses load it lazily."""
        raise NotImplementedError("a constrained backend provides its tokenizer")

    def _generate_constrained(self, prompt: str, observation: Observation,
                              allowed_ids: Sequence[int], max_new_tokens: int,
                              *, sample: bool = False) -> str:
        raise NotImplementedError("a constrained backend implements this")

    @property
    def loaded(self) -> bool:
        """True once the weights are in memory - they load lazily, on the first decide()."""
        return getattr(self, "_model", None) is not None

    def _generate(self, prompt: str, observation: Observation) -> str:
        raise NotImplementedError(
            "subclass ModelBrain and implement _generate(prompt, observation)"
        )