diff --git a/docs/RL_FRAMEWORK.md b/docs/RL_FRAMEWORK.md index 9094ff2..caf3fed 100644 --- a/docs/RL_FRAMEWORK.md +++ b/docs/RL_FRAMEWORK.md @@ -62,6 +62,24 @@ raises on model output and separates format validity from constraint validity; the evaluator and the reward adapter both consume it so there is exactly one contract implementation. +Phase 6 introduced the prompt-facing choice protocol +(PromptChoiceRegistry, global Stage 1 ids and local Stage 2 ids), keeping +the canonical category_id out of prompts and model outputs. + +Phase 7 adds the shared choice-aware parser/decode layer in the same parser +module (`check_stage1_choices` / `check_stage2_choices` -> +`ChoiceParseResult`): model rollouts answer with choice ids (Stage 1 global +ids, Stage 2 local bundle ids), which this layer validates (JSON/schema, +count, uniqueness, known ids / 1..5) and decodes to canonical category_ids +BEFORE the existing evaluation (`evaluate_stage1_choices` / +`evaluate_stage2_choices`) and reward entry points +(`reward_stage1_choices` / `reward_stage2_choices`) apply, completing the +task-level choice-to-canonical reward contract. Choice validation is +implemented exactly once; `RewardConfig`, reward numbers, canonical +registry/ground truth/candidates are unchanged. `RewardResult` gained a +`constraint_valid` flag (additive) so "valid but wrong" (partial credit) is +distinguishable from "invalid" (zero). + ### Training adapters `agent.training.sft` owns messages-Parquet export, validation, the temporary @@ -117,6 +135,12 @@ Still landing with the first real RL vertical slice (M4): reward loop; 2. a small GPU smoke launcher. +The task-level choice-to-canonical contract is now complete on master: Phase +6 (prompt-facing choice protocol, PromptChoiceRegistry with global Stage 1 +ids and local Stage 2 ids) and Phase 7 (shared choice-aware parser/decode +layer consumed by evaluation and reward; see the Evaluation module) are both +merged. + ## When to add training-algorithm RL code Do not create empty GRPO/PPO/DAPO packages or `cfg/verl/rl` placeholders. diff --git a/src/agent/evaluation/classification.py b/src/agent/evaluation/classification.py index 3204c9a..84904c3 100644 --- a/src/agent/evaluation/classification.py +++ b/src/agent/evaluation/classification.py @@ -10,30 +10,26 @@ evaluate_stage1/2 consume the unified parser (check_stage1_output / check_stage2_output) so the reward adapter and the evaluator share one contract implementation. evaluate_stage1_choices / evaluate_stage2_choices -are thin adapters: they parse and decode the choice protocol BEFORE -delegating to the canonical evaluators, so choice ids never leak into -correctness logic and no canonical check logic is duplicated. +consume the SHARED choice-aware layer (check_stage1_choices / +check_stage2_choices) and apply the same canonical correctness facts on +the decoded category ids, so choice validation is implemented exactly +once across evaluation and reward. """ from __future__ import annotations from dataclasses import dataclass -import json from typing import Sequence from agent.task import LeafRegistry from agent.task.parser import ( - PredictionFormatError, + Stage2Output, + check_stage1_choices, check_stage1_output, + check_stage2_choices, check_stage2_output, - parse_stage1_output, - parse_stage2_output, -) -from agent.task.prompt_choices import ( - PromptChoiceError, - PromptChoiceRegistry, - decode_stage2_answer, ) +from agent.task.prompt_choices import PromptChoiceRegistry @dataclass(frozen=True) @@ -87,28 +83,23 @@ def evaluate_stage1_choices( ) -> Stage1Evaluation: """Evaluate a choice-id Stage 1 output; decode BEFORE canonical logic. - The model answers with global choice ids ("1".."N"); the returned - prediction is the decoded canonical category_id tuple. Invalid choice - ids / wrong counts / duplicates yield an explicit invalid result with - no name or fuzzy fallback. Delegates to evaluate_stage1 with the - decoded canonical payload, so the canonical contract logic is never - duplicated. + Consumes the shared choice-aware check (check_stage1_choices) so + evaluation and reward share one choice validation implementation; the + returned prediction is the decoded canonical category_id tuple. Invalid + choice ids / wrong counts / duplicates yield an explicit invalid result + with no name or fuzzy fallback. """ if ground_truth not in registry.ids: raise ValueError("ground_truth must belong to the leaf registry") choices = choices or PromptChoiceRegistry.from_registry(registry) - try: - output = parse_stage1_output(solution) - except PredictionFormatError as exc: - return Stage1Evaluation(None, False, False, False, (str(exc),)) - try: - decoded = choices.decode_candidates(output.candidates) - except PromptChoiceError as exc: - return Stage1Evaluation(None, True, False, False, (str(exc),)) - return evaluate_stage1( - json.dumps({"candidates": list(decoded)}, ensure_ascii=False, separators=(",", ":")), - ground_truth=ground_truth, - registry=registry, + result = check_stage1_choices(solution, choices=choices) + if not result.format_valid: + return Stage1Evaluation(None, False, False, False, result.errors) + if not result.constraint_valid: + return Stage1Evaluation(None, True, False, False, result.errors) + assert result.decoded is not None + return Stage1Evaluation( + result.decoded, True, True, ground_truth in result.decoded, () ) @@ -150,10 +141,11 @@ def evaluate_stage2_choices( """Evaluate a local-id Stage 2 output; decode BEFORE canonical logic. The model answers with a LOCAL bundle id ("1".."5" in candidate order); - the returned prediction is the decoded canonical category_id. Anything - but an exact local id yields an explicit invalid result with no name or - fuzzy fallback. Delegates to evaluate_stage2 with the decoded canonical - payload, so the canonical contract logic is never duplicated. + the returned prediction is the decoded canonical category_id. Consumes + the shared choice-aware check (check_stage2_choices) so evaluation and + reward share one choice validation implementation. Anything but an + exact local id yields an explicit invalid result with no name or fuzzy + fallback. """ if ground_truth not in registry.ids: raise ValueError("ground_truth must belong to the leaf registry") @@ -163,17 +155,15 @@ def evaluate_stage2_choices( or any(candidate not in registry.ids for candidate in candidates) ): raise ValueError("candidates must be 5 unique IDs from the leaf registry") - try: - output = parse_stage2_output(solution) - except PredictionFormatError as exc: - return Stage2Evaluation(None, False, False, False, (str(exc),)) - try: - decoded = decode_stage2_answer(output.answer, tuple(candidates)) - except PromptChoiceError as exc: - return Stage2Evaluation(output.answer, True, False, False, (str(exc),)) - return evaluate_stage2( - json.dumps({"answer": decoded}, ensure_ascii=False, separators=(",", ":")), - ground_truth=ground_truth, - candidates=candidates, - registry=registry, + result = check_stage2_choices(solution, candidates=candidates) + if not result.format_valid: + return Stage2Evaluation(None, False, False, False, result.errors) + if not result.constraint_valid: + assert isinstance(result.output, Stage2Output) + return Stage2Evaluation( + result.output.answer, True, False, False, result.errors + ) + assert result.decoded is not None + return Stage2Evaluation( + result.decoded, True, True, result.decoded == ground_truth, () ) diff --git a/src/agent/task/__init__.py b/src/agent/task/__init__.py index 4df4299..fba2b96 100644 --- a/src/agent/task/__init__.py +++ b/src/agent/task/__init__.py @@ -10,11 +10,14 @@ from .dataset_config import BUILTIN_DATASET_CONFIGS, DatasetConfig from .identity import code_leaf_map, compact, leaf_registry_from_corpus, qualified_category_id from .parser import ( + ChoiceParseResult, ParseResult, PredictionFormatError, Stage1Output, Stage2Output, + check_stage1_choices, check_stage1_output, + check_stage2_choices, check_stage2_output, parse_stage1_output, parse_stage2_output, @@ -62,8 +65,11 @@ "Stage1Output", "Stage2Output", "ParseResult", + "ChoiceParseResult", "check_stage1_output", "check_stage2_output", + "check_stage1_choices", + "check_stage2_choices", "parse_stage1_output", "parse_stage2_output", "PromptChoice", diff --git a/src/agent/task/parser.py b/src/agent/task/parser.py index 0ffe524..3a21c7e 100644 --- a/src/agent/task/parser.py +++ b/src/agent/task/parser.py @@ -1,6 +1,6 @@ """Strict parsers for model outputs in the two-stage classification task. -Two layers, sharing the exact JSON-shape parsers: +Three layers, sharing the exact JSON-shape parsers: - ``parse_stage1_output`` / ``parse_stage2_output`` raise ``PredictionFormatError`` on malformed JSON/schema (used by SFT @@ -11,6 +11,14 @@ uniqueness / registry membership, or answer membership). RL rewards and the evaluator consume this layer so malformed model output can never crash a training loop. +- ``check_stage1_choices`` / ``check_stage2_choices`` are the SHARED + choice-aware layer: they run the same JSON/schema parsing, validate the + choice protocol (stage1: exact count, uniqueness, known choice ids; + stage2: local id 1..5) and decode to canonical category ids, returning a + ``ChoiceParseResult``. Both the evaluation adapters and the RL reward + entry points consume this layer, so choice validation is implemented + exactly once and the model's choice ids never leak into canonical + correctness logic. """ from __future__ import annotations @@ -20,6 +28,7 @@ from typing import Any, Sequence from .contracts import LeafRegistry +from .prompt_choices import PromptChoiceRegistry class PredictionFormatError(ValueError): @@ -167,3 +176,150 @@ def check_stage2_output(text: str, *, candidates: Sequence[str]) -> ParseResult: valid = output.answer in candidate_ids errors = () if valid else ("stage2 answer must be one of the candidates",) return ParseResult(True, valid, output, errors) + + +_STAGE2_LOCAL_IDS = tuple(str(index) for index in range(1, 6)) + + +@dataclass(frozen=True) +class ChoiceParseResult: + """Structured never-raised outcome of the choice-aware parse + decode. + + The choice protocol (``agent.task.prompt_choices``) maps model-facing + choice ids to canonical category ids. This result mirrors + ``ParseResult``'s shape contract and additionally carries ``decoded``: + the canonical category ids the model's choice ids map to (stage1: the + candidate tuple; stage2: the single answer id). ``output`` always keeps + the model-level parsed shape (choice ids) when format_valid, so + consumers can see exactly what the model said before decoding. + + - format_valid: the output is a JSON object with exactly the required + schema (identical semantics to ParseResult.format_valid). + - constraint_valid: the parsed choice ids satisfy the choice protocol + (stage1: exactly ``expected_count`` unique known choice ids; stage2: + answer is one of the local ids 1..5). + - decoded: the canonical category ids when constraint_valid (stage1: + tuple of candidate ids; stage2: the single answer id), else None. + - errors: human-readable failure reasons; empty when both valid. + """ + + format_valid: bool + constraint_valid: bool + decoded: tuple[str, ...] | str | None = None + output: "Stage1Output | Stage2Output | None" = None + errors: tuple[str, ...] = () + + @property + def ok(self) -> bool: + return self.format_valid and self.constraint_valid + + def __post_init__(self) -> None: + if not isinstance(self.format_valid, bool) or not isinstance( + self.constraint_valid, bool + ): + raise ValueError("ChoiceParseResult validity flags must be bool") + if self.errors and self.format_valid and self.constraint_valid: + raise ValueError("a fully valid ChoiceParseResult must carry no errors") + if not isinstance(self.errors, tuple) or not all( + isinstance(error, str) for error in self.errors + ): + raise ValueError("ChoiceParseResult errors must be a tuple of strings") + if self.constraint_valid and self.decoded is None: + raise ValueError( + "a constraint-valid ChoiceParseResult must carry decoded canonical ids" + ) + + def canonical_view(self) -> ParseResult: + """Equivalent canonical ``ParseResult`` for the shared eval/reward cores. + + When constraints hold the output is the decoded canonical shape + (Stage1Output of decoded ids / Stage2Output of the decoded answer); + when only the format held, the model-level shape is kept so + consumers can still inspect what the model said. + """ + if not self.format_valid: + return ParseResult(False, False, None, self.errors) + if self.constraint_valid: + output: Stage1Output | Stage2Output | None = ( + Stage1Output(self.decoded) + if isinstance(self.decoded, tuple) + else Stage2Output(self.decoded) + ) + else: + output = self.output + return ParseResult(True, self.constraint_valid, output, self.errors) + + +def check_stage1_choices( + text: str, + *, + choices: PromptChoiceRegistry, + expected_count: int = 5, +) -> ChoiceParseResult: + """Shared choice-aware Stage 1 parser + decode; never raises on model output. + + Validates, in order: JSON object, exact schema (only ``candidates``), + candidate count, uniqueness, and choice-id membership in the prompt + catalog; then decodes the choice ids to canonical category ids. + Malformed model output maps to a structured ``ChoiceParseResult``; only + a missing/invalid ``choices`` mapping is a programming error and + raises. + """ + if expected_count < 1: + raise ValueError("expected_count must be positive") + if not isinstance(choices, PromptChoiceRegistry): + raise ValueError("choices must be a PromptChoiceRegistry") + try: + output = parse_stage1_output(text) + except PredictionFormatError as exc: + return ChoiceParseResult(False, False, None, None, (str(exc),)) + errors: list[str] = [] + if len(output.candidates) != expected_count: + errors.append( + f"stage1 prediction must contain exactly {expected_count} candidates" + ) + if len(set(output.candidates)) != len(output.candidates): + errors.append("stage1 candidates must be unique") + for choice_id in output.candidates: + if not choices.contains_choice_id(choice_id): + errors.append(f"choice id {choice_id!r} is not in the prompt catalog") + if errors: + return ChoiceParseResult(True, False, None, output, tuple(errors)) + return ChoiceParseResult( + True, + True, + tuple(choices.category_id_of(choice_id) for choice_id in output.candidates), + output, + (), + ) + + +def check_stage2_choices(text: str, *, candidates: Sequence[str]) -> ChoiceParseResult: + """Shared choice-aware Stage 2 parser + decode; never raises on model output. + + Validates: JSON object, exact schema (only ``answer``), and answer + membership in the local ids 1..5; then decodes the local id to the + canonical category id of the candidate at that position. Candidates + originate from the dataset, so an invalid candidate list is a + programming error and raises (mirrors check_stage2_output). + """ + if isinstance(candidates, (str, bytes)) or len(candidates) != 5: + raise ValueError("stage2 requires exactly 5 candidates for local-id decode") + candidate_ids = tuple(candidates) + if not all(isinstance(candidate, str) and candidate for candidate in candidate_ids): + raise ValueError("stage2 candidates must be non-empty strings") + try: + output = parse_stage2_output(text) + except PredictionFormatError as exc: + return ChoiceParseResult(False, False, None, None, (str(exc),)) + if output.answer not in _STAGE2_LOCAL_IDS: + return ChoiceParseResult( + True, + False, + None, + output, + (f"stage2 answer {output.answer!r} must be one of 1..5",), + ) + return ChoiceParseResult( + True, True, candidate_ids[int(output.answer) - 1], output, () + ) diff --git a/src/agent/training/rl/__init__.py b/src/agent/training/rl/__init__.py index fd51852..3eb9988 100644 --- a/src/agent/training/rl/__init__.py +++ b/src/agent/training/rl/__init__.py @@ -20,9 +20,12 @@ STAGE2_PARTIAL_DEFAULT, RewardConfig, RewardResult, + reward_for_choice_result, reward_for_parse_result, reward_stage1, + reward_stage1_choices, reward_stage2, + reward_stage2_choices, ) from .sample import ( RewardMeta, @@ -49,6 +52,9 @@ "reward_stage1", "reward_stage2", "reward_for_parse_result", + "reward_stage1_choices", + "reward_stage2_choices", + "reward_for_choice_result", "RL_SPLITS", "VERL_RL_COLUMNS", "export_rl_dataset", diff --git a/src/agent/training/rl/reward.py b/src/agent/training/rl/reward.py index 9b642d5..cced1ff 100644 --- a/src/agent/training/rl/reward.py +++ b/src/agent/training/rl/reward.py @@ -24,6 +24,14 @@ candidates — a Stage 1 recall failure (ground_truth not in candidates) is a legal task state, not a Stage 2 contract error. A valid candidate answer still computes a reward without raising. + +Choice protocol: reward_stage1_choices / reward_stage2_choices / +reward_for_choice_result consume the SHARED choice-aware parser +(check_stage1_choices / check_stage2_choices), which decodes the model's +choice ids to canonical category ids BEFORE the SAME reward table below +applies. Choice validation, partial-credit values and malformed-output +handling are therefore exactly the canonical ones; RewardConfig and the +reward numbers are unchanged. """ from __future__ import annotations @@ -32,12 +40,16 @@ from agent.task.contracts import LeafRegistry from agent.task.parser import ( + ChoiceParseResult, ParseResult, Stage1Output, Stage2Output, + check_stage1_choices, check_stage1_output, + check_stage2_choices, check_stage2_output, ) +from agent.task.prompt_choices import PromptChoiceRegistry FULL_REWARD = 1.0 INVALID_REWARD = 0.0 @@ -75,8 +87,13 @@ class RewardResult: - reward: the scalar task reward. - reason: human-readable justification (safe to log). - - parsed_output: the parsed model output when format-valid, else None. + - parsed_output: the parsed model output when format-valid (for choice + outputs: the decoded canonical shape when decodable, else the raw + model-level shape). - format_valid: whether the output was valid JSON with the exact schema. + - constraint_valid: whether the output satisfied the task constraints + (Stage 1 candidate rules / Stage 2 answer membership) — distinguishes + "valid but wrong" (partial credit) from "invalid" (zero). - task_correct: whether the output achieved the task goal (Stage 1: ground truth recalled; Stage 2: correct answer). """ @@ -85,11 +102,12 @@ class RewardResult: reason: str parsed_output: Stage1Output | Stage2Output | None format_valid: bool + constraint_valid: bool task_correct: bool def _invalid(reason: str) -> RewardResult: - return RewardResult(INVALID_REWARD, reason, None, False, False) + return RewardResult(INVALID_REWARD, reason, None, False, False, False) def reward_for_parse_result( @@ -115,15 +133,17 @@ def reward_for_parse_result( result.output, True, False, + False, ) assert isinstance(result.output, Stage1Output) if ground_truth in result.output.candidates: - return RewardResult(FULL_REWARD, "stage1 candidates contain the ground truth", result.output, True, True) + return RewardResult(FULL_REWARD, "stage1 candidates contain the ground truth", result.output, True, True, True) return RewardResult( config.stage1_valid_miss, "stage1 output is valid but misses the ground truth", result.output, True, + True, False, ) if stage == "stage2": @@ -136,15 +156,17 @@ def reward_for_parse_result( result.output, True, False, + False, ) assert isinstance(result.output, Stage2Output) if result.output.answer == ground_truth: - return RewardResult(FULL_REWARD, "stage2 answer is correct", result.output, True, True) + return RewardResult(FULL_REWARD, "stage2 answer is correct", result.output, True, True, True) return RewardResult( config.stage2_partial, "stage2 answer is valid but wrong", result.output, True, + True, False, ) raise ValueError(f"reward stage must be stage1 or stage2, got {stage!r}") @@ -193,6 +215,80 @@ def reward_stage2( return reward_for_parse_result("stage2", result, ground_truth=ground_truth, config=config) +def reward_for_choice_result( + stage: str, + result: ChoiceParseResult, + *, + ground_truth: str, + config: RewardConfig | None = None, +) -> RewardResult: + """Map a shared choice-aware ``ChoiceParseResult`` to a task reward. + + Single implementation of the reward tables: the decoded canonical view + feeds ``reward_for_parse_result``, so choice validation, partial-credit + values and malformed-output handling are exactly the canonical ones. + Never raises on model output. + """ + return reward_for_parse_result( + stage, result.canonical_view(), ground_truth=ground_truth, config=config + ) + + +def reward_stage1_choices( + solution: str, + *, + ground_truth: str, + registry: LeafRegistry, + choices: PromptChoiceRegistry | None = None, + config: RewardConfig | None = None, +) -> RewardResult: + """Stage 1 task reward over the choice protocol (global choice ids). + + The model answers with choice ids; the shared choice-aware parser + (check_stage1_choices) decodes them to canonical category ids BEFORE + the unchanged reward table applies. Never raises on model output; + ``ground_truth`` outside the registry is a programming error and + raises (mirrors reward_stage1). + """ + if ground_truth not in registry.ids: + raise ValueError("ground_truth must belong to the leaf registry") + choices = choices or PromptChoiceRegistry.from_registry(registry) + result = check_stage1_choices(solution, choices=choices) + return reward_for_choice_result( + "stage1", result, ground_truth=ground_truth, config=config + ) + + +def reward_stage2_choices( + solution: str, + *, + ground_truth: str, + candidates: tuple[str, ...] | list[str], + registry: LeafRegistry, + config: RewardConfig | None = None, +) -> RewardResult: + """Stage 2 task reward over the choice protocol (local bundle id 1..5). + + The model answers with a local id; the shared choice-aware parser + (check_stage2_choices) decodes it against the candidate bundle BEFORE + the unchanged reward table applies. Never raises on model output; + ``candidates`` must be 5 unique registry IDs (programming contract, + mirrors reward_stage2). + """ + if ground_truth not in registry.ids: + raise ValueError("ground_truth must belong to the leaf registry") + if isinstance(candidates, (str, bytes)) or ( + len(candidates) != 5 + or len(set(candidates)) != 5 + or any(candidate not in registry.ids for candidate in candidates) + ): + raise ValueError("candidates must be 5 unique IDs from the leaf registry") + result = check_stage2_choices(solution, candidates=candidates) + return reward_for_choice_result( + "stage2", result, ground_truth=ground_truth, config=config + ) + + __all__ = [ "FULL_REWARD", "INVALID_REWARD", @@ -203,4 +299,7 @@ def reward_stage2( "reward_stage1", "reward_stage2", "reward_for_parse_result", + "reward_stage1_choices", + "reward_stage2_choices", + "reward_for_choice_result", ] diff --git a/tests/rl/test_rl_reward_choices.py b/tests/rl/test_rl_reward_choices.py new file mode 100644 index 0000000..ab65956 --- /dev/null +++ b/tests/rl/test_rl_reward_choices.py @@ -0,0 +1,233 @@ +"""Choice-protocol reward (Phase 7): reward_stage1/2_choices. + +RL rollouts answer with choice ids; the shared choice-aware parser decodes +them to canonical category ids before the unchanged reward table applies. +These tests pin the reward behavior (correct / valid-wrong / invalid / +malformed) and the agreement between evaluation and reward on the validity +of the same output. +""" + +from __future__ import annotations + +import pytest + +from agent.evaluation import evaluate_stage1_choices, evaluate_stage2_choices +from agent.task import LeafRegistry, PromptChoiceRegistry +from agent.training.rl import ( + FULL_REWARD, + INVALID_REWARD, + STAGE1_VALID_MISS_DEFAULT, + STAGE2_PARTIAL_DEFAULT, + RewardConfig, + RewardResult, + reward_stage1_choices, + reward_stage2_choices, +) + + +def _registry() -> LeafRegistry: + return LeafRegistry.from_mapping(["A", "B", "C", "D", "E", "F"]) + + +def _choices() -> PromptChoiceRegistry: + return PromptChoiceRegistry.from_registry(_registry()) + + +# local ids in candidate order: C=1 A=2 B=3 D=4 E=5 (ground truth = C) +CANDIDATES = ("C", "A", "B", "D", "E") + + +def test_stage1_choice_correct_scores_full() -> None: + # choice ids: A=1 B=2 C=3 D=4 E=5 F=6 -> candidates C A B D E recall C + result = reward_stage1_choices( + '{"candidates":["3","1","2","4","5"]}', + ground_truth="C", + registry=_registry(), + choices=_choices(), + ) + assert result.reward == pytest.approx(FULL_REWARD) + assert result.format_valid is True + assert result.constraint_valid is True + assert result.task_correct is True + + +def test_stage1_choice_valid_but_miss_scores_valid_miss() -> None: + # A B D E F do not include C: valid output, partial credit + result = reward_stage1_choices( + '{"candidates":["1","2","4","5","6"]}', + ground_truth="C", + registry=_registry(), + choices=_choices(), + ) + assert result.reward == pytest.approx(STAGE1_VALID_MISS_DEFAULT) + assert result.format_valid is True + assert result.constraint_valid is True + assert result.task_correct is False + + +def test_stage1_choice_invalid_scores_zero() -> None: + for text in ( + '{"candidates":["3","1","2","4","9"]}', # unknown choice id + '{"candidates":["1","1","2","4","5"]}', # duplicate + '{"candidates":["1","2","3"]}', # wrong count + ): + result = reward_stage1_choices( + text, ground_truth="C", registry=_registry(), choices=_choices() + ) + assert result.reward == pytest.approx(INVALID_REWARD), text + assert result.format_valid is True, text + assert result.constraint_valid is False, text + assert result.task_correct is False + + +def test_stage2_choice_correct_scores_full() -> None: + result = reward_stage2_choices( + '{"answer":"1"}', ground_truth="C", candidates=CANDIDATES, registry=_registry() + ) + assert result.reward == pytest.approx(FULL_REWARD) + assert result.format_valid is True + assert result.constraint_valid is True + assert result.task_correct is True + + +def test_stage2_choice_valid_but_wrong_scores_partial() -> None: + result = reward_stage2_choices( + '{"answer":"2"}', ground_truth="C", candidates=CANDIDATES, registry=_registry() + ) + assert result.reward == pytest.approx(STAGE2_PARTIAL_DEFAULT) + assert result.format_valid is True + assert result.constraint_valid is True + assert result.task_correct is False + + +def test_stage2_choice_invalid_answer_scores_zero() -> None: + for answer in ("0", "6", "A"): + result = reward_stage2_choices( + f'{{"answer":"{answer}"}}', + ground_truth="C", + candidates=CANDIDATES, + registry=_registry(), + ) + assert result.reward == pytest.approx(INVALID_REWARD), answer + assert result.format_valid is True, answer + assert result.constraint_valid is False, answer + assert result.task_correct is False + + +def test_malformed_json_scores_zero_without_raising() -> None: + for text in ("", "{", "[]", "not json", '{"candidates": 3}', '{"answer": null}'): + result = reward_stage1_choices( + text, ground_truth="C", registry=_registry() + ) # default choices path + assert isinstance(result, RewardResult), text + assert result.reward == pytest.approx(INVALID_REWARD), text + assert result.format_valid is False, text + result2 = reward_stage2_choices( + text, ground_truth="C", candidates=CANDIDATES, registry=_registry() + ) + assert isinstance(result2, RewardResult), text + assert result2.reward == pytest.approx(INVALID_REWARD), text + assert result2.format_valid is False, text + + +def test_choice_reward_config_is_unchanged_and_applied() -> None: + config = RewardConfig(stage1_valid_miss=0.2, stage2_partial=0.25) + miss = reward_stage1_choices( + '{"candidates":["1","2","4","5","6"]}', + ground_truth="C", + registry=_registry(), + choices=_choices(), + config=config, + ) + wrong = reward_stage2_choices( + '{"answer":"2"}', + ground_truth="C", + candidates=CANDIDATES, + registry=_registry(), + config=config, + ) + assert miss.reward == pytest.approx(0.2) + assert wrong.reward == pytest.approx(0.25) + # defaults remain the contract values + defaults = RewardConfig() + assert defaults.stage1_valid_miss == pytest.approx(STAGE1_VALID_MISS_DEFAULT) + assert defaults.stage2_partial == pytest.approx(STAGE2_PARTIAL_DEFAULT) + + +def test_choice_reward_result_carries_contract_fields() -> None: + result = reward_stage1_choices( + '{"candidates":["3","1","2","4","5"]}', + ground_truth="C", + registry=_registry(), + choices=_choices(), + ) + for field in ( + "reward", + "reason", + "parsed_output", + "format_valid", + "constraint_valid", + "task_correct", + ): + assert hasattr(result, field) + + +def test_programming_errors_still_raise_for_choice_reward() -> None: + with pytest.raises(ValueError): + reward_stage1_choices( + '{"candidates":["1","2","3","4","5"]}', + ground_truth="Z", + registry=_registry(), + choices=_choices(), + ) + with pytest.raises(ValueError): + reward_stage2_choices( + '{"answer":"1"}', ground_truth="C", candidates=("C", "A"), registry=_registry() + ) + + +STAGE1_OUTPUTS = ( + '{"candidates":["3","1","2","4","5"]}', # valid, ground truth recalled + '{"candidates":["1","2","4","5","6"]}', # valid, ground truth missed + '{"candidates":["3","1","2","4","9"]}', # unknown choice id + '{"candidates":["1","1","2","4","5"]}', # duplicate + '{"candidates":["1","2","3"]}', # wrong count + "not json", + "{", + '{"candidates":["1","2",3,"4","5"]}', # non-string member +) + +STAGE2_OUTPUTS = ( + '{"answer":"1"}', # correct + '{"answer":"2"}', # valid wrong + '{"answer":"0"}', + '{"answer":"6"}', + '{"answer":"A"}', + "garbage", + "{", + '{"answer":5}', # non-string answer +) + + +@pytest.mark.parametrize("text", STAGE1_OUTPUTS) +def test_evaluation_and_reward_agree_on_stage1_validity(text: str) -> None: + evaluation = evaluate_stage1_choices( + text, ground_truth="C", registry=_registry(), choices=_choices() + ) + reward = reward_stage1_choices( + text, ground_truth="C", registry=_registry(), choices=_choices() + ) + assert evaluation.format_valid == reward.format_valid, text + assert evaluation.contract_valid == reward.constraint_valid, text + + +@pytest.mark.parametrize("text", STAGE2_OUTPUTS) +def test_evaluation_and_reward_agree_on_stage2_validity(text: str) -> None: + evaluation = evaluate_stage2_choices( + text, ground_truth="C", candidates=CANDIDATES, registry=_registry() + ) + reward = reward_stage2_choices( + text, ground_truth="C", candidates=CANDIDATES, registry=_registry() + ) + assert evaluation.format_valid == reward.format_valid, text + assert evaluation.contract_valid == reward.constraint_valid, text diff --git a/tests/task/test_parser_choices.py b/tests/task/test_parser_choices.py new file mode 100644 index 0000000..b6dff54 --- /dev/null +++ b/tests/task/test_parser_choices.py @@ -0,0 +1,179 @@ +"""Shared choice-aware parser layer (Phase 7): check_stage1/2_choices. + +Model outputs speak choice ids; this layer validates the choice protocol +and decodes to canonical category ids in ONE shared implementation that +both evaluation adapters and RL reward entry points consume. +""" + +from __future__ import annotations + +import pytest + +from agent.task import ( + ChoiceParseResult, + LeafRegistry, + ParseResult, + PromptChoiceRegistry, + check_stage1_choices, + check_stage2_choices, +) + + +def _registry() -> LeafRegistry: + return LeafRegistry.from_mapping(["A", "B", "C", "D", "E", "F"]) + + +def _choices() -> PromptChoiceRegistry: + return PromptChoiceRegistry.from_registry(_registry()) + + +# canonical candidates for stage2 fixtures; local ids: C=1 A=2 B=3 D=4 E=5 +CANDIDATES = ("C", "A", "B", "D", "E") + + +def test_stage1_valid_choice_decodes_to_canonical_ids() -> None: + # choice ids: A=1 B=2 C=3 D=4 E=5 F=6 + result = check_stage1_choices('{"candidates":["3","1","2","4","5"]}', choices=_choices()) + assert isinstance(result, ChoiceParseResult) + assert result.ok is True + assert result.format_valid is True + assert result.constraint_valid is True + assert result.decoded == ("C", "A", "B", "D", "E") + assert result.errors == () + # the model-level shape is preserved for inspection + assert result.output is not None + assert result.output.candidates == ("3", "1", "2", "4", "5") + + +def test_stage1_unknown_choice_id_fails_constraints() -> None: + result = check_stage1_choices('{"candidates":["3","1","2","4","9"]}', choices=_choices()) + assert result.format_valid is True + assert result.constraint_valid is False + assert result.ok is False + assert result.decoded is None + assert any("9" in error and "prompt catalog" in error for error in result.errors) + + +def test_stage1_duplicate_and_wrong_count_fail_constraints() -> None: + duplicate = check_stage1_choices( + '{"candidates":["1","1","2","4","5"]}', choices=_choices() + ) + short = check_stage1_choices('{"candidates":["1","2","3"]}', choices=_choices()) + long = check_stage1_choices( + '{"candidates":["1","2","3","4","5","6"]}', choices=_choices() + ) + assert duplicate.constraint_valid is False + assert any("unique" in error for error in duplicate.errors) + assert short.constraint_valid is False + assert long.constraint_valid is False + assert any("exactly 5" in error for error in short.errors) + assert any("exactly 5" in error for error in long.errors) + # expected_count is configurable, like check_stage1_output + two = check_stage1_choices('{"candidates":["1","2"]}', choices=_choices(), expected_count=2) + assert two.ok is True + assert two.decoded == ("A", "B") + + +def test_stage1_malformed_json_is_format_failure() -> None: + for text in ( + "not json", + "{", + "[]", + '{"candidates":["1","2",3,"4","5"]}', # non-string member + '{"candidates":["1","2","3","4","5"],"extra":1}', # extra key + '{"answer":"1"}', # wrong key + ): + result = check_stage1_choices(text, choices=_choices()) + assert result.format_valid is False, text + assert result.ok is False + assert result.decoded is None + + +def test_stage1_valid_choice_without_ground_truth_still_decodes() -> None: + result = check_stage1_choices('{"candidates":["1","2","4","5","6"]}', choices=_choices()) + assert result.ok is True + assert result.decoded == ("A", "B", "D", "E", "F") + + +def test_stage2_valid_local_answer_decodes_positionally() -> None: + result = check_stage2_choices('{"answer":"1"}', candidates=CANDIDATES) + assert result.ok is True + assert result.decoded == "C" + assert result.output is not None and result.output.answer == "1" + assert check_stage2_choices('{"answer":"5"}', candidates=CANDIDATES).decoded == "E" + + +def test_stage2_answer_outside_local_ids_fails_constraints() -> None: + for answer in ("0", "6", "A", "01"): + result = check_stage2_choices(f'{{"answer":"{answer}"}}', candidates=CANDIDATES) + assert result.format_valid is True, answer + assert result.constraint_valid is False, answer + assert result.decoded is None, answer + assert any("one of 1..5" in error for error in result.errors), answer + + +def test_stage2_malformed_json_is_format_failure() -> None: + for text in ( + "garbage", + "{", + '{"answer":5}', # non-string answer + '{"candidates":["C","A","B","D","E"]}', # wrong key + '{"answer":"1","why":"x"}', # extra key + ): + result = check_stage2_choices(text, candidates=CANDIDATES) + assert result.format_valid is False, text + assert result.ok is False + + +def test_choice_checks_never_raise_on_model_output() -> None: + garbage = [ + "", + None, # type: ignore[arg-type] + "{", + "[]", + '{"candidates": 3}', + '{"answer": null}', + "{}", + "\x00\x01binary", + ] + for text in garbage: + result = check_stage1_choices(text, choices=_choices()) # type: ignore[arg-type] + assert isinstance(result, ChoiceParseResult) + result2 = check_stage2_choices(text, candidates=CANDIDATES) # type: ignore[arg-type] + assert isinstance(result2, ChoiceParseResult) + + +def test_canonical_view_matches_parse_result_semantics() -> None: + # constraint-valid -> decoded canonical shape + view = check_stage1_choices( + '{"candidates":["3","1","2","4","5"]}', choices=_choices() + ).canonical_view() + assert isinstance(view, ParseResult) + assert view.ok is True + assert view.output is not None and view.output.candidates == ("C", "A", "B", "D", "E") + # format-valid but constraint-invalid -> model-level shape kept + bad = check_stage1_choices( + '{"candidates":["3","1","2","4","9"]}', choices=_choices() + ).canonical_view() + assert bad.format_valid is True and bad.constraint_valid is False + assert bad.output is not None and bad.output.candidates == ("3", "1", "2", "4", "9") + # format-invalid -> None output + junk = check_stage1_choices("nope", choices=_choices()).canonical_view() + assert junk.format_valid is False and junk.output is None + # stage2 canonical view carries the decoded answer + view2 = check_stage2_choices('{"answer":"2"}', candidates=CANDIDATES).canonical_view() + assert view2.ok is True + assert view2.output is not None and view2.output.answer == "A" + + +def test_programming_errors_still_raise() -> None: + with pytest.raises(ValueError): + check_stage1_choices( + '{"candidates":["1","2","3","4","5"]}', choices=None # type: ignore[arg-type] + ) + with pytest.raises(ValueError): + check_stage1_choices('{"candidates":["1","2"]}', choices=_choices(), expected_count=0) + with pytest.raises(ValueError): + check_stage2_choices('{"answer":"1"}', candidates=("C", "A")) # not 5 + with pytest.raises(ValueError): + check_stage2_choices('{"answer":"1"}', candidates="C") # bare string