"""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 (````) 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": , "reason": ""}', ' {"kind": "done", "reason": ""}', "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)" )