"""Constrained decoding: make an echo impossible instead of merely detected.""" import pytest from tiny_computer_agent import ( Action, BrainError, DecisionMenu, Element, GoalUnreachable, ModelBrain, Observation, ) from tiny_computer_agent.decoding import ( MenuMask, brevity_ladder, choice_token_ids, digit_token_ids, majority_choice, read_choice, ) FRAME = Observation(elements=( Element("Files", x=10, y=6, width=60, height=20), Element("Save", x=200, y=270, width=90, height=30), Element("Cancel", x=400, y=270, width=90, height=30), ), text="Untitled document - unsaved changes") class FakeTokenizer: """A tokenizer where each digit is one token (bare and space-prefixed).""" def __init__(self, expressible=True): self.expressible = expressible self.table = {} if expressible: for digit in range(10): self.table[str(digit)] = 100 + digit self.table[" " + str(digit)] = 200 + digit for index, word in enumerate(("yes", "Yes", "no", "No")): self.table[word] = 300 + index self.table[" " + word] = 320 + index def encode(self, text, add_special_tokens=False): if text in self.table: return [self.table[text]] return [999, 999] # never a single token # -- token ids --------------------------------------------------------------- def test_digit_ids_cover_every_digit_exactly_once(): ids = digit_token_ids(FakeTokenizer()) assert len(ids) == 20 # 0-9, bare and space-prefixed assert set(range(100, 110)) <= set(ids) def test_choice_ids_cover_only_the_legal_answers(): assert sorted(choice_token_ids(FakeTokenizer(), 3)) == sorted([100, 101, 102, 103, 200, 201, 202, 203]) assert sorted(choice_token_ids(FakeTokenizer(), 2, allow_done=False)) == sorted([101, 102, 201, 202]) def test_a_tokenizer_that_cannot_spell_a_number_yields_no_ids(): assert digit_token_ids(FakeTokenizer(expressible=False)) == [] assert choice_token_ids(FakeTokenizer(expressible=False), 5) == [] # -- the mask ---------------------------------------------------------------- def test_the_mask_forbids_everything_that_is_not_a_menu_number(): torch = pytest.importorskip("torch") scores = torch.zeros(1, 8) masked = MenuMask([1, 3])(None, scores) assert masked[0, 1].item() == 0 and masked[0, 3].item() == 0 for blocked in (0, 2, 4, 5, 6, 7): assert masked[0, blocked].item() == float("-inf") assert MenuMask([1, 3]).allowed_ids == (1, 3) def test_the_mask_keeps_the_end_of_sequence_reachable_after_one_token(): torch = pytest.importorskip("torch") process = MenuMask([1], eos_token_id=7) first = process(None, torch.zeros(1, 8)) assert first[0, 7].item() == float("-inf"), "stopping immediately is not an answer" assert first[0, 1].item() == 0 second = process(None, torch.zeros(1, 8)) assert second[0, 7].item() == 0 # ...but it can stop after answering assert second[0, 2].item() == float("-inf") def test_the_mask_can_be_told_not_to_force_a_token(): torch = pytest.importorskip("torch") process = MenuMask([1], eos_token_id=7, force_first=False) assert process(None, torch.zeros(1, 8))[0, 7].item() == 0 # -- reading an answer ------------------------------------------------------- def test_read_choice_reads_only_the_digits(): assert read_choice("3", 12) == 3 assert read_choice(" 3 ", 12) == 3 assert read_choice("3.", 12) == 3 assert read_choice("12", 12) == 12 assert read_choice("0", 12) == 0 # 0 is "already done" assert read_choice("I would click item 2", 12) == 2 def test_read_choice_refuses_what_is_not_exactly_one_answer(): assert read_choice("", 12) is None assert read_choice("no idea", 12) is None assert read_choice("13", 12) is None # outside the menu assert read_choice("2 or 3", 12) is None # 23: ambiguous, so refused assert read_choice("0", 12, allow_done=False) is None # -- the ladder and the vote ------------------------------------------------- def test_the_brevity_ladder_grows_then_stops(): assert brevity_ladder(3) == (0, 1, 2) assert brevity_ladder(1) == (0,) assert brevity_ladder(0) == (0,) def test_majority_choice_ignores_failures_and_breaks_ties_by_first_seen(): assert majority_choice([1, 1, 2]) == 1 assert majority_choice([None, 3, None]) == 3 assert majority_choice([2, 1]) == 2 # a tie goes to the first seen assert majority_choice([]) is None assert majority_choice([None, None]) is None # -- the brain, in constrained mode ------------------------------------------ class Scripted(ModelBrain): """A backend whose replies are scripted, so the ladder can be tested.""" name = "scripted" supports_constrained = True def __init__(self, replies, tokenizer=None, **kwargs): super().__init__("fake/model", **kwargs) self._replies = list(replies) self._tokenizer = tokenizer or FakeTokenizer() self.prompts = [] def _tokenizer_for_choice(self): return self._tokenizer def _generate_constrained(self, prompt, observation, allowed_ids, max_new_tokens, *, sample=False): self.prompts.append(prompt) assert allowed_ids, "the mask must be non-empty" return self._replies.pop(0) def _generate(self, prompt, observation): self.prompts.append(prompt) return self._replies.pop(0) def test_choice_mode_turns_a_number_into_a_click(): brain = Scripted(["2"], decision_mode="choice") action = brain.decide("save the document", FRAME) assert action.kind == "click" and action.label == "Save" assert "constrained choice 2" in action.reason assert brain.last_choice == 2 def test_choice_mode_reads_zero_as_already_done(): action = Scripted(["0"], decision_mode="choice").decide("save", FRAME) assert action.kind == "done" def test_the_echoed_prompt_is_not_usable_even_if_the_model_returns_it(): """The old failure mode, fed straight in: it cannot be mistaken for an answer.""" brain = Scripted(["2"], decision_mode="choice") echo = brain.build_prompt("save", DecisionMenu(FRAME), FRAME, brevity=0) assert read_choice(echo, 12) is None, "an echoed menu must not parse as a choice" # the prompt contains the menu numbers, spaced apart: 1, 2, 3 -> "123" assert read_choice(brain.build_prompt("save", DecisionMenu(FRAME), FRAME), 3) is None def test_a_useless_reply_shortens_the_prompt_and_tries_again(): brain = Scripted(["", "3"], decision_mode="choice", retries=3) action = brain.decide("discard the change", FRAME) assert action.label == "Cancel" assert len(brain.prompts) == 2 assert len(brain.prompts[1]) < len(brain.prompts[0]), "the retry must be shorter" def test_the_prompt_shrinks_all_the_way_when_nothing_works(): brain = Scripted([""] * 5, decision_mode="choice", retries=3) with pytest.raises(BrainError, match="not a menu number"): brain.decide("save", FRAME) lengths = [len(prompt) for prompt in brain.prompts] assert lengths[0] > lengths[-1] and len(lengths) == 3 assert "Text on screen" not in brain.prompts[-1] def test_votes_take_a_majority_over_samples(): brain = Scripted(["2", "3", "2"], decision_mode="choice", votes=3) action = brain.decide("save", FRAME) assert action.label == "Save" # 2 wins over 3 assert len(brain.prompts) == 3 # only one ladder rung used def test_a_tokenizer_that_cannot_spell_numbers_is_refused_clearly(): brain = Scripted(["1"], tokenizer=FakeTokenizer(expressible=False), decision_mode="choice") with pytest.raises(BrainError, match="cannot express a menu number"): brain.decide("save", FRAME) def test_an_empty_screen_is_a_hard_stop_in_choice_mode(): brain = Scripted(["1"], decision_mode="choice") with pytest.raises(GoalUnreachable): brain.decide("save", Observation(elements=())) class ScanBackend(ModelBrain): """Answers yes only for the labels it was told to like.""" name = "scan-backend" supports_constrained = True def __init__(self, likes, tokenizer=None, **kwargs): super().__init__("fake/model", **kwargs) self.likes = set(likes) self.tokenizer = tokenizer or FakeTokenizer() self.prompts = [] def _tokenizer_for_choice(self): return self.tokenizer def _generate_constrained(self, prompt, observation, allowed_ids, max_new_tokens, *, sample=False): self.prompts.append(prompt) label = prompt.split("Candidate item: ", 1)[1].splitlines()[0] return "yes" if label in self.likes else "no" def test_scanning_asks_one_question_per_item_and_takes_the_first_yes(): brain = ScanBackend(["Save"], decision_mode="scan") action = brain.decide("click 'Save' to save the document", FRAME) assert action.kind == "click" and action.label == "Save" assert "said yes to 'Save'" in action.reason assert len(brain.prompts) == 2 # Files, Save - then it stopped assert all("yes or no" in prompt for prompt in brain.prompts) assert all("Candidate item:" in prompt for prompt in brain.prompts) def test_scanning_refuses_when_every_item_is_rejected(): brain = ScanBackend([], decision_mode="scan") with pytest.raises(BrainError, match="said no to all"): brain.decide("save", FRAME) assert len(brain.prompts) == 3 # every item was asked about def test_scanning_refuses_a_reply_that_is_neither_yes_nor_no(): class Mute(ScanBackend): def _generate_constrained(self, prompt, observation, allowed_ids, max_new_tokens, *, sample=False): return "maybe" with pytest.raises(BrainError, match="neither yes nor no|masked reply"): Mute(["Save"], decision_mode="scan").decide("save", FRAME) def test_read_verdict_reads_what_the_mask_allows(): from tiny_computer_agent.decoding import pair_token_ids, read_verdict assert read_verdict("yes") == "yes" assert read_verdict(" Yes ") == "yes" assert read_verdict("no") == "no" assert read_verdict("No.") == "no" assert read_verdict("perhaps") is None assert read_verdict("") is None yes_ids, no_ids = pair_token_ids(FakeTokenizer()) assert yes_ids and no_ids and not (set(yes_ids) & set(no_ids)) def test_json_mode_is_untouched(): """The default must not change: the existing lenient parsing still applies.""" brain = Scripted(['{"kind": "click", "index": 2}']) assert brain.decision_mode == "json" assert brain.decide("save", FRAME).label == "Save" def test_choice_mode_is_only_entered_when_the_backend_supports_it(): class Plain(Scripted): supports_constrained = False brain = Plain(['{"kind": "click", "index": 2}'], decision_mode="choice") assert brain.decide("save", FRAME).label == "Save" # fell back to JSON parsing