From e881a6a418b78d6164cf062f35332ae21b830516 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:11:22 -0400 Subject: [PATCH 1/9] Calculator features one at a time: tool v3 and ablation arms The calculator's six features (names, placement, overrides, explain, concise, flat) are each on or off, so a study can change one at a time. "pe2" keeps tool version 2 (the first four, its notice byte for byte), "pe3" offers all six, and any other set runs as "pe:" plus its features, e.g. "pe:names+flat". - concise: an explanation gives the value, its direct inputs and up to five parameters unless `detail` is "full". - flat: the household is described as a state, people with a role (head, spouse or dependent) and facts, and household facts; the tool builds the units the benchmark's own situations have and places each household fact on its entity. Tests: the v2 notice is reproduced from its features; each feature adds only its own behavior; flat builds the benchmark's units and refuses what it cannot place; and on the installed engine eight frozen-board households described flat compute every board output exactly as Scenario.to_pe_household does. Five mutants of the flat and arm logic are all caught. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 522 +++++++++++++++++++++++++++++---- tests/test_tool_pilot_argon.py | 1 + tests/test_tool_v3.py | 377 ++++++++++++++++++++++++ 3 files changed, 850 insertions(+), 50 deletions(-) create mode 100644 tests/test_tool_v3.py diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index c8a1559b..88d8c4d6 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -152,18 +152,75 @@ # with no tools), so a prompt version other than the board's has its own # no-tools baseline. Repairs, retries and truncation then follow this pilot's # loop, not the board's; they matter only when a first answer does not parse. -ARMS = ("none", "pe", "web", "web+pe", "pe2") +ARMS = ("none", "pe", "web", "web+pe", "pe2", "pe3") ARM_SUFFIX = { "none": "+no-tools", "pe": "+pe-tool", "web": "+web", "web+pe": "+web+pe-tool", "pe2": "+pe-tool-v2", + "pe3": "+pe-tool-v3", } -# "pe2" is the calculator, tool version 2 (Max, 2026-10-23: "yes go"), alone: -# it reads the names the prompt uses, places inputs set on the wrong entity, -# flags inputs that replace a formula, and explains a value on request. -TOOL_VERSION_BY_ARM = {"pe": 1, "web+pe": 1, "pe2": 2} +# The calculator's features, each on or off, so a study can change one at a +# time (Max, 2026-10-23: "measuring this as an experiment could be +# publishable"; "do all these tests with a cheap model"): +# - names: a name written with spaces, or a benchmark output key, is read as +# its engine variable; +# - placement: an input set on another entity moves to the one unit holding +# its people; +# - overrides: a realized contribution is read as its _desired input, and an +# input that replaces a formula is flagged; +# - explain: `explain` traces up to three variables; +# - concise: an explanation is brief unless `detail` is "full"; +# - flat: the household is described as people with roles and facts +# (`household`), and the tool builds its units as the benchmark does. +# "pe2" is tool version 2 (the first four); "pe3" is all six. Any other set +# runs as an arm named "pe:" plus its features joined by "+", e.g. +# "pe:names+flat". +NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( + "names", + "placement", + "overrides", + "explain", + "concise", + "flat", +) +TOOL_FEATURES = (NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT) +V2_FEATURES = frozenset({NAMES, PLACEMENT, OVERRIDES, EXPLAIN}) +FEATURES_BY_ARM = { + "pe": frozenset(), + "web+pe": frozenset(), + "pe2": V2_FEATURES, + "pe3": frozenset(TOOL_FEATURES), +} +_FEATURE_ARM = re.compile(r"^pe:([a-z]+(?:\+[a-z]+)*)$") + + +def arm_features(arm: str) -> frozenset | None: + """The calculator features an arm offers, or None for an arm without the + calculator. Refuses an unknown feature, or concise without explain.""" + if arm in FEATURES_BY_ARM: + return FEATURES_BY_ARM[arm] + match = _FEATURE_ARM.match(arm or "") + if match is None: + return None + features = frozenset(match.group(1).split("+")) + unknown = sorted(features - set(TOOL_FEATURES)) + if unknown: + raise SystemExit( + f"unknown tool features {unknown}; the features are {TOOL_FEATURES}" + ) + if CONCISE in features and EXPLAIN not in features: + raise SystemExit("the concise feature needs explain") + return features + + +def arm_suffix(arm: str) -> str: + if arm in ARM_SUFFIX: + return ARM_SUFFIX[arm] + features = arm_features(arm) + return f"+pe-tool[{'+'.join(f for f in TOOL_FEATURES if f in features)}]" + TOOL_NOTICE_V2 = ( "\n\nYou also have a `calculate_household` function. It runs PolicyEngine US " @@ -183,6 +240,76 @@ "instructed above.", "Then give your final answers as the single JSON object described above.", ) + + +def _listing(items: list[str]) -> str: + if len(items) <= 2: + return " and ".join(items) + return ", ".join(items[:-1]) + ", and " + items[-1] + + +def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: + """The calculator notice for a set of features: the opening, a sentence + per feature, then the turn limit and the answer instruction. No feature + gives TOOL_NOTICE and the version 2 features give TOOL_NOTICE_V2, byte for + byte.""" + if not features: + return TOOL_NOTICE if tool_contract else JSON_TOOL_NOTICE + what = ( + "household you describe" + if FLAT in features + else "household situation you build" + ) + text = ( + "\n\nYou also have a `calculate_household` function. It runs PolicyEngine US " + "(policyengine-us 2.41.0), an open-source rules engine for US federal and " + f"state taxes and benefits, on a {what} from the facts above." + ) + if FLAT in features: + text += ( + " Describe it as `household`: each person with their role (head, spouse " + "or dependent) and facts, and the facts about the household as a " + "whole; it builds the household's units as the facts imply." + ) + if NAMES in features: + text += ( + " The facts above are labeled with its variable names, written with " + "spaces; it accepts them as written, and it accepts the output keys you " + "are asked for. Set every listed fact." + ) + notes = [] + if NAMES in features: + notes.append("which names it read") + if PLACEMENT in features: + notes.append("which inputs it moved to the entity they belong to") + if OVERRIDES in features: + notes.append("which inputs replace a formula") + explain = "" + if EXPLAIN in features: + explain = ( + "ask it to `explain` up to three variables to see in brief how it " + "computed them, with the parameters and laws it used, or every step " + 'with `detail` set to "full"' + if CONCISE in features + else "ask it to `explain` up to three variables to see how it computed " + "them, with the parameters and laws it used" + ) + if notes and explain: + text += f" Its result says {_listing(notes)}; {explain}." + elif notes: + text += f" Its result says {_listing(notes)}." + elif explain: + text += f" You may {explain}." + text += " You may call it in up to 8 turns. " + text += ( + "Then call `submit_outputs` exactly once with your final answers, as " + "instructed above." + if tool_contract + else "Then give your final answers as the single JSON object described above." + ) + return text + + # The engine version the notices and the tool description state: the board's # (set by command_run), so a household set on a newer engine says so. The # frozen board's 2.41.0 leaves every text byte for byte as before. @@ -511,19 +638,86 @@ def refuse_policyengine(label: str, value) -> None: MAX_EXPLAIN = 3 -def calculate_tool_schema(version: int = 1) -> dict: - if version == 2: +FLAT_DESCRIPTION = ( + "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " + "the requested variables. `household` is a JSON object: `state` (the " + "two-letter state code), `people` (a list, each with `name`, `role` -- " + 'head, spouse or dependent -- and `facts`, for example {"age": 40, ' + '"employment income": 30000}), and `household_facts` (facts about the ' + "household as a whole: its tax unit, family, SPM unit or home). The head " + "and spouse are one couple and everyone else is on their own; everyone " + "shares one tax unit, family, SPM unit and household. A fact is a value " + 'for `period`, or {period: value} with a period such as "2026" or ' + '"2026-01". Facts you do not give take the engine\'s defaults. `variables` ' + 'lists the variables to compute for `period` (default "2026"). Each ' + "result gives the variable's entity, label and value for each member of " + "that entity; a monthly variable computed for a year is summed over its " + "months." +) + + +def calculate_tool_schema(version: int = 1, features: frozenset | None = None) -> dict: + if features is None: + features = V2_FEATURES if version == 2 else frozenset() + if features: schema = calculate_tool_schema(1) function = schema["function"] - function["description"] = engine_text(CALCULATE_DESCRIPTION_V2) - function["parameters"]["properties"]["explain"] = { - "type": "array", - "items": {"type": "string"}, - "description": ( - f"Up to {MAX_EXPLAIN} variables to explain: how the engine " - "computed them." - ), - } + properties = function["parameters"]["properties"] + if features == V2_FEATURES: + function["description"] = engine_text(CALCULATE_DESCRIPTION_V2) + else: + description = ( + FLAT_DESCRIPTION if FLAT in features else CALCULATE_DESCRIPTION + ) + notes = [] + if NAMES in features: + notes.append( + "A name may also be written with spaces (as the facts are " + "labeled) or be one of the output keys you are asked for." + ) + if PLACEMENT in features and FLAT not in features: + notes.append( + "An input set on the wrong entity is moved to the one unit it " + "belongs to." + ) + if OVERRIDES in features: + notes.append( + "Setting a variable the engine computes replaces its formula; a " + "contribution set as the amount contributed is read as the " + "amount asked for before limits (`..._desired`)." + ) + if notes: + notes.append("The result's `notes` say what it did.") + function["description"] = engine_text(" ".join([description, *notes])) + if FLAT in features: + del properties["situation"] + properties["household"] = { + "type": "string", + "description": ( + "The household as a JSON object (state, people with role and " + "facts, household_facts), written as a JSON string." + ), + } + function["parameters"]["required"] = ["household", "variables"] + if EXPLAIN in features: + properties["explain"] = { + "type": "array", + "items": {"type": "string"}, + "description": ( + f"Up to {MAX_EXPLAIN} variables to explain: how the engine " + "computed them." + ), + } + if CONCISE in features: + properties["detail"] = { + "type": "string", + "enum": ["brief", "full"], + "description": ( + 'How much an explanation shows: "brief" (the default: the ' + 'value, what it came from, and the main parameters) or "full" ' + "(every step)." + ), + } return schema return { "type": "function", @@ -733,7 +927,12 @@ def _normalize_inputs( for name, value in data.items(): if allow_members and name == "members": continue - if v2 is not None and isinstance(name, str) and name not in registry: + if ( + v2 is not None + and NAMES in v2["features"] + and isinstance(name, str) + and name not in registry + ): resolved, note = resolve_name(name, registry, output_names=False) if resolved is not None: v2["notes"].append(f"{where}: {note}") @@ -745,7 +944,7 @@ def _normalize_inputs( ) continue info = registry[name] - if v2 is not None: + if v2 is not None and OVERRIDES in v2["features"]: desired = f"{name}{DESIRED_SUFFIX}" if ( info.computed @@ -763,7 +962,7 @@ def _normalize_inputs( "replaces its formula" ) if info.entity != entity: - if v2 is not None: + if v2 is not None and PLACEMENT in v2["features"]: v2["misplaced"].append((where, entity, name, value)) continue problems.append( @@ -782,7 +981,10 @@ def _normalize_inputs( def normalize_request( - args, registry: Mapping[str, VariableInfo], version: int = 1 + args, + registry: Mapping[str, VariableInfo], + version: int = 1, + features: frozenset | None = None, ) -> tuple[dict | None, dict | None]: """Check a calculate_household call; return (request, None) or (None, error). @@ -795,7 +997,9 @@ def normalize_request( ``explain``; the request's ``notes`` say what it did. """ problems: list[str] = [] - v2 = {"notes": [], "misplaced": []} if version == 2 else None + if features is None: + features = V2_FEATURES if version == 2 else frozenset() + v2 = {"notes": [], "misplaced": [], "features": features} if features else None def fail() -> tuple[None, dict]: return None, { @@ -818,9 +1022,10 @@ def fail() -> tuple[None, dict]: if not isinstance(args, dict): problems.append("arguments must be an object") return fail() - allowed = {"situation", "variables", "period"} | ( - {"explain"} if version == 2 else set() - ) + allowed = {"variables", "period"} + allowed |= {"household"} if FLAT in features else {"situation"} + allowed |= {"explain"} if EXPLAIN in features else set() + allowed |= {"detail"} if CONCISE in features else set() unknown = sorted(set(args) - allowed) if unknown: problems.append(f"unknown arguments: {', '.join(map(str, unknown))}") @@ -836,6 +1041,13 @@ def fail() -> tuple[None, dict]: ) period = str(YEAR) + if FLAT in features: + args = dict(args) + args["situation"] = flat_situation( + args.get("household"), registry, problems, v2 + ) + if args["situation"] is None: + return fail() situation = args.get("situation") if isinstance(situation, str): if len(situation) > MAX_SITUATION_CHARS: @@ -943,11 +1155,16 @@ def fail() -> tuple[None, dict]: variables = [] elif len(variables) > MAX_VARIABLES: problems.append(f"at most {MAX_VARIABLES} variables per call") - if v2 is not None: + if v2 is not None and PLACEMENT in features: _place_misplaced(v2, normalized, situation, registry, period, problems) requested = [] for name in variables: - if v2 is not None and isinstance(name, str) and name not in registry: + if ( + v2 is not None + and NAMES in features + and isinstance(name, str) + and name not in registry + ): resolved, note = resolve_name(name, registry, output_names=True) if resolved is not None: v2["notes"].append(f"variables: {note}") @@ -959,7 +1176,7 @@ def fail() -> tuple[None, dict]: elif name not in requested: requested.append(name) explain = [] - if v2 is not None and args.get("explain") is not None: + if EXPLAIN in features and args.get("explain") is not None: names = args["explain"] if isinstance(names, str): names = [names] @@ -969,7 +1186,7 @@ def fail() -> tuple[None, dict]: elif len(names) > MAX_EXPLAIN: problems.append(f"explain at most {MAX_EXPLAIN} variables per call") for name in names: - if name not in registry: + if NAMES in features and name not in registry: resolved, note = resolve_name(name, registry, output_names=True) if resolved is not None: v2["notes"].append(f"explain: {note}") @@ -981,15 +1198,154 @@ def fail() -> tuple[None, dict]: elif name not in explain: explain.append(name) + detail = "full" + if CONCISE in features: + detail = args.get("detail") or "brief" + if detail not in ("brief", "full"): + problems.append('detail must be "brief" or "full"') + if problems: return fail() request = {"situation": normalized, "variables": requested, "period": period} if v2 is not None: request["explain"] = explain request["notes"] = v2["notes"] + request["detail"] = detail return request, None +FLAT_ROLES = ("head", "spouse", "dependent") + + +def flat_situation(household, registry, problems: list[str], v2: dict | None): + """The flat feature's household -- state, people with role and facts, and + household facts -- as a situation, with the units the benchmark's own + reference situations have (policybench.scenarios.Scenario.to_pe_household): + the head and spouse one marital unit and everyone else alone, and one tax + unit, SPM unit, family and household holding everyone. A household fact + goes to the unit of its variable's entity. Returns None after recording + problems.""" + if isinstance(household, str): + if len(household) > MAX_SITUATION_CHARS: + problems.append("household is too large") + return None + try: + household = json.loads(household) + except (ValueError, RecursionError) as error: + problems.append(f"household is not valid JSON: {error}"[:MAX_MESSAGE_CHARS]) + return None + if not isinstance(household, dict): + problems.append( + "household must be a JSON object with state, people and household_facts" + ) + return None + extra = sorted(set(household) - {"state", "people", "household_facts"}) + if extra: + problems.append( + "household accepts state, people and household_facts; remove " + + ", ".join(map(str, extra)) + ) + people = household.get("people") + if not isinstance(people, list) or not people: + problems.append("household.people must be a non-empty list of people") + return None + if len(people) > MAX_PEOPLE: + problems.append(f"at most {MAX_PEOPLE} people") + return None + situation: dict[str, Any] = {"people": {}} + roles: dict[str, str] = {} + for index, person in enumerate(people): + if not isinstance(person, dict): + problems.append( + f"people[{index}] must be an object with name, role and facts" + ) + continue + name, role, facts = ( + person.get("name"), + person.get("role"), + person.get("facts", {}), + ) + if not isinstance(name, str) or not name or len(name) > 64: + problems.append(f"people[{index}].name must be a short string") + continue + if name in roles: + problems.append(f"people[{index}]: {name!r} is listed twice") + continue + if role not in FLAT_ROLES: + problems.append( + f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" + ) + continue + if not isinstance(facts, dict): + problems.append(f"people[{index}].facts must be an object") + continue + roles[name] = role + situation["people"][name] = dict(facts) + heads = [n for n, r in roles.items() if r == "head"] + spouses = [n for n, r in roles.items() if r == "spouse"] + if len(heads) != 1: + problems.append("the household needs exactly one head") + if len(spouses) > 1: + problems.append("the household has at most one spouse") + if problems: + return None + for name, role in roles.items(): + facts = situation["people"][name] + facts.setdefault("is_tax_unit_head", role == "head") + facts.setdefault("is_tax_unit_spouse", role == "spouse") + everyone = list(situation["people"]) + couple = heads + spouses + groups = [couple] + [[n] for n in everyone if n not in couple] + situation["marital_units"] = { + f"marital_unit_{i}": {"members": members} + for i, members in enumerate(groups, start=1) + } + situation["tax_units"] = {"tax_unit": {"members": list(everyone)}} + situation["spm_units"] = {"spm_unit": {"members": list(everyone)}} + situation["families"] = {"family": {"members": list(everyone)}} + situation["households"] = {"household": {"members": list(everyone)}} + state = household.get("state") + if state is not None: + situation["households"]["household"]["state_code"] = state + unit_of = { + "tax_unit": situation["tax_units"]["tax_unit"], + "spm_unit": situation["spm_units"]["spm_unit"], + "family": situation["families"]["family"], + "household": situation["households"]["household"], + "marital_unit": situation["marital_units"]["marital_unit_1"], + } + facts = household.get("household_facts") or {} + if not isinstance(facts, dict): + problems.append("household_facts must be an object") + return None + for name, value in facts.items(): + resolved = name + if ( + isinstance(name, str) + and name not in registry + and v2 is not None + and NAMES in v2["features"] + ): + resolved, note = resolve_name(name, registry, output_names=False) + if resolved is not None: + v2["notes"].append(f"household_facts: {note}") + if not isinstance(resolved, str) or resolved not in registry: + problems.append( + f"household_facts: unknown variable {name!r}." + + (_suggest(name, registry) if isinstance(name, str) else "") + ) + continue + entity = registry[resolved].entity + if entity == "person": + problems.append( + f"household_facts: {resolved} is a person's fact; give it under " + "that person" + ) + continue + unit_of[entity][resolved] = value + return None if problems else situation + + DESIRED_SUFFIX = "_desired" _LABEL_SEPARATORS = re.compile(r"[\s\-]+") @@ -1152,12 +1508,40 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic out["notes"] = list(request["notes"]) if request.get("explain"): out["explanations"] = { - name: explain_variable(system, request, name, simulation_cls) + name: ( + brief_explanation( + explain_variable(system, request, name, simulation_cls) + ) + if request.get("detail") == "brief" + else explain_variable(system, request, name, simulation_cls) + ) for name in request["explain"] } return out +BRIEF_INPUTS = 8 +BRIEF_PARAMETERS = 5 + + +def brief_explanation(full: dict) -> dict: + """The concise feature's explanation: the value, the variables it came + from directly (one level down, zero-valued ones left out), and the first + parameters it read with their sources.""" + if "error" in full: + return full + steps = full["steps"] + direct = [ + s.strip() for s in steps[1:] if s.startswith(" ") and not s.startswith(" ") + ] + return { + "value": steps[0].strip() if steps else "", + "from": direct[:BRIEF_INPUTS], + "parameters": full["parameters"][:BRIEF_PARAMETERS], + "more": 'set detail to "full" for every step', + } + + EXPLAIN_MAX_LINES = 60 EXPLAIN_MAX_DEPTH = 6 EXPLAIN_MAX_PARAMETERS = 15 @@ -1356,10 +1740,13 @@ def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> return if args is None: return - version = 1 - if isinstance(args, dict) and args.get("__tool_version__") == 2: - version, args = 2, args.get("arguments") - request, error = normalize_request(args, registry, version) + features = frozenset() + if isinstance(args, dict) and "__tool_features__" in args: + features = frozenset(args.get("__tool_features__") or ()) + args = args.get("arguments") + elif isinstance(args, dict) and args.get("__tool_version__") == 2: + features, args = V2_FEATURES, args.get("arguments") + request, error = normalize_request(args, registry, features=features) if error is not None: conn.send(("ok", error)) continue @@ -1379,9 +1766,15 @@ def __init__( acquire_timeout: float | None = None, board: Path | None = None, tool_version: int = 1, + tool_features: frozenset | None = None, ): self._ctx = mp.get_context("spawn") self._tool_version = tool_version + self._tool_features = ( + tool_features + if tool_features is not None + else (V2_FEATURES if tool_version == 2 else frozenset()) + ) self._fix_path = str(fix_path) # A worker gets the board only when it is not the frozen one, so a # three-argument worker target still starts on the default. @@ -1448,8 +1841,13 @@ def call(self, args) -> tuple[dict, float]: time.perf_counter() - started, ) process, conn = worker - if self._tool_version == 2: + if self._tool_features == V2_FEATURES: args = {"__tool_version__": 2, "arguments": args} + elif self._tool_features: + args = { + "__tool_features__": sorted(self._tool_features), + "arguments": args, + } try: conn.send(args) if conn.poll(self.timeout): @@ -1607,6 +2005,8 @@ def notice(self) -> str: notice = WEB_TOOL_NOTICE if tool else JSON_WEB_TOOL_NOTICE elif self.arm == "pe2": return engine_text(TOOL_NOTICE_V2 if tool else JSON_TOOL_NOTICE_V2) + elif self.arm != "pe" and arm_features(self.arm) is not None: + return engine_text(compose_tool_notice(arm_features(self.arm), tool)) else: return engine_text(TOOL_NOTICE if tool else JSON_TOOL_NOTICE) return engine_text(notice) + ( @@ -1615,11 +2015,15 @@ def notice(self) -> str: @property def uses_calculator(self) -> bool: - return self.arm in TOOL_VERSION_BY_ARM + return arm_features(self.arm) is not None + + @property + def tool_features(self) -> frozenset: + return arm_features(self.arm) or frozenset() @property def tool_version(self) -> int: - return TOOL_VERSION_BY_ARM.get(self.arm, 1) + return 2 if self.tool_features == V2_FEATURES else 1 @property def uses_web(self) -> bool: @@ -1637,7 +2041,7 @@ def arm1_label(self) -> str: def arm2_label(self) -> str: version = "" if self.prompt_version == "v1" else f"@{self.prompt_version}" rule = "+no-benchmark-search" if self.benchmark_search_rule else "" - return f"{self.name}{ARM_SUFFIX[self.arm]}{version}{rule}" + return f"{self.name}{arm_suffix(self.arm)}{version}{rule}" def provider_family(model_id: str) -> str: @@ -1663,8 +2067,11 @@ def model_profile( _uses_responses_api, ) - if arm not in ARMS: - raise SystemExit(f"unknown arm {arm!r}; the arms are {', '.join(ARMS)}") + if arm not in ARMS and arm_features(arm) is None: + raise SystemExit( + f"unknown arm {arm!r}; the arms are {', '.join(ARMS)}, or pe: plus " + f"calculator features joined by + ({', '.join(TOOL_FEATURES)})" + ) if prompt_version not in PROMPT_VERSIONS: raise SystemExit( f"unknown prompt {prompt_version!r}; the prompts are " @@ -2298,7 +2705,7 @@ def run_household( sent_board_prompt_sha256 = base.pop("_sent_board_prompt_sha256") ceiling = base.pop("_budget_ceiling") answer_tool = base.pop("_answer_tool") - calculate_tool = calculate_tool_schema(profile.tool_version) + calculate_tool = calculate_tool_schema(features=profile.tool_features) offered = ([web_search_tool(profile)] if profile.uses_web else []) + ( [calculate_tool] if profile.uses_calculator else [] ) @@ -3133,7 +3540,10 @@ def command_run(args) -> None: "api": profile.api, "answer_contract": profile.contract, "tool_notice": profile.notice, - "calculate_tool": calculate_tool_schema(profile.tool_version) + "calculate_tool": calculate_tool_schema(features=profile.tool_features) + if profile.uses_calculator + else None, + "tool_features": sorted(profile.tool_features) if profile.uses_calculator else None, "tool_version": profile.tool_version if profile.uses_calculator else None, @@ -3164,7 +3574,12 @@ def command_run(args) -> None: return # Arms without the calculator start no engine workers. pool = ( - ToolPool(size=args.workers, board=snapshot, tool_version=profile.tool_version) + ToolPool( + size=args.workers, + board=snapshot, + tool_version=profile.tool_version, + tool_features=profile.tool_features, + ) if profile.uses_calculator else None ) @@ -3879,6 +4294,12 @@ def _successful_calls(record: dict) -> list[dict]: def _person_names(arguments) -> list[str]: try: args = json.loads(arguments) if isinstance(arguments, str) else arguments + if "household" in args: # the flat feature's people list + household = args["household"] + household = ( + json.loads(household) if isinstance(household, str) else household + ) + return [person["name"] for person in household["people"]] situation = args["situation"] situation = json.loads(situation) if isinstance(situation, str) else situation return list(situation["people"]) @@ -4486,7 +4907,7 @@ def main_call(record: dict, scenario, outputs: list[str]) -> dict | None: def situation_fidelity( - records, board, system, registry, defaults, out: Path, version: int = 1 + records, board, system, registry, defaults, out: Path, features=frozenset() ) -> None: """How each household's main situation differs from the reference's. @@ -4508,7 +4929,7 @@ def situation_fidelity( if call is None: rows.append({"scenario_id": sid, "differences": None}) continue - request, _ = normalize_request(call["arguments"], registry, version) + request, _ = normalize_request(call["arguments"], registry, features=features) if request is None: rows.append({"scenario_id": sid, "differences": None}) continue @@ -4561,7 +4982,7 @@ def command_diagnose(args) -> None: out = Path(args.out_dir) records = load_records(out) # A run's calls are read as its own tool version read them. - version = TOOL_VERSION_BY_ARM.get(records_arm(records), 1) + features = arm_features(records_arm(records)) or frozenset() rows_path = out / "scored_rows.csv" if not rows_path.exists(): raise SystemExit("run score first") @@ -4577,7 +4998,7 @@ def command_diagnose(args) -> None: name: system.variables[name].default_value for name in system.variables } if registry is not None: - situation_fidelity(records, board, system, registry, defaults, out, version) + situation_fidelity(records, board, system, registry, defaults, out, features) out_rows = [] if misses.empty: print("no arm-2 misses") @@ -4595,7 +5016,7 @@ def command_diagnose(args) -> None: if registry is not None: def board_value_for(arguments, scenario=scenario, output=row.variable): - request, _ = normalize_request(arguments, registry, version) + request, _ = normalize_request(arguments, registry, features=features) if request is None: return None return output_value(system, request["situation"], scenario, output) @@ -4619,7 +5040,7 @@ def board_value_for(arguments, scenario=scenario, output=row.variable): attribution = {} arguments = verdict.get("call_arguments") if arguments is not None and registry is not None: - request, _ = normalize_request(arguments, registry, version) + request, _ = normalize_request(arguments, registry, features=features) if request is not None: reference_situation = scenario.to_pe_household() diff = situation_diff( @@ -4716,9 +5137,10 @@ def main(argv: list[str] | None = None) -> None: ) p.add_argument( "--arm", - choices=ARMS, default="pe", - help="the tools: none, PolicyEngine (pe), web search (web), or both", + help="the tools: none, PolicyEngine (pe, pe2, pe3, or pe: plus " + "calculator features joined by +), web search (web), or both " + "(web+pe)", ) p.add_argument( "--prompt", diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index 986a89d9..4e586374 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -2294,6 +2294,7 @@ def test_web_notices_name_search_and_keep_the_answer_instruction(): "web": "gpt-6-sol+web", "web+pe": "gpt-6-sol+web+pe-tool", "pe2": "gpt-6-sol+pe-tool-v2", + "pe3": "gpt-6-sol+pe-tool-v3", } with pytest.raises(SystemExit): pilot.model_profile(SOL, "search") diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py new file mode 100644 index 00000000..5c018fd1 --- /dev/null +++ b/tests/test_tool_v3.py @@ -0,0 +1,377 @@ +"""The calculator's features one at a time (tool v3 and the ablation arms). + +Each feature (names, placement, overrides, explain, concise, flat) must work +on its own and change nothing else, so a study can turn one on at a time. The +version 2 feature set must reproduce tool v2's notice byte for byte, and the +flat household format must build the units the benchmark's own reference +situations have. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) + +import tool_pilot_argon as pilot # noqa: E402 + +V = pilot.VariableInfo +REGISTRY = { + "age": V("person", "int", "age", "year"), + "employment_income": V("person", "float", "employment income", "year"), + "is_tax_unit_head": V("person", "bool", "head", "year"), + "is_tax_unit_spouse": V("person", "bool", "spouse", "year"), + "traditional_401k_contributions": V( + "person", "float", "401(k)", "year", computed=True + ), + "traditional_401k_contributions_desired": V("person", "float", "asked", "year"), + "snap": V("spm_unit", "float", "SNAP", "month", computed=True), + "pre_subsidy_rent": V("spm_unit", "float", "rent", "year"), + "qualified_passenger_vehicle_loan_interest": V( + "household", "float", "car loan", "year" + ), + "state_code": V("household", "Enum", "state", "year", ("CA", "LA")), + "tax_unit_childcare_expenses": V("tax_unit", "float", "childcare", "year"), + "marital_unit_flag": V("marital_unit", "bool", "flag", "year"), +} +F = pilot + + +def _situation(): + return { + "people": {"head": {"age": 40}, "kid": {"age": 8}}, + "spm_units": {"spm": {"members": ["head", "kid"]}}, + "households": {"h": {"members": ["head", "kid"], "state_code": "CA"}}, + } + + +def _norm(args, *features): + return pilot.normalize_request(args, REGISTRY, features=frozenset(features)) + + +def test_the_version_2_set_reproduces_tool_v2s_notice(): + assert F.compose_tool_notice(F.V2_FEATURES, True) == F.TOOL_NOTICE_V2 + assert F.compose_tool_notice(F.V2_FEATURES, False) == F.JSON_TOOL_NOTICE_V2 + assert F.compose_tool_notice(frozenset(), True) == F.TOOL_NOTICE + assert F.compose_tool_notice(frozenset(), False) == F.JSON_TOOL_NOTICE + + +def test_each_feature_adds_its_own_sentence_only(): + base = F.compose_tool_notice(frozenset({F.EXPLAIN}), True) + assert "`explain`" in base and "names it read" not in base + assert "written with spaces" not in base + names = F.compose_tool_notice(frozenset({F.NAMES}), True) + assert "written with spaces" in names and "`explain`" not in names + flat = F.compose_tool_notice(frozenset({F.FLAT}), True) + assert "household you describe" in flat and "`household`" in flat + concise = F.compose_tool_notice(frozenset({F.EXPLAIN, F.CONCISE}), True) + assert '`detail` set to "full"' in concise + + +def test_arms_name_their_features(): + assert F.arm_features("pe") == frozenset() + assert F.arm_features("pe2") == F.V2_FEATURES + assert F.arm_features("pe3") == frozenset(F.TOOL_FEATURES) + assert F.arm_features("pe:flat+names") == frozenset({"flat", "names"}) + assert F.arm_features("none") is None and F.arm_features("web") is None + assert F.arm_suffix("pe:flat+names") == "+pe-tool[names+flat]" + with pytest.raises(SystemExit, match="unknown tool features"): + F.arm_features("pe:names+telepathy") + with pytest.raises(SystemExit, match="concise feature needs explain"): + F.arm_features("pe:concise") + profile = F.ModelProfile("m", "gpt-x", "chat", "tool", arm="pe:names") + assert profile.uses_calculator and profile.tool_features == frozenset({"names"}) + assert profile.tool_version == 1 + assert "written with spaces" in profile.notice + assert profile.arm2_label == "m+pe-tool[names]" + + +def test_names_alone_reads_labels_but_places_and_redirects_nothing(): + situation = _situation() + situation["people"]["head"]["Employment income"] = 30000 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.NAMES) + assert error is None + assert request["situation"]["people"]["head"]["employment_income"] == { + "2026": 30000 + } + misplaced = _situation() + misplaced["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 5 + request, error = _norm({"situation": misplaced, "variables": ["snap"]}, F.NAMES) + assert request is None and any( + "set it under households" in p for p in error["problems"] + ) + realized = _situation() + realized["people"]["head"]["traditional_401k_contributions"] = 100 + request, error = _norm({"situation": realized, "variables": ["snap"]}, F.NAMES) + assert "traditional_401k_contributions" in request["situation"]["people"]["head"] + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, F.NAMES + ) + assert request is None and any( + "unknown arguments: explain" in p for p in error["problems"] + ) + + +def test_placement_alone_moves_inputs_but_reads_no_labels(): + situation = _situation() + situation["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 5 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.PLACEMENT) + assert error is None + assert request["situation"]["households"]["h"][ + "qualified_passenger_vehicle_loan_interest" + ] == {"2026": 5} + labeled = _situation() + labeled["people"]["head"]["employment income"] = 1 + request, error = _norm({"situation": labeled, "variables": ["snap"]}, F.PLACEMENT) + assert request is None + + +def test_overrides_alone_redirects_and_flags(): + situation = _situation() + situation["people"]["head"]["traditional_401k_contributions"] = 100 + situation["spm_units"]["spm"]["snap"] = 5 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.OVERRIDES) + assert error is None + head = request["situation"]["people"]["head"] + assert head == { + "age": {"2026": 40}, + "traditional_401k_contributions_desired": {"2026": 100}, + } + assert any("snap is computed" in n for n in request["notes"]) + + +def test_concise_takes_a_detail_and_defaults_to_brief(): + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, + F.EXPLAIN, + F.CONCISE, + ) + assert error is None and request["detail"] == "brief" + request, _ = _norm( + { + "situation": _situation(), + "variables": ["snap"], + "explain": ["snap"], + "detail": "full", + }, + F.EXPLAIN, + F.CONCISE, + ) + assert request["detail"] == "full" + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "detail": "short"}, + F.EXPLAIN, + F.CONCISE, + ) + assert request is None + request, _ = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, + F.EXPLAIN, + ) + assert request["detail"] == "full" + + +def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_parameters(): + full = { + "steps": [ + "snap (2026) = 288", + " snap_normal_allotment (2026) = 288", + " snap_net_income (2026) = 1700", + " is_snap_eligible (2026) = True", + ], + "parameters": [{"name": f"p{i}", "value": i, "sources": []} for i in range(9)], + } + brief = F.brief_explanation(full) + assert brief["value"] == "snap (2026) = 288" + assert brief["from"] == [ + "snap_normal_allotment (2026) = 288", + "is_snap_eligible (2026) = True", + ] + assert [p["name"] for p in brief["parameters"]] == ["p0", "p1", "p2", "p3", "p4"] + assert F.brief_explanation({"error": "x"}) == {"error": "x"} + + +def _flat(**overrides): + household = { + "state": "LA", + "people": [ + { + "name": "Head", + "role": "head", + "facts": {"age": 40, "employment_income": 30000}, + }, + {"name": "Spouse", "role": "spouse", "facts": {"age": 38}}, + {"name": "Kid", "role": "dependent", "facts": {"age": 6}}, + ], + "household_facts": { + "pre_subsidy_rent": 9000, + "tax_unit_childcare_expenses": 1200, + "qualified_passenger_vehicle_loan_interest": 300, + }, + } + household.update(overrides) + return {"household": json.dumps(household), "variables": ["snap"]} + + +def test_flat_builds_the_benchmarks_units(): + request, error = _norm(_flat(), F.FLAT) + assert error is None, error + s = request["situation"] + assert s["marital_units"] == { + "marital_unit_1": {"members": ["Head", "Spouse"]}, + "marital_unit_2": {"members": ["Kid"]}, + } + for group, unit in ( + ("tax_units", "tax_unit"), + ("spm_units", "spm_unit"), + ("families", "family"), + ): + assert s[group][unit]["members"] == ["Head", "Spouse", "Kid"] + assert s["households"]["household"]["state_code"] == {"2026": "LA"} + assert s["spm_units"]["spm_unit"]["pre_subsidy_rent"] == {"2026": 9000} + assert s["tax_units"]["tax_unit"]["tax_unit_childcare_expenses"] == {"2026": 1200} + assert s["households"]["household"][ + "qualified_passenger_vehicle_loan_interest" + ] == {"2026": 300} + assert s["people"]["Head"]["is_tax_unit_head"] == {"2026": True} + assert s["people"]["Spouse"]["is_tax_unit_spouse"] == {"2026": True} + assert s["people"]["Kid"]["is_tax_unit_head"] == {"2026": False} + + +def test_flat_refuses_what_it_cannot_place(): + people = [ + {"name": "A", "role": "head", "facts": {}}, + {"name": "B", "role": "head", "facts": {}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("exactly one head" in p for p in error["problems"]) + request, error = _norm(_flat(household_facts={"age": 3}), F.FLAT) + assert any("person's fact" in p for p in error["problems"]) + people = [{"name": "A", "role": "boss", "facts": {}}] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("role must be one of" in p for p in error["problems"]) + request, error = _norm({"situation": _situation(), "variables": ["snap"]}, F.FLAT) + assert any("unknown arguments: situation" in p for p in error["problems"]) + request, error = _norm(_flat(), F.NAMES) + assert any("unknown arguments: household" in p for p in error["problems"]) + + +def test_flat_with_names_reads_labeled_household_facts(): + request, error = _norm( + _flat(household_facts={"pre subsidy rent": 9000}), F.FLAT, F.NAMES + ) + assert error is None + assert request["situation"]["spm_units"]["spm_unit"]["pre_subsidy_rent"] == { + "2026": 9000 + } + assert any("as pre_subsidy_rent" in n for n in request["notes"]) + + +def test_the_schema_offers_each_features_arguments(): + flat = F.calculate_tool_schema(features=frozenset({F.FLAT}))["function"] + props = flat["parameters"]["properties"] + assert "household" in props and "situation" not in props + assert flat["parameters"]["required"] == ["household", "variables"] + concise = F.calculate_tool_schema(features=frozenset({F.EXPLAIN, F.CONCISE})) + props = concise["function"]["parameters"]["properties"] + assert "explain" in props and props["detail"]["enum"] == ["brief", "full"] + assert F.calculate_tool_schema(features=F.V2_FEATURES) == F.calculate_tool_schema(2) + assert F.calculate_tool_schema(features=frozenset()) == F.calculate_tool_schema(1) + + +def _echo_worker(conn, fix_path, reform): + conn.send(("ready", None)) + while True: + args = conn.recv() + if args is None: + return + conn.send(("ok", {"received": args})) + + +def test_the_pool_tells_its_workers_the_features(): + pool = F.ToolPool( + size=1, + timeout=5.0, + worker_target=_echo_worker, + tool_features=frozenset({F.FLAT, F.NAMES}), + ) + try: + result, _ = pool.call({"x": 1}) + assert result == { + "received": {"__tool_features__": ["flat", "names"], "arguments": {"x": 1}} + } + finally: + pool.close() + + +@pytest.mark.slow +def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): + """Differential: on the installed engine, each frozen-board household + described in the flat format computes every board output exactly as the + benchmark's own situation (Scenario.to_pe_household) does.""" + import household_set as hs + import pandas as pd + from policyengine_us import Simulation + + from policybench.ground_truth import _extract_person_value, _pe_variable_for_output + + system = F.build_tool_system(reform=True) + registry = F.build_registry(system) + reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") + variables = reference.groupby("scenario_id").variable.agg(list) + relations = {"is_tax_unit_head", "is_tax_unit_spouse"} + checked = 0 + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:8]: + situation = hs.situation_for(scenario) + couple = scenario.marital_couple() or () + head = couple[0] if couple else scenario.adults[0].name + people = [ + { + "name": name, + "role": "head" + if name == head + else ("spouse" if name in couple else "dependent"), + "facts": {k: v for k, v in facts.items() if k not in relations}, + } + for name, facts in situation["people"].items() + ] + household_facts = {} + for group in ("tax_units", "spm_units", "families", "households"): + for unit in situation[group].values(): + household_facts.update( + { + k: v + for k, v in unit.items() + if k not in ("members", "state_code") + } + ) + household = { + "state": scenario.state, + "people": people, + "household_facts": household_facts, + } + request, error = F.normalize_request( + {"household": json.dumps(household), "variables": ["snap"]}, + registry, + features=frozenset({F.FLAT}), + ) + assert error is None, (scenario.id, error) + own = Simulation(tax_benefit_system=system, situation=situation) + flat = Simulation(tax_benefit_system=system, situation=request["situation"]) + for variable in variables[scenario.id]: + name = _pe_variable_for_output(variable, "us") + values = [ + float( + _extract_person_value( + simulation.calculate(name, 2026), scenario, variable + ) + ) + for simulation in (own, flat) + ] + assert values[0] == values[1], (scenario.id, variable, values) + checked += 1 + assert checked > 80 From b3ec406e62b19d59f285150ddc8a6207a2cf0bf0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:25:47 -0400 Subject: [PATCH 2/9] Read facts as the prompt labels them (labels feature) The prompts label some facts in words rather than as the variable name with spaces: employment_income is "gross wages and salaries", hours_worked_last_week "usual weekly hours worked", has_esi "has employer-sponsored insurance" (policybench.prompts.INPUT_LABEL_OVERRIDES). In the 20-household test set nearly every household has at least one; hours appear in 17, employer coverage and premiums in 11 each. Tool version 2's names feature cannot read these, and its notice says the facts carry the variable names. A Flash-Lite smoke test of the flat format met three such errors and fell back to setting employer_sponsored_insurance_premiums to 0 and weekly_hours_worked, which the SNAP work rules do not read. The new labels feature reads a fact under its prompt label (case and punctuation aside, renames applied) and, as the reference situations do, also sets stated usual weekly hours as weekly_hours_worked_before_lsr (PE_INPUT_ALIASES). pe3 now offers all seven features. pe2 is unchanged and keeps its notice as it ran; every other feature set's notice says only what the tool reads. Tests: label reading, the alias (and leaving a set value alone), pe2 unchanged, the variable-name claim only in pe2's notice across every feature set, every prompt label reading back to its engine name, and a differential: twelve frozen-board households described flat with each fact as its prompt writes it compute every board output as the benchmark does. Three mutants (no alias, no renames, strict label matching) are caught. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 130 +++++++++++++++++++++------ tests/test_tool_v3.py | 175 ++++++++++++++++++++++++++++++++++++ 2 files changed, 279 insertions(+), 26 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 88d8c4d6..58ec78bb 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -51,6 +51,7 @@ import copy import decimal import difflib +import functools import hashlib import importlib.metadata import importlib.util @@ -166,6 +167,12 @@ # publishable"; "do all these tests with a cheap model"): # - names: a name written with spaces, or a benchmark output key, is read as # its engine variable; +# - labels: a fact written under the label the prompt gives it (a prompt +# labels some facts in words, e.g. employment_income as "gross wages and +# salaries"; policybench.prompts.INPUT_LABEL_OVERRIDES) is read as its +# variable, and a fact the reference also sets under another name +# (policybench.scenarios.PE_INPUT_ALIASES: usual weekly hours) is set there +# too; # - placement: an input set on another entity moves to the one unit holding # its people; # - overrides: a realized contribution is read as its _desired input, and an @@ -174,18 +181,20 @@ # - concise: an explanation is brief unless `detail` is "full"; # - flat: the household is described as people with roles and facts # (`household`), and the tool builds its units as the benchmark does. -# "pe2" is tool version 2 (the first four); "pe3" is all six. Any other set +# "pe2" is tool version 2 (names, placement, overrides and explain); "pe3" is +# all seven. Any other set # runs as an arm named "pe:" plus its features joined by "+", e.g. # "pe:names+flat". -NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( +NAMES, LABELS, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( "names", + "labels", "placement", "overrides", "explain", "concise", "flat", ) -TOOL_FEATURES = (NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT) +TOOL_FEATURES = (NAMES, LABELS, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT) V2_FEATURES = frozenset({NAMES, PLACEMENT, OVERRIDES, EXPLAIN}) FEATURES_BY_ARM = { "pe": frozenset(), @@ -251,10 +260,14 @@ def _listing(items: list[str]) -> str: def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: """The calculator notice for a set of features: the opening, a sentence per feature, then the turn limit and the answer instruction. No feature - gives TOOL_NOTICE and the version 2 features give TOOL_NOTICE_V2, byte for - byte.""" + gives TOOL_NOTICE and the version 2 features TOOL_NOTICE_V2, as those arms + ran. TOOL_NOTICE_V2 says the facts are labeled with the engine's variable + names, which holds for most facts but not those the prompt labels in words + ("gross wages and salaries"); every other set says only what it reads.""" if not features: return TOOL_NOTICE if tool_contract else JSON_TOOL_NOTICE + if features == V2_FEATURES: + return TOOL_NOTICE_V2 if tool_contract else JSON_TOOL_NOTICE_V2 what = ( "household you describe" if FLAT in features @@ -271,14 +284,24 @@ def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: "or dependent) and facts, and the facts about the household as a " "whole; it builds the household's units as the facts imply." ) - if NAMES in features: + if LABELS in features and NAMES in features: + text += ( + " It accepts the facts above under the labels they have there, its " + "variable names written with spaces, and the output keys you are asked " + "for. Set every listed fact." + ) + elif LABELS in features: + text += ( + " It accepts the facts above under the labels they have there. Set " + "every listed fact." + ) + elif NAMES in features: text += ( - " The facts above are labeled with its variable names, written with " - "spaces; it accepts them as written, and it accepts the output keys you " - "are asked for. Set every listed fact." + " It accepts its variable names written with spaces, and the output " + "keys you are asked for. Set every listed fact." ) notes = [] - if NAMES in features: + if NAMES in features or LABELS in features: notes.append("which names it read") if PLACEMENT in features: notes.append("which inputs it moved to the entity they belong to") @@ -670,10 +693,15 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - FLAT_DESCRIPTION if FLAT in features else CALCULATE_DESCRIPTION ) notes = [] + if LABELS in features: + notes.append( + "A fact may be written under the label it has in the facts " + "you were given." + ) if NAMES in features: notes.append( - "A name may also be written with spaces (as the facts are " - "labeled) or be one of the output keys you are asked for." + "A name may also be written with spaces or be one of the " + "output keys you are asked for." ) if PLACEMENT in features and FLAT not in features: notes.append( @@ -927,13 +955,8 @@ def _normalize_inputs( for name, value in data.items(): if allow_members and name == "members": continue - if ( - v2 is not None - and NAMES in v2["features"] - and isinstance(name, str) - and name not in registry - ): - resolved, note = resolve_name(name, registry, output_names=False) + if v2 is not None and isinstance(name, str) and name not in registry: + resolved, note = read_fact_name(name, registry, v2["features"]) if resolved is not None: v2["notes"].append(f"{where}: {note}") name = resolved @@ -1157,6 +1180,8 @@ def fail() -> tuple[None, dict]: problems.append(f"at most {MAX_VARIABLES} variables per call") if v2 is not None and PLACEMENT in features: _place_misplaced(v2, normalized, situation, registry, period, problems) + if v2 is not None and LABELS in features: + _set_aliases(v2, normalized["people"], registry) requested = [] for name in variables: if ( @@ -1320,13 +1345,8 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None for name, value in facts.items(): resolved = name - if ( - isinstance(name, str) - and name not in registry - and v2 is not None - and NAMES in v2["features"] - ): - resolved, note = resolve_name(name, registry, output_names=False) + if isinstance(name, str) and name not in registry and v2 is not None: + resolved, note = read_fact_name(name, registry, v2["features"]) if resolved is not None: v2["notes"].append(f"household_facts: {note}") if not isinstance(resolved, str) or resolved not in registry: @@ -1375,6 +1395,64 @@ def output_reading(name: str) -> tuple[str, str] | None: return engine, how +def _label_key(text: str) -> str: + """A label as compared: lower case, letters and digits only, so "Roth + 401(k) contributions" and "roth 401k contributions" read alike.""" + return re.sub(r"[^0-9a-z]", "", text.casefold()) + + +@functools.lru_cache(maxsize=1) +def prompt_labels() -> dict[str, str]: + """The labels the prompts give facts in words, keyed by _label_key, as the + engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied). A + label two facts share is left out.""" + from policybench.prompts import INPUT_LABEL_OVERRIDES + from policybench.scenarios import PE_INPUT_RENAMES + + labels: dict[str, str] = {} + shared = set() + for source, label in INPUT_LABEL_OVERRIDES.items(): + key = _label_key(label) + name = PE_INPUT_RENAMES.get(source, source) + if key in labels and labels[key] != name: + shared.add(key) + labels[key] = name + for source, name in PE_INPUT_RENAMES.items(): + labels.setdefault(_label_key(source), name) + return {key: name for key, name in labels.items() if key not in shared} + + +def read_fact_name( + name: str, registry: Mapping[str, VariableInfo], features: frozenset +) -> tuple[str | None, str | None]: + """A fact's engine variable under the labels feature (the label the prompt + gives it) and then the names feature (spaces for underscores). Returns + (engine name, note) or (None, None).""" + if LABELS in features: + labeled = prompt_labels().get(_label_key(name)) + if labeled is not None and labeled in registry: + return labeled, f"read {name!r} as {labeled}" + if NAMES in features: + return resolve_name(name, registry, output_names=False) + return None, None + + +def _set_aliases(v2: dict, people: dict, registry: Mapping[str, VariableInfo]): + """Set each fact the reference also sets under another name there too + (policybench.scenarios.PE_INPUT_ALIASES: the SNAP work rules read stated + usual weekly hours as weekly_hours_worked_before_lsr), unless it is set.""" + from policybench.scenarios import PE_INPUT_ALIASES + + for person, inputs in people.items(): + for source, target in PE_INPUT_ALIASES.items(): + if source in inputs and target in registry and target not in inputs: + inputs[target] = dict(inputs[source]) + v2["notes"].append( + f"people.{person}: also set {target} from {source}, as the " + "benchmark's reference does" + ) + + def resolve_name( name: str, registry: Mapping[str, VariableInfo], output_names: bool ) -> tuple[str | None, str | None]: diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index 5c018fd1..c3a5eb36 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -9,6 +9,7 @@ from __future__ import annotations +import itertools import json import sys from pathlib import Path @@ -37,6 +38,10 @@ "state_code": V("household", "Enum", "state", "year", ("CA", "LA")), "tax_unit_childcare_expenses": V("tax_unit", "float", "childcare", "year"), "marital_unit_flag": V("marital_unit", "bool", "flag", "year"), + "has_esi": V("person", "bool", "ESI", "year"), + "hours_worked_last_week": V("person", "float", "hours", "year"), + "weekly_hours_worked_before_lsr": V("person", "float", "hours", "year"), + "partnership_self_employment_net_earnings": V("person", "float", "p", "year"), } F = pilot @@ -76,6 +81,7 @@ def test_arms_name_their_features(): assert F.arm_features("pe") == frozenset() assert F.arm_features("pe2") == F.V2_FEATURES assert F.arm_features("pe3") == frozenset(F.TOOL_FEATURES) + assert len(F.TOOL_FEATURES) == 7 and F.LABELS in F.arm_features("pe3") assert F.arm_features("pe:flat+names") == frozenset({"flat", "names"}) assert F.arm_features("none") is None and F.arm_features("web") is None assert F.arm_suffix("pe:flat+names") == "+pe-tool[names+flat]" @@ -195,6 +201,82 @@ def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_paramete assert F.brief_explanation({"error": "x"}) == {"error": "x"} +def test_labels_reads_the_facts_as_the_prompt_labels_them(): + situation = _situation() + situation["people"]["head"].update( + { + "gross wages and salaries": 30000, + "Has Employer-Sponsored Insurance": True, + "usual weekly hours worked": 40, + "self-employment partnership income": 5, + } + ) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None, error + head = request["situation"]["people"]["head"] + assert head["employment_income"] == {"2026": 30000} + assert head["has_esi"] == {"2026": True} + assert head["hours_worked_last_week"] == {"2026": 40} + assert head["partnership_self_employment_net_earnings"] == {"2026": 5} + # The reference also sets stated hours where the SNAP work rules read them. + assert head["weekly_hours_worked_before_lsr"] == {"2026": 40} + assert any("also set weekly_hours_worked_before_lsr" in n for n in request["notes"]) + # A name written with spaces is the names feature's, not this one's. + spaced = _situation() + spaced["people"]["head"]["employment income"] = 1 + request, error = _norm({"situation": spaced, "variables": ["snap"]}, F.LABELS) + assert request is None and "unknown variable 'employment income'" in str(error) + + +def test_labels_leaves_hours_already_set_where_the_rules_read_them(): + situation = _situation() + situation["people"]["head"].update( + {"hours_worked_last_week": 40, "weekly_hours_worked_before_lsr": 20} + ) + request, _ = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + head = request["situation"]["people"]["head"] + assert head["weekly_hours_worked_before_lsr"] == {"2026": 20} + + +def test_pe2_reads_no_labels_and_sets_no_alias_as_it_ran(): + situation = _situation() + situation["people"]["head"]["gross wages and salaries"] = 30000 + request, error = _norm( + {"situation": situation, "variables": ["snap"]}, *F.V2_FEATURES + ) + assert request is None and "gross wages and salaries" in str(error) + situation = _situation() + situation["people"]["head"]["hours_worked_last_week"] = 40 + request, _ = _norm({"situation": situation, "variables": ["snap"]}, *F.V2_FEATURES) + assert ( + "weekly_hours_worked_before_lsr" not in request["situation"]["people"]["head"] + ) + + +def test_only_pe2_says_the_facts_carry_variable_names(): + claim = "labeled with its variable names" + assert claim in F.TOOL_NOTICE_V2 + for size in range(1, len(F.TOOL_FEATURES) + 1): + for features in itertools.combinations(F.TOOL_FEATURES, size): + features = frozenset(features) + if F.CONCISE in features and F.EXPLAIN not in features: + continue + notice = F.compose_tool_notice(features, True) + assert (claim in notice) == (features == F.V2_FEATURES), features + assert notice.endswith("as instructed above.") + assert ("names it read" in notice) == bool(features & {F.NAMES, F.LABELS}) + + +def test_every_prompt_label_reads_back_to_its_engine_name(): + from policybench.prompts import INPUT_LABEL_OVERRIDES + from policybench.scenarios import PE_INPUT_RENAMES + + labels = F.prompt_labels() + for field, label in INPUT_LABEL_OVERRIDES.items(): + assert labels[F._label_key(label)] == PE_INPUT_RENAMES.get(field, field) + assert labels[F._label_key("roth 401k contributions")] == "roth_401k_contributions" + + def _flat(**overrides): household = { "state": "LA", @@ -375,3 +457,96 @@ def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): assert values[0] == values[1], (scenario.id, variable, values) checked += 1 assert checked > 80 + + +@pytest.mark.slow +def test_a_household_described_under_the_prompts_labels_computes_the_same(): + """Differential: each frozen-board household described flat with every + fact as its prompt writes it (a label in words, or the name with spaces), + and stated hours only where the prompt states them, computes every board + output as the benchmark does. Take-up conventions, which the prompt does + not list, keep their engine names.""" + import household_set as hs + import pandas as pd + from policyengine_us import Simulation + + from policybench.ground_truth import _extract_person_value, _pe_variable_for_output + from policybench.prompts import _humanize_input_label + from policybench.scenarios import PE_INPUT_ALIASES, PE_INPUT_RENAMES + + system = F.build_tool_system(reform=True) + registry = F.build_registry(system) + reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") + variables = reference.groupby("scenario_id").variable.agg(list) + field_of = {name: field for field, name in PE_INPUT_RENAMES.items()} + dropped = {"is_tax_unit_head", "is_tax_unit_spouse", *PE_INPUT_ALIASES.values()} + + worded = [] + + def labeled(facts): + out = {} + for k, v in facts.items(): + if k in dropped: + continue + if k.startswith(F.CONVENTION_PREFIXES): + out[k] = v + continue + label = _humanize_input_label(field_of.get(k, k)) + if label != k.replace("_", " "): + worded.append(label) + out[label] = v + return out + + checked = 0 + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:12]: + situation = hs.situation_for(scenario) + couple = scenario.marital_couple() or () + head = couple[0] if couple else scenario.adults[0].name + people = [ + { + "name": name, + "role": "head" + if name == head + else ("spouse" if name in couple else "dependent"), + "facts": labeled(facts), + } + for name, facts in situation["people"].items() + ] + household_facts = {} + for group in ("tax_units", "spm_units", "families", "households"): + for unit in situation[group].values(): + household_facts.update( + labeled( + { + k: v + for k, v in unit.items() + if k not in ("members", "state_code") + } + ) + ) + household = { + "state": scenario.state, + "people": people, + "household_facts": household_facts, + } + request, error = F.normalize_request( + {"household": json.dumps(household), "variables": ["snap"]}, + registry, + features=frozenset({F.FLAT, F.LABELS, F.NAMES}), + ) + assert error is None, (scenario.id, error) + own = Simulation(tax_benefit_system=system, situation=situation) + flat = Simulation(tax_benefit_system=system, situation=request["situation"]) + for variable in variables[scenario.id]: + name = _pe_variable_for_output(variable, "us") + values = [ + float( + _extract_person_value( + simulation.calculate(name, 2026), scenario, variable + ) + ) + for simulation in (own, flat) + ] + assert values[0] == values[1], (scenario.id, variable, values) + checked += 1 + assert checked > 120 and len(worded) > 10, (checked, worded) From fd3f5d096562754da767819bf6334a84449acf73 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:31:44 -0400 Subject: [PATCH 3/9] Labels: read the v2-pilot relabelings; flat: names optional The v2-pilot prompt relabels two facts itself ("home energy assistance received" for spm_unit_energy_subsidy, "employer-paid health insurance premiums ..." for employer_sponsored_insurance_premiums); the labels feature now reads them, derived from V2_PILOT_CONVENTIONS so the two stay in step. In the flat format a person given no name is named by role (head, spouse, dependent_1, ...): Flash-Lite's first flat call in two smoke households left the names out. The label differential now writes each fact as the v2-pilot prompt does. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 46 +++++++++++++++++++++++++++---------- tests/test_tool_v3.py | 44 +++++++++++++++++++++++++++++++++-- 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 58ec78bb..775725c9 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -664,9 +664,10 @@ def refuse_policyengine(label: str, value) -> None: FLAT_DESCRIPTION = ( "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " "the requested variables. `household` is a JSON object: `state` (the " - "two-letter state code), `people` (a list, each with `name`, `role` -- " - 'head, spouse or dependent -- and `facts`, for example {"age": 40, ' - '"employment income": 30000}), and `household_facts` (facts about the ' + "two-letter state code), `people` (a list, each with `role` -- head, " + 'spouse or dependent --, `facts`, for example {"age": 40, "employment ' + 'income": 30000}, and optionally `name`, by default the role, numbered for ' + "dependents), and `household_facts` (facts about the " "household as a whole: its tax unit, family, SPM unit or home). The head " "and spouse are one couple and everyone else is on their own; everyone " "shares one tax unit, family, SPM unit and household. A fact is a value " @@ -1248,7 +1249,8 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): reference situations have (policybench.scenarios.Scenario.to_pe_household): the head and spouse one marital unit and everyone else alone, and one tax unit, SPM unit, family and household holding everyone. A household fact - goes to the unit of its variable's entity. Returns None after recording + goes to the unit of its variable's entity. A person given no name is named + by role (head, spouse, dependent_1, ...). Returns None after recording problems.""" if isinstance(household, str): if len(household) > MAX_SITUATION_CHARS: @@ -1279,6 +1281,7 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None situation: dict[str, Any] = {"people": {}} roles: dict[str, str] = {} + dependents = 0 for index, person in enumerate(people): if not isinstance(person, dict): problems.append( @@ -1290,17 +1293,23 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): person.get("role"), person.get("facts", {}), ) + if role not in FLAT_ROLES: + problems.append( + f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" + ) + continue + if role == "dependent": + dependents += 1 + if name is None: + name = role if role != "dependent" else f"dependent_{dependents}" + if v2 is not None: + v2["notes"].append(f"people[{index}]: named {name}") if not isinstance(name, str) or not name or len(name) > 64: problems.append(f"people[{index}].name must be a short string") continue if name in roles: problems.append(f"people[{index}]: {name!r} is listed twice") continue - if role not in FLAT_ROLES: - problems.append( - f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" - ) - continue if not isinstance(facts, dict): problems.append(f"people[{index}].facts must be an object") continue @@ -1404,14 +1413,15 @@ def _label_key(text: str) -> str: @functools.lru_cache(maxsize=1) def prompt_labels() -> dict[str, str]: """The labels the prompts give facts in words, keyed by _label_key, as the - engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied). A - label two facts share is left out.""" + engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied): + the board's labels and the v2-pilot prompt's relabelings. A label two facts + share is left out.""" from policybench.prompts import INPUT_LABEL_OVERRIDES from policybench.scenarios import PE_INPUT_RENAMES labels: dict[str, str] = {} shared = set() - for source, label in INPUT_LABEL_OVERRIDES.items(): + for source, label in {**INPUT_LABEL_OVERRIDES, **v2_pilot_labels()}.items(): key = _label_key(label) name = PE_INPUT_RENAMES.get(source, source) if key in labels and labels[key] != name: @@ -1422,6 +1432,18 @@ def prompt_labels() -> dict[str, str]: return {key: name for key, name in labels.items() if key not in shared} +def v2_pilot_labels() -> dict[str, str]: + """The facts the v2-pilot prompt relabels (a V2_PILOT_CONVENTIONS rewrite + of ": " to ": "), by input name.""" + relabeled = {} + for pattern, replacement in V2_PILOT_CONVENTIONS: + old = pattern.pattern + if pattern.groups or not old.endswith(": ") or not replacement.endswith(": "): + continue + relabeled[old[:-2].replace(" ", "_")] = replacement[:-2] + return relabeled + + def read_fact_name( name: str, registry: Mapping[str, VariableInfo], features: frozenset ) -> tuple[str | None, str | None]: diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index c3a5eb36..d8903c93 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -274,6 +274,16 @@ def test_every_prompt_label_reads_back_to_its_engine_name(): labels = F.prompt_labels() for field, label in INPUT_LABEL_OVERRIDES.items(): assert labels[F._label_key(label)] == PE_INPUT_RENAMES.get(field, field) + # The v2-pilot prompt's own relabelings, as that prompt writes them. + relabeled = F.v2_pilot_labels() + assert set(relabeled) == { + "employer_sponsored_insurance_premiums", + "spm_unit_energy_subsidy", + } + for field, label in relabeled.items(): + written = F.v2_pilot_text(f"- {field.replace('_', ' ')}: $5") + assert written == f"- {label}: $5" + assert labels[F._label_key(label)] == field assert labels[F._label_key("roth 401k contributions")] == "roth_401k_contributions" @@ -324,6 +334,33 @@ def test_flat_builds_the_benchmarks_units(): assert s["people"]["Kid"]["is_tax_unit_head"] == {"2026": False} +def test_flat_names_a_person_by_role_when_no_name_is_given(): + people = [ + {"role": "head", "facts": {"age": 40}}, + {"role": "dependent", "facts": {"age": 6}}, + {"name": "Kid", "role": "dependent", "facts": {"age": 4}}, + {"role": "dependent", "facts": {"age": 2}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert error is None, error + assert list(request["situation"]["people"]) == [ + "head", + "dependent_1", + "Kid", + "dependent_3", + ] + assert "people[0]: named head" in request["notes"] + people = [ + {"role": "head", "facts": {}}, + {"name": "head", "role": "dependent", "facts": {}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("listed twice" in p for p in error["problems"]) + people = [{"name": "", "role": "head", "facts": {}}] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("name must be a short string" in p for p in error["problems"]) + + def test_flat_refuses_what_it_cannot_place(): people = [ {"name": "A", "role": "head", "facts": {}}, @@ -462,7 +499,8 @@ def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): @pytest.mark.slow def test_a_household_described_under_the_prompts_labels_computes_the_same(): """Differential: each frozen-board household described flat with every - fact as its prompt writes it (a label in words, or the name with spaces), + fact as its v2-pilot prompt writes it (a label in words, or the name with + spaces), and stated hours only where the prompt states them, computes every board output as the benchmark does. Take-up conventions, which the prompt does not list, keep their engine names.""" @@ -491,7 +529,9 @@ def labeled(facts): if k.startswith(F.CONVENTION_PREFIXES): out[k] = v continue - label = _humanize_input_label(field_of.get(k, k)) + label = F.v2_pilot_text(f"- {_humanize_input_label(field_of.get(k, k))}: ")[ + 2:-2 + ] if label != k.replace("_", " "): worded.append(label) out[label] = v From 5c0f3ec5843b4a839b8255f9836a359e4d6f3c80 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:39:43 -0400 Subject: [PATCH 4/9] Flat: a household of one's person fact goes to that person Two of three Flash-Lite smoke households gave the head's rent as a household fact. pre_subsidy_rent is a person variable in policyengine-us, so the flat format refused it. Under the placement feature a person fact given for a household of one now goes to that person, as placement moves a person variable off a one-person unit; with two or more people it is still refused, now saying to give it under the person whose facts list it. The test registry's entities now match the engine's (pre_subsidy_rent is a person variable). Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 18 ++++++++++++++++- tests/test_tool_v3.py | 39 +++++++++++++++++++++++++++++-------- 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 775725c9..e3e3f1b4 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1366,9 +1366,25 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): continue entity = registry[resolved].entity if entity == "person": + # As placement moves a person variable off a one-person unit: only + # a household of one says whose fact it is. + if v2 is not None and PLACEMENT in v2["features"] and len(everyone) == 1: + facts_of = situation["people"][everyone[0]] + if resolved in facts_of and facts_of[resolved] != value: + problems.append( + f"household_facts: {resolved} is also set for " + f"{everyone[0]} with a different value" + ) + continue + facts_of[resolved] = value + v2["notes"].append( + f"household_facts: moved {resolved}, a person's fact, to " + f"{everyone[0]}" + ) + continue problems.append( f"household_facts: {resolved} is a person's fact; give it under " - "that person" + "the person whose facts list it" ) continue unit_of[entity][resolved] = value diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index d8903c93..be60efdf 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -31,7 +31,8 @@ ), "traditional_401k_contributions_desired": V("person", "float", "asked", "year"), "snap": V("spm_unit", "float", "SNAP", "month", computed=True), - "pre_subsidy_rent": V("spm_unit", "float", "rent", "year"), + "pre_subsidy_rent": V("person", "float", "rent", "year"), + "spm_unit_energy_subsidy": V("spm_unit", "float", "energy", "year"), "qualified_passenger_vehicle_loan_interest": V( "household", "float", "car loan", "year" ), @@ -300,7 +301,7 @@ def _flat(**overrides): {"name": "Kid", "role": "dependent", "facts": {"age": 6}}, ], "household_facts": { - "pre_subsidy_rent": 9000, + "spm_unit_energy_subsidy": 900, "tax_unit_childcare_expenses": 1200, "qualified_passenger_vehicle_loan_interest": 300, }, @@ -324,7 +325,7 @@ def test_flat_builds_the_benchmarks_units(): ): assert s[group][unit]["members"] == ["Head", "Spouse", "Kid"] assert s["households"]["household"]["state_code"] == {"2026": "LA"} - assert s["spm_units"]["spm_unit"]["pre_subsidy_rent"] == {"2026": 9000} + assert s["spm_units"]["spm_unit"]["spm_unit_energy_subsidy"] == {"2026": 900} assert s["tax_units"]["tax_unit"]["tax_unit_childcare_expenses"] == {"2026": 1200} assert s["households"]["household"][ "qualified_passenger_vehicle_loan_interest" @@ -361,6 +362,29 @@ def test_flat_names_a_person_by_role_when_no_name_is_given(): assert any("name must be a short string" in p for p in error["problems"]) +def test_flat_placement_gives_a_household_of_ones_person_fact_to_that_person(): + one = [{"name": "Solo", "role": "head", "facts": {"age": 70}}] + args = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, error + solo = request["situation"]["people"]["Solo"] + assert solo["pre_subsidy_rent"] == {"2026": 9000} + assert "pre_subsidy_rent" not in request["situation"]["spm_units"]["spm_unit"] + assert any("moved pre_subsidy_rent" in n for n in request["notes"]) + one[0]["facts"]["pre_subsidy_rent"] = 100 + conflict = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(conflict, F.FLAT, F.PLACEMENT) + assert any("different value" in p for p in error["problems"]) + # Two people: whose rent it is is not said, so it is refused. + request, error = _norm( + _flat(household_facts={"pre_subsidy_rent": 9000}), F.FLAT, F.PLACEMENT + ) + assert any("whose facts list it" in p for p in error["problems"]) + # Without placement, even a household of one is refused. + request, error = _norm(args, F.FLAT) + assert any("whose facts list it" in p for p in error["problems"]) + + def test_flat_refuses_what_it_cannot_place(): people = [ {"name": "A", "role": "head", "facts": {}}, @@ -381,13 +405,12 @@ def test_flat_refuses_what_it_cannot_place(): def test_flat_with_names_reads_labeled_household_facts(): request, error = _norm( - _flat(household_facts={"pre subsidy rent": 9000}), F.FLAT, F.NAMES + _flat(household_facts={"spm unit energy subsidy": 900}), F.FLAT, F.NAMES ) assert error is None - assert request["situation"]["spm_units"]["spm_unit"]["pre_subsidy_rent"] == { - "2026": 9000 - } - assert any("as pre_subsidy_rent" in n for n in request["notes"]) + spm_unit = request["situation"]["spm_units"]["spm_unit"] + assert spm_unit["spm_unit_energy_subsidy"] == {"2026": 900} + assert any("as spm_unit_energy_subsidy" in n for n in request["notes"]) def test_the_schema_offers_each_features_arguments(): From 9c43f4f7ea0e576e8686af00b962a1896fe458a3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:43:24 -0400 Subject: [PATCH 5/9] Keep tool version 2's request as it was: detail only under concise A differential against 5bbd571a (the commit the test batches ran) found version 2 requests carrying detail "full"; only execute_request reads it, through .get, so results were unchanged, but the request now matches exactly. 4,000 generated calls through versions 1 and 2, the notices, descriptions, schemas and 15 model profiles are identical to 5bbd571a. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 3 ++- tests/test_tool_v3.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index e3e3f1b4..2712b046 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1224,7 +1224,7 @@ def fail() -> tuple[None, dict]: elif name not in explain: explain.append(name) - detail = "full" + detail = None if CONCISE in features: detail = args.get("detail") or "brief" if detail not in ("brief", "full"): @@ -1236,6 +1236,7 @@ def fail() -> tuple[None, dict]: if v2 is not None: request["explain"] = explain request["notes"] = v2["notes"] + if detail is not None: request["detail"] = detail return request, None diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index be60efdf..e0064022 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -179,7 +179,7 @@ def test_concise_takes_a_detail_and_defaults_to_brief(): {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, F.EXPLAIN, ) - assert request["detail"] == "full" + assert "detail" not in request # every step, as tool version 2 explains def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_parameters(): From 6478e8f40f01b2253245f23a8e872d4a11784642 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 18:35:53 -0400 Subject: [PATCH 6/9] Tool v3 review r1: frame, resume, labels, flat and text fixes GPT-6.1 Sol's review (REQUEST_CHANGES, 7 findings) at 9c43f4f7: 1. A model's arguments could turn on features its arm does not offer: the worker read "__tool_features__" (and, as at 5bbd571a, "__tool_version__") from the arguments dict. A pool now sends features in a tuple frame no JSON can be, the worker validates them (names, concise beside explain), and keys like "__tool_version__" in a model's arguments are refused as unknown arguments. No run record ever carried them. 2. A resume could mix treatments. Each household record now keeps its tool features and calculator schema hash; a resume compares features, notice and schema, and refuses a non-legacy record without them. 3. Labels missed what the prompt renders. The label index is now built from the prompt's own renderer (_format_input_line, then v2_pilot_text) for every promptable, labeled and reference-stated input: fallback labels, the Marketplace plan line, "state", partnership or S-corp income. Facts the reference drops (PE_DROPPED_INPUTS) are ignored with a note. A label that names two of the engine's inputs is not read. 4. Flat refuses person keys other than name, role and facts. 5. Flat stores household facts under the name written, so normalization compares them once normalized; roles set the relationship inputs and a stated one that contradicts its role is refused. 6. The flat example uses employment_income; placement replaces the wrong-entity error sentence; "full" detail is the trace, capped. 7. Diagnostics read flat people by their default names. Evidence: 210 passed, 4 frozen-engine skips. A labels-only differential builds 30 frozen households from their prompts' facts and matches the reference situation on every output, and so do all 100 frozen and all 20 test-set households in a one-off run, with or without take-up inputs. pe and pe2 are identical to 5bbd571a over 4,000 calls and 2,000 pool messages. 11 mutants of the fixes are caught. An end-to-end Flash-Lite pe3 smoke run had 0 tool errors on 3 households. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 454 +++++++++++++++++++++++++++--------- tests/test_tool_v2.py | 7 +- tests/test_tool_v3.py | 354 ++++++++++++++++++++++------ 3 files changed, 629 insertions(+), 186 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 2712b046..5f586522 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -311,8 +311,8 @@ def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: if EXPLAIN in features: explain = ( "ask it to `explain` up to three variables to see in brief how it " - "computed them, with the parameters and laws it used, or every step " - 'with `detail` set to "full"' + "computed them, with the parameters and laws it used, or the fuller " + 'trace with `detail` set to "full"' if CONCISE in features else "ask it to `explain` up to three variables to see how it computed " "them, with the parameters and laws it used" @@ -665,9 +665,9 @@ def refuse_policyengine(label: str, value) -> None: "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " "the requested variables. `household` is a JSON object: `state` (the " "two-letter state code), `people` (a list, each with `role` -- head, " - 'spouse or dependent --, `facts`, for example {"age": 40, "employment ' - 'income": 30000}, and optionally `name`, by default the role, numbered for ' - "dependents), and `household_facts` (facts about the " + 'spouse or dependent --, `facts`, for example {"age": 40, ' + '"employment_income": 30000}, and optionally `name`, by default the role, ' + "numbered for dependents), and `household_facts` (facts about the " "household as a whole: its tax unit, family, SPM unit or home). The head " "and spouse are one couple and everyone else is on their own; everyone " "shares one tax unit, family, SPM unit and household. A fact is a value " @@ -676,7 +676,13 @@ def refuse_policyengine(label: str, value) -> None: 'lists the variables to compute for `period` (default "2026"). Each ' "result gives the variable's entity, label and value for each member of " "that entity; a monthly variable computed for a year is summed over its " - "months." + "months. An unknown name or a malformed household returns an error with " + "suggestions instead of results." +) +# The version 1 description's last sentence, which placement amends. +_ENTITY_ERROR_SENTENCE = ( + "An unknown name, a variable set on the wrong entity, or a malformed " + "situation returns an error with suggestions instead of results." ) @@ -693,6 +699,12 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - description = ( FLAT_DESCRIPTION if FLAT in features else CALCULATE_DESCRIPTION ) + if PLACEMENT in features and FLAT not in features: + description = description.replace( + _ENTITY_ERROR_SENTENCE, + "An unknown name or a malformed situation returns an error " + "with suggestions instead of results.", + ) notes = [] if LABELS in features: notes.append( @@ -706,8 +718,9 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - ) if PLACEMENT in features and FLAT not in features: notes.append( - "An input set on the wrong entity is moved to the one unit it " - "belongs to." + "An input set on the wrong entity is moved to the one unit " + "holding its people; where no single unit holds them, it " + "returns an error." ) if OVERRIDES in features: notes.append( @@ -744,7 +757,8 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - "description": ( 'How much an explanation shows: "brief" (the default: the ' 'value, what it came from, and the main parameters) or "full" ' - "(every step)." + f"(the trace, up to {EXPLAIN_MAX_LINES} steps and " + f"{EXPLAIN_MAX_PARAMETERS} parameters)." ), } return schema @@ -821,13 +835,9 @@ class VariableInfo: def stated_inputs() -> frozenset[str]: """Every variable a household fact the benchmark states can set.""" - from policybench.scenarios import ( - PE_INPUT_ALIASES, - PE_INPUT_RENAMES, - get_promptable_input_specs, - ) + from policybench.scenarios import PE_INPUT_ALIASES, PE_INPUT_RENAMES - names = {spec.source_name for spec in get_promptable_input_specs()} + names = {spec.source_name for spec in _promptable_specs()} names |= BASE_STATED_INPUTS names |= set(PE_INPUT_ALIASES.values()) | set(PE_INPUT_RENAMES.values()) return frozenset(names) @@ -961,6 +971,9 @@ def _normalize_inputs( if resolved is not None: v2["notes"].append(f"{where}: {note}") name = resolved + elif (field := dropped_fact(name, registry, v2["features"])) is not None: + v2["notes"].append(_dropped_note(where, name, field)) + continue if not isinstance(name, str) or name not in registry: problems.append( f"{where}: unknown variable {name!r}." @@ -1068,7 +1081,7 @@ def fail() -> tuple[None, dict]: if FLAT in features: args = dict(args) args["situation"] = flat_situation( - args.get("household"), registry, problems, v2 + args.get("household"), registry, problems, v2, period ) if args["situation"] is None: return fail() @@ -1242,17 +1255,67 @@ def fail() -> tuple[None, dict]: FLAT_ROLES = ("head", "spouse", "dependent") +FLAT_PERSON_KEYS = ("name", "role", "facts") +# The relationship inputs a role sets (as Scenario.to_pe_household's people +# carry them). +FLAT_RELATIONS = {"is_tax_unit_head": "head", "is_tax_unit_spouse": "spouse"} + + +def flat_names(people) -> list: + """The name each listed person goes by: the one given, or, for a person + given none, the role (head, spouse, dependent_1, ...: dependents are + numbered by their place among the dependents listed).""" + names, dependents = [], 0 + for person in people if isinstance(people, list) else []: + if not isinstance(person, dict): + names.append(None) + continue + role = person.get("role") + if role == "dependent": + dependents += 1 + name = person.get("name") + if name is None and role in FLAT_ROLES: + name = role if role != "dependent" else f"dependent_{dependents}" + names.append(name) + return names -def flat_situation(household, registry, problems: list[str], v2: dict | None): +def _same_value(registry, name: str, a, b, period: str) -> bool | None: + """Whether two values of an input are the same once normalized (None when + either is malformed; normalization reports it).""" + scratch: list[str] = [] + info = registry[name] + first = _normalize_value("", name, a, info, period, scratch) + second = _normalize_value("", name, b, info, period, scratch) + return None if scratch else first == second + + +def flat_situation( + household, + registry, + problems: list[str], + v2: dict | None, + period: str = str(YEAR), +): """The flat feature's household -- state, people with role and facts, and household facts -- as a situation, with the units the benchmark's own reference situations have (policybench.scenarios.Scenario.to_pe_household): the head and spouse one marital unit and everyone else alone, and one tax - unit, SPM unit, family and household holding everyone. A household fact - goes to the unit of its variable's entity. A person given no name is named - by role (head, spouse, dependent_1, ...). Returns None after recording - problems.""" + unit, SPM unit, family and household holding everyone. Roles set the + relationship inputs; a stated one that contradicts its role is refused. + A household fact goes to the unit of its variable's entity, under the + name it was written with, so normalization reads and compares it as any + other input. A person given no name is named by role (flat_names). + Returns None after recording problems.""" + features = v2["features"] if v2 is not None else frozenset() + + def engine_name(name): + if isinstance(name, str) and name in registry: + return name + if isinstance(name, str): + return read_fact_name(name, registry, features)[0] + return None + if isinstance(household, str): if len(household) > MAX_SITUATION_CHARS: problems.append("household is too large") @@ -1282,29 +1345,27 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None situation: dict[str, Any] = {"people": {}} roles: dict[str, str] = {} - dependents = 0 - for index, person in enumerate(people): + for index, (person, name) in enumerate(zip(people, flat_names(people))): if not isinstance(person, dict): problems.append( f"people[{index}] must be an object with name, role and facts" ) continue - name, role, facts = ( - person.get("name"), - person.get("role"), - person.get("facts", {}), - ) + stray = sorted(map(str, set(person) - set(FLAT_PERSON_KEYS))) + if stray: + problems.append( + f"people[{index}] takes only name, role and facts; give " + f"{', '.join(stray)} under facts" + ) + continue + role, facts = person.get("role"), person.get("facts", {}) if role not in FLAT_ROLES: problems.append( f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" ) continue - if role == "dependent": - dependents += 1 - if name is None: - name = role if role != "dependent" else f"dependent_{dependents}" - if v2 is not None: - v2["notes"].append(f"people[{index}]: named {name}") + if person.get("name") is None and v2 is not None: + v2["notes"].append(f"people[{index}]: named {name}") if not isinstance(name, str) or not name or len(name) > 64: problems.append(f"people[{index}].name must be a short string") continue @@ -1326,8 +1387,21 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None for name, role in roles.items(): facts = situation["people"][name] - facts.setdefault("is_tax_unit_head", role == "head") - facts.setdefault("is_tax_unit_spouse", role == "spouse") + for relation, relation_role in FLAT_RELATIONS.items(): + if relation not in registry: + continue + implied = role == relation_role + stated = [k for k in facts if engine_name(k) == relation] + if not stated: + facts[relation] = implied + elif any( + _same_value(registry, relation, facts[k], implied, period) is False + for k in stated + ): + problems.append( + f"people.{name}: {relation} contradicts the role {role}; the " + "role sets it" + ) everyone = list(situation["people"]) couple = heads + spouses groups = [couple] + [[n] for n in everyone if n not in couple] @@ -1354,12 +1428,16 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): problems.append("household_facts must be an object") return None for name, value in facts.items(): - resolved = name - if isinstance(name, str) and name not in registry and v2 is not None: - resolved, note = read_fact_name(name, registry, v2["features"]) - if resolved is not None: - v2["notes"].append(f"household_facts: {note}") - if not isinstance(resolved, str) or resolved not in registry: + resolved = engine_name(name) + if resolved is None: + field = ( + dropped_fact(name, registry, features) + if isinstance(name, str) + else None + ) + if field is not None: + v2["notes"].append(_dropped_note("household_facts", name, field)) + continue problems.append( f"household_facts: unknown variable {name!r}." + (_suggest(name, registry) if isinstance(name, str) else "") @@ -1369,26 +1447,25 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): if entity == "person": # As placement moves a person variable off a one-person unit: only # a household of one says whose fact it is. - if v2 is not None and PLACEMENT in v2["features"] and len(everyone) == 1: - facts_of = situation["people"][everyone[0]] - if resolved in facts_of and facts_of[resolved] != value: - problems.append( - f"household_facts: {resolved} is also set for " - f"{everyone[0]} with a different value" - ) - continue - facts_of[resolved] = value - v2["notes"].append( - f"household_facts: moved {resolved}, a person's fact, to " - f"{everyone[0]}" + if v2 is None or PLACEMENT not in features or len(everyone) != 1: + problems.append( + f"household_facts: {resolved} is a person's fact; give it " + "under the person whose facts list it" ) continue - problems.append( - f"household_facts: {resolved} is a person's fact; give it under " - "the person whose facts list it" + target = situation["people"][everyone[0]] + v2["notes"].append( + f"household_facts: moved {resolved}, a person's fact, to {everyone[0]}" ) + else: + target = unit_of[entity] + if name in target: + if _same_value(registry, resolved, target[name], value, period) is False: + problems.append( + f"household_facts: {resolved} is set twice with different values" + ) continue - unit_of[entity][resolved] = value + target[name] = value return None if problems else situation @@ -1427,55 +1504,123 @@ def _label_key(text: str) -> str: return re.sub(r"[^0-9a-z]", "", text.casefold()) +# The prompt's household block gives the state as "state: CT". +PROMPT_EXTRA_LABELS = {"state": "state_code"} + + +@functools.lru_cache(maxsize=1) +def _promptable_specs(): + from policybench.scenarios import get_promptable_input_specs + + return get_promptable_input_specs() + + +def _sample_value(value_type: str): + if value_type == "bool": + return True + if value_type in ("str", "Enum"): + return "x" + return 1.0 + + +def rendered_labels(field: str, value_type: str = "float") -> tuple[str, ...]: + """The labels a fact is shown under: in the board's prompt + (policybench.prompts._format_input_line, special renderers included) and + in the v2-pilot prompt (v2_pilot_text applied to that line).""" + from policybench.prompts import _format_input_line, _humanize_input_label + + line = _format_input_line(field, _sample_value(value_type)) + label = _humanize_input_label(field) + if not (line == f"- {label}" or line.startswith(f"- {label}: ")): + label = line[2:].split(": ", 1)[0] + labels = [label] + suffix = line[2 + len(label) :] + rewritten = v2_pilot_text(line) + if rewritten != line and rewritten.startswith("- ") and rewritten.endswith(suffix): + labels.append(rewritten[2 : len(rewritten) - len(suffix)]) + return tuple(labels) + + @functools.lru_cache(maxsize=1) -def prompt_labels() -> dict[str, str]: - """The labels the prompts give facts in words, keyed by _label_key, as the - engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied): - the board's labels and the v2-pilot prompt's relabelings. A label two facts - share is left out.""" +def prompt_labels() -> dict[str, frozenset[tuple[str, bool]]]: + """Every label the prompts show a fact under, keyed by _label_key, each + with the facts it may name as (engine name, dropped): the promptable + inputs, every labeled input (policybench.prompts.INPUT_LABEL_OVERRIDES), + the inputs the reference sets besides them (BASE_STATED_INPUTS), inputs + the engine has renamed and inputs the reference leaves out + (policybench.scenarios.PE_INPUT_RENAMES, PE_DROPPED_INPUTS), each in the + board's and the v2-pilot prompt's wording. Dropped facts are shown but are + not engine inputs. A label may name facts of another country; a lookup + reads only the one its engine has.""" from policybench.prompts import INPUT_LABEL_OVERRIDES - from policybench.scenarios import PE_INPUT_RENAMES - - labels: dict[str, str] = {} - shared = set() - for source, label in {**INPUT_LABEL_OVERRIDES, **v2_pilot_labels()}.items(): - key = _label_key(label) - name = PE_INPUT_RENAMES.get(source, source) - if key in labels and labels[key] != name: - shared.add(key) - labels[key] = name - for source, name in PE_INPUT_RENAMES.items(): - labels.setdefault(_label_key(source), name) - return {key: name for key, name in labels.items() if key not in shared} - - -def v2_pilot_labels() -> dict[str, str]: - """The facts the v2-pilot prompt relabels (a V2_PILOT_CONVENTIONS rewrite - of ": " to ": "), by input name.""" - relabeled = {} - for pattern, replacement in V2_PILOT_CONVENTIONS: - old = pattern.pattern - if pattern.groups or not old.endswith(": ") or not replacement.endswith(": "): - continue - relabeled[old[:-2].replace(" ", "_")] = replacement[:-2] - return relabeled + from policybench.scenarios import PE_DROPPED_INPUTS, PE_INPUT_RENAMES + + fields = {spec.output_name: spec.value_type for spec in _promptable_specs()} + dropped = {f for names in PE_DROPPED_INPUTS.values() for f in names} + for source in ( + "age", + *BASE_STATED_INPUTS, + *INPUT_LABEL_OVERRIDES, + *PE_INPUT_RENAMES, + *dropped, + ): + fields.setdefault(source, "int" if source == "age" else "float") + labels: dict[str, set[tuple[str, bool]]] = {} + for source, value_type in fields.items(): + target = (PE_INPUT_RENAMES.get(source, source), source in dropped) + keys = {_label_key(label) for label in rendered_labels(source, value_type)} + if source in PE_INPUT_RENAMES or source in dropped: + keys.add(_label_key(source)) + for key in keys: + labels.setdefault(key, set()).add(target) + for label, name in PROMPT_EXTRA_LABELS.items(): + labels.setdefault(_label_key(label), set()).add((name, False)) + return {key: frozenset(targets) for key, targets in labels.items()} def read_fact_name( name: str, registry: Mapping[str, VariableInfo], features: frozenset ) -> tuple[str | None, str | None]: - """A fact's engine variable under the labels feature (the label the prompt - gives it) and then the names feature (spaces for underscores). Returns - (engine name, note) or (None, None).""" + """A fact's engine variable under the labels feature (any label the + prompt shows it under, when exactly one fact it may name is the engine's) + and then the names feature (spaces for underscores). Returns (engine + name, note) or (None, None).""" if LABELS in features: - labeled = prompt_labels().get(_label_key(name)) - if labeled is not None and labeled in registry: - return labeled, f"read {name!r} as {labeled}" + found = [ + engine + for engine, dropped in prompt_labels().get(_label_key(name), ()) + if not dropped and engine in registry + ] + if len(found) == 1: + return found[0], f"read {name!r} as {found[0]}" if NAMES in features: return resolve_name(name, registry, output_names=False) return None, None +def dropped_fact( + name: str, registry: Mapping[str, VariableInfo], features: frozenset +) -> str | None: + """Under the labels feature, the input a fact names when the prompt shows + it but the reference leaves it out (policybench.scenarios.PE_DROPPED_INPUTS: + no longer an engine input); otherwise None.""" + if LABELS not in features: + return None + found = [ + engine + for engine, dropped in prompt_labels().get(_label_key(name), ()) + if dropped + ] + return found[0] if len(found) == 1 else None + + +def _dropped_note(where: str, name: str, field: str) -> str: + return ( + f"{where}: {name!r} ({field}) is not an input of this engine; the " + "benchmark's reference leaves it out too, so it was ignored" + ) + + def _set_aliases(v2: dict, people: dict, registry: Mapping[str, VariableInfo]): """Set each fact the reference also sets under another name there too (policybench.scenarios.PE_INPUT_ALIASES: the SNAP work rules read stated @@ -1643,7 +1788,7 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic def brief_explanation(full: dict) -> dict: """The concise feature's explanation: the value, the variables it came - from directly (one level down, zero-valued ones left out), and the first + from directly (one level down, the first BRIEF_INPUTS), and the first parameters it read with their sources.""" if "error" in full: return full @@ -1655,7 +1800,7 @@ def brief_explanation(full: dict) -> dict: "value": steps[0].strip() if steps else "", "from": direct[:BRIEF_INPUTS], "parameters": full["parameters"][:BRIEF_PARAMETERS], - "more": 'set detail to "full" for every step', + "more": 'set detail to "full" for the trace', } @@ -1843,6 +1988,46 @@ def build_tool_system( return CountryTaxBenefitSystem(reform=load_fix_reform(path, snapshot)) +# How a pool tells its worker the calculator's features: a tuple, which a +# model's JSON arguments can never be, so a model cannot turn on a feature its +# arm does not offer. A pool without features sends the arguments as they are. +CALL_FRAME = "calculate" + + +def tool_call_frame(features: frozenset, arguments) -> tuple: + return (CALL_FRAME, tuple(sorted(features)), arguments) + + +def frame_features(names) -> frozenset: + """A frame's features, refused (ValueError) unless they are a tuple of + known feature names with concise only beside explain.""" + if not isinstance(names, tuple) or not all(isinstance(n, str) for n in names): + raise ValueError("the features must be a tuple of names") + features = frozenset(names) + if features - set(TOOL_FEATURES): + raise ValueError( + f"unknown tool features {sorted(features - set(TOOL_FEATURES))}" + ) + if CONCISE in features and EXPLAIN not in features: + raise ValueError("the concise feature needs explain") + return features + + +def worker_request(message, registry) -> tuple[dict | None, dict | None]: + """A worker's reading of one pool message: a call frame's arguments under + its features, or anything else as a model's arguments with no features + (keys such as "__tool_features__" or "__tool_version__" in them are the + model's, and refused as unknown arguments).""" + features = frozenset() + if isinstance(message, tuple) and len(message) == 3 and message[0] == CALL_FRAME: + try: + features = frame_features(message[1]) + except ValueError as error: + return None, {"error": "configuration", "message": str(error)} + message = message[2] + return normalize_request(message, registry, features=features) + + def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> None: for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"): os.environ[name] = "1" @@ -1857,13 +2042,7 @@ def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> return if args is None: return - features = frozenset() - if isinstance(args, dict) and "__tool_features__" in args: - features = frozenset(args.get("__tool_features__") or ()) - args = args.get("arguments") - elif isinstance(args, dict) and args.get("__tool_version__") == 2: - features, args = V2_FEATURES, args.get("arguments") - request, error = normalize_request(args, registry, features=features) + request, error = worker_request(args, registry) if error is not None: conn.send(("ok", error)) continue @@ -1958,13 +2137,8 @@ def call(self, args) -> tuple[dict, float]: time.perf_counter() - started, ) process, conn = worker - if self._tool_features == V2_FEATURES: - args = {"__tool_version__": 2, "arguments": args} - elif self._tool_features: - args = { - "__tool_features__": sorted(self._tool_features), - "arguments": args, - } + if self._tool_features: + args = tool_call_frame(self._tool_features, args) try: conn.send(args) if conn.poll(self.timeout): @@ -2863,6 +3037,7 @@ def run_household( "benchmark_search_rule": profile.benchmark_search_rule, "sent_board_prompt_sha256": sent_board_prompt_sha256, "tool_notice_sha256": hashlib.sha256(profile.notice.encode()).hexdigest(), + **profile_treatment(profile), "turns": [], "status": "running", } @@ -3571,6 +3746,52 @@ def records_wording(records: Mapping[str, dict]) -> str | None: return hashes.pop() if hashes else None +# Arms whose calculator features never changed, so a record written before +# households recorded their features has these. +LEGACY_TOOL_ARMS = {"pe": frozenset(), "web+pe": frozenset(), "pe2": V2_FEATURES} + + +def profile_treatment(profile: "ModelProfile") -> dict: + """What a household record keeps of the treatment it ran under, beyond + the notice hash: the calculator's features and its schema's hash.""" + if not profile.uses_calculator: + return {"tool_features": None, "calculate_tool_sha256": None} + schema = calculate_tool_schema(features=profile.tool_features) + return { + "tool_features": sorted(profile.tool_features), + "calculate_tool_sha256": hashlib.sha256( + json.dumps(schema, sort_keys=True).encode() + ).hexdigest(), + } + + +def record_treatment_problem(record: dict, profile: "ModelProfile") -> str | None: + """Why a household record ran under another treatment than the profile's + (features, notice or calculator schema), or None. A record written before + a field existed counts as matching only for an arm whose treatment never + changed (none, web, and LEGACY_TOOL_ARMS); otherwise it is unrecorded.""" + sid = record.get("scenario_id") + want = profile_treatment(profile) + legacy = profile.arm in ("none", "web") or profile.arm in LEGACY_TOOL_ARMS + if "tool_features" in record: + if record["tool_features"] != want["tool_features"]: + return f"{sid} ran with tool features {record['tool_features']}" + elif not legacy: + return f"{sid} has no recorded tool features" + notice = record.get("tool_notice_sha256") + if notice is None and not legacy: + return f"{sid} has no recorded tool notice" + if ( + notice is not None + and notice != hashlib.sha256(profile.notice.encode()).hexdigest() + ): + return f"{sid} ran under another tool notice" + schema = record.get("calculate_tool_sha256") + if schema is not None and schema != want["calculate_tool_sha256"]: + return f"{sid} ran with another calculator schema" + return None + + def records_arm(records: Mapping[str, dict]) -> str: """The one arm a run directory's records belong to; records written before the web arms are the PE tool's.""" @@ -3629,6 +3850,17 @@ def command_run(args) -> None: f"{households_dir} holds records sent under other v2-pilot wording " f"({records_wording(previous)}); use a fresh --out-dir" ) + mixed = [ + problem + for record in previous.values() + if (problem := record_treatment_problem(record, profile)) is not None + ] + if mixed: + raise SystemExit( + f"{households_dir} holds records from another treatment (" + + "; ".join(mixed[:3]) + + "); use a fresh --out-dir" + ) spend = Spend(args.budget_usd, spent_before) todo = [sid for sid in ids if sid not in done] manifest = { @@ -4416,7 +4648,7 @@ def _person_names(arguments) -> list[str]: household = ( json.loads(household) if isinstance(household, str) else household ) - return [person["name"] for person in household["people"]] + return [name or "" for name in flat_names(household["people"])] situation = args["situation"] situation = json.loads(situation) if isinstance(situation, str) else situation return list(situation["people"]) diff --git a/tests/test_tool_v2.py b/tests/test_tool_v2.py index d6e381b5..68248090 100644 --- a/tests/test_tool_v2.py +++ b/tests/test_tool_v2.py @@ -261,8 +261,13 @@ def test_the_pool_tells_its_workers_the_tool_version(): ) try: result, _ = pool.call({"situation": "{}"}) + # A frame no model's JSON arguments can be (tool v3 review r1). assert result == { - "received": {"__tool_version__": 2, "arguments": {"situation": "{}"}} + "received": ( + "calculate", + ("explain", "names", "overrides", "placement"), + {"situation": "{}"}, + ) } finally: pool.close() diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index e0064022..4b1e9c99 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -179,7 +179,7 @@ def test_concise_takes_a_detail_and_defaults_to_brief(): {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, F.EXPLAIN, ) - assert "detail" not in request # every step, as tool version 2 explains + assert "detail" not in request # the trace, as tool version 2 explains def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_parameters(): @@ -268,24 +268,75 @@ def test_only_pe2_says_the_facts_carry_variable_names(): assert ("names it read" in notice) == bool(features & {F.NAMES, F.LABELS}) -def test_every_prompt_label_reads_back_to_its_engine_name(): - from policybench.prompts import INPUT_LABEL_OVERRIDES - from policybench.scenarios import PE_INPUT_RENAMES +def test_every_label_a_prompt_shows_reads_back_to_its_engine_name(): + from policybench.prompts import _format_input_line + from policybench.scenarios import PE_DROPPED_INPUTS, PE_INPUT_RENAMES labels = F.prompt_labels() - for field, label in INPUT_LABEL_OVERRIDES.items(): - assert labels[F._label_key(label)] == PE_INPUT_RENAMES.get(field, field) - # The v2-pilot prompt's own relabelings, as that prompt writes them. - relabeled = F.v2_pilot_labels() - assert set(relabeled) == { - "employer_sponsored_insurance_premiums", - "spm_unit_energy_subsidy", + fields = {spec.output_name: spec.value_type for spec in F._promptable_specs()} + fields |= {"age": "int", "employment_income": "float"} + dropped = {f for names in PE_DROPPED_INPUTS.values() for f in names} + for field, value_type in {**fields, **dict.fromkeys(dropped, "float")}.items(): + for label in F.rendered_labels(field, value_type): + want = (PE_INPUT_RENAMES.get(field, field), field in dropped) + assert want in labels[F._label_key(label)], (field, label) + # Fallback labels (the name with spaces) are the prompt's too. + assert ("first_home_mortgage_balance", False) in labels[ + F._label_key("first home mortgage balance") + ] + # Special renderers: the Marketplace plan line is labeled on its own. + line = _format_input_line("selected_marketplace_plan_benchmark_ratio", 0.8) + assert line.startswith("- selected Marketplace plan: ") + assert ("selected_marketplace_plan_benchmark_ratio", False) in labels[ + F._label_key("selected Marketplace plan") + ] + # The v2-pilot prompt's relabelings, as that prompt writes them. + for field, label in ( + ("spm_unit_energy_subsidy", "home energy assistance received"), + ( + "employer_sponsored_insurance_premiums", + "employer-paid health insurance premiums (paid by the employer, " + "not by the household)", + ), + ): + assert label in F.rendered_labels(field) + assert (field, False) in labels[F._label_key(label)] + assert ("state_code", False) in labels[F._label_key("state")] + assert ("first_home_mortgage_interest", True) in labels[ + F._label_key("first home mortgage interest") + ] + assert ("roth_401k_contributions_desired", False) in labels[ + F._label_key("Roth 401k contributions desired") + ] + assert ("partnership_s_corp_income", False) in labels[ + F._label_key("partnership or S-corp income") + ] + + +def test_labels_reads_no_label_that_names_two_of_the_engines_inputs(monkeypatch): + two = frozenset({("has_esi", False), ("age", False)}) + monkeypatch.setattr(F, "prompt_labels", lambda: {F._label_key("covered"): two}) + situation = _situation() + situation["people"]["head"]["covered"] = True + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert request is None and "unknown variable 'covered'" in str(error) + one = frozenset({("has_esi", False), ("uk_only_input", False)}) + monkeypatch.setattr(F, "prompt_labels", lambda: {F._label_key("covered"): one}) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None and request["situation"]["people"]["head"]["has_esi"] + + +def test_labels_ignores_a_fact_the_reference_leaves_out(): + situation = _situation() + situation["tax_units"] = { + "t": {"members": ["head", "kid"], "first home mortgage interest": 900} } - for field, label in relabeled.items(): - written = F.v2_pilot_text(f"- {field.replace('_', ' ')}: $5") - assert written == f"- {label}: $5" - assert labels[F._label_key(label)] == field - assert labels[F._label_key("roth 401k contributions")] == "roth_401k_contributions" + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None, error + assert request["situation"]["tax_units"]["t"] == {"members": ["head", "kid"]} + assert any("leaves it out too" in n for n in request["notes"]) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.NAMES) + assert request is None and "first home mortgage interest" in str(error) def _flat(**overrides): @@ -385,6 +436,74 @@ def test_flat_placement_gives_a_household_of_ones_person_fact_to_that_person(): assert any("whose facts list it" in p for p in error["problems"]) +def test_flat_refuses_facts_given_beside_the_person(): + people = [{"role": "head", "age": 40, "employment_income": 30000}] + request, error = _norm(_flat(people=people), F.FLAT) + assert request is None + assert any( + "takes only name, role and facts; give age, employment_income under facts" in p + for p in error["problems"] + ) + + +def test_flat_reads_a_fact_written_two_ways_once_whatever_the_order(): + for facts in ( + {"spm_unit_energy_subsidy": 900, "spm unit energy subsidy": 800}, + {"spm unit energy subsidy": 800, "spm_unit_energy_subsidy": 900}, + ): + request, error = _norm(_flat(household_facts=facts), F.FLAT, F.NAMES) + assert request is None + assert any("set twice with different values" in p for p in error["problems"]) + same = {"spm_unit_energy_subsidy": 900, "spm unit energy subsidy": {"2026": 900}} + request, error = _norm(_flat(household_facts=same), F.FLAT, F.NAMES) + assert error is None, error + + +def test_flat_roles_set_the_relationship_inputs_and_refuse_a_contradiction(): + for written in ("is_tax_unit_head", "is tax unit head"): + people = [{"name": "A", "role": "head", "facts": {written: False}}] + request, error = _norm(_flat(people=people), F.FLAT, F.NAMES) + assert request is None, written + assert any("contradicts the role head" in p for p in error["problems"]) + people = [{"name": "A", "role": "head", "facts": {written: True}}] + request, error = _norm(_flat(people=people), F.FLAT, F.NAMES) + assert error is None, (written, error) + head = request["situation"]["people"]["A"] + assert head["is_tax_unit_head"] == {"2026": True} + assert head["is_tax_unit_spouse"] == {"2026": False} + + +def test_flat_placement_compares_values_once_normalized(): + one = [ + {"name": "Solo", "role": "head", "facts": {"pre_subsidy_rent": {"2026": 9000}}} + ] + args = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, error + assert request["situation"]["people"]["Solo"]["pre_subsidy_rent"] == {"2026": 9000} + + +def test_the_descriptions_promise_only_what_the_features_do(): + flat = F.calculate_tool_schema(features=frozenset({F.FLAT}))["function"] + assert '"employment_income": 30000' in flat["description"] + example = [ + {"role": "head", "facts": {"age": 40, "employment_income": 30000}}, + ] + request, error = _norm(_flat(people=example), F.FLAT) + assert error is None, error + entity_error = "a variable set on the wrong entity" + placed = F.calculate_tool_schema(features=frozenset({F.PLACEMENT}))["function"] + assert entity_error not in placed["description"] + assert "moved to the one unit holding its people" in placed["description"] + named = F.calculate_tool_schema(features=frozenset({F.NAMES}))["function"] + assert entity_error in named["description"] + concise = F.calculate_tool_schema(features=frozenset({F.EXPLAIN, F.CONCISE})) + detail = concise["function"]["parameters"]["properties"]["detail"]["description"] + assert f"up to {F.EXPLAIN_MAX_LINES} steps" in detail + for features in (frozenset({F.EXPLAIN, F.CONCISE}), frozenset(F.TOOL_FEATURES)): + assert "every step" not in F.compose_tool_notice(features, True) + + def test_flat_refuses_what_it_cannot_place(): people = [ {"name": "A", "role": "head", "facts": {}}, @@ -434,7 +553,7 @@ def _echo_worker(conn, fix_path, reform): conn.send(("ok", {"received": args})) -def test_the_pool_tells_its_workers_the_features(): +def test_the_pool_tells_its_workers_the_features_in_a_frame(): pool = F.ToolPool( size=1, timeout=5.0, @@ -443,13 +562,91 @@ def test_the_pool_tells_its_workers_the_features(): ) try: result, _ = pool.call({"x": 1}) - assert result == { - "received": {"__tool_features__": ["flat", "names"], "arguments": {"x": 1}} - } + assert result == {"received": ("calculate", ("flat", "names"), {"x": 1})} finally: pool.close() +def test_a_model_cannot_turn_on_a_feature_its_arm_does_not_offer(): + forged = { + "__tool_features__": ["names"], + "arguments": { + "situation": {"people": {"head": {"Employment income": 12345}}}, + "variables": ["age"], + }, + } + for message in (forged, {"__tool_version__": 2, "arguments": forged["arguments"]}): + request, error = F.worker_request(message, REGISTRY) + assert request is None + assert "unknown arguments: " in error["problems"][0] + request, error = F.worker_request( + F.tool_call_frame(frozenset({F.NAMES}), forged["arguments"]), REGISTRY + ) + assert error is None + assert request["situation"]["people"]["head"]["employment_income"] == { + "2026": 12345 + } + for names in (("telepathy",), ("concise",), ({},), ["names"]): + request, error = F.worker_request(("calculate", names, {}), REGISTRY) + assert request is None and error["error"] == "configuration", names + + +def _profile(arm): + return F.ModelProfile("m", "gpt-x", "chat", "tool", arm=arm) + + +def _record(profile, **changes): + record = { + "scenario_id": "mt_000", + "tool_notice_sha256": __import__("hashlib") + .sha256(profile.notice.encode()) + .hexdigest(), + **F.profile_treatment(profile), + } + record.update(changes) + return record + + +def test_a_resume_refuses_households_run_under_another_treatment(): + pe3 = _profile("pe3") + assert F.record_treatment_problem(_record(pe3), pe3) is None + six = sorted(set(F.TOOL_FEATURES) - {F.LABELS}) + assert "tool features" in F.record_treatment_problem( + _record(pe3, tool_features=six), pe3 + ) + assert "another tool notice" in F.record_treatment_problem( + _record(pe3, tool_notice_sha256="0" * 64), pe3 + ) + assert "calculator schema" in F.record_treatment_problem( + _record(pe3, calculate_tool_sha256="0" * 64), pe3 + ) + old = _record(pe3) + del old["tool_features"] + assert "no recorded tool features" in F.record_treatment_problem(old, pe3) + # Arms whose treatment never changed accept records from before the fields. + for arm in ("pe", "pe2", "web+pe", "none", "web"): + profile = _profile(arm) + legacy = _record(profile) + for key in ("tool_features", "calculate_tool_sha256", "tool_notice_sha256"): + legacy.pop(key) + assert F.record_treatment_problem(legacy, profile) is None, arm + pe2 = _profile("pe2") + assert "tool features" in F.record_treatment_problem(_record(_profile("pe")), pe2) + + +def test_diagnostics_read_flat_people_by_their_default_names(): + household = { + "people": [ + {"role": "head", "facts": {}}, + {"role": "spouse", "facts": {}}, + {"name": "Kid", "role": "dependent", "facts": {}}, + {"role": "dependent", "facts": {}}, + ] + } + names = F._person_names({"household": json.dumps(household)}) + assert names == ["head", "spouse", "Kid", "dependent_2"] + + @pytest.mark.slow def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): """Differential: on the installed engine, each frozen-board household @@ -520,73 +717,72 @@ def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): @pytest.mark.slow -def test_a_household_described_under_the_prompts_labels_computes_the_same(): - """Differential: each frozen-board household described flat with every - fact as its v2-pilot prompt writes it (a label in words, or the name with - spaces), - and stated hours only where the prompt states them, computes every board - output as the benchmark does. Take-up conventions, which the prompt does - not list, keep their engine names.""" +def test_a_household_described_from_its_prompt_computes_the_same(): + """Differential, with the labels feature alone: each of 30 frozen-board + households described flat from the facts its v2-pilot prompt shows, each + under the label the prompt shows it under (dropped facts included), with + the state and roles the prompt gives, computes every board output as the + benchmark's own situation does. The take-up conventions, which the prompt + states in words, are set under their engine names.""" import household_set as hs import pandas as pd from policyengine_us import Simulation from policybench.ground_truth import _extract_person_value, _pe_variable_for_output - from policybench.prompts import _humanize_input_label - from policybench.scenarios import PE_INPUT_ALIASES, PE_INPUT_RENAMES + from policybench.scenarios import ( + DEFAULT_TAKEUP_INPUTS, + is_excluded_prompt_input_name, + ) system = F.build_tool_system(reform=True) registry = F.build_registry(system) reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") variables = reference.groupby("scenario_id").variable.agg(list) - field_of = {name: field for field, name in PE_INPUT_RENAMES.items()} - dropped = {"is_tax_unit_head", "is_tax_unit_spouse", *PE_INPUT_ALIASES.values()} + shown = set() - worded = [] - - def labeled(facts): - out = {} - for k, v in facts.items(): - if k in dropped: - continue - if k.startswith(F.CONVENTION_PREFIXES): - out[k] = v - continue - label = F.v2_pilot_text(f"- {_humanize_input_label(field_of.get(k, k))}: ")[ - 2:-2 - ] - if label != k.replace("_", " "): - worded.append(label) - out[label] = v - return out + def label(field, value): + value_type = ( + "bool" + if isinstance(value, bool) + else ("str" if isinstance(value, str) else "float") + ) + written = F.rendered_labels(field, value_type)[-1] + shown.add((field, written)) + return written + + def prompt_facts(inputs): + return { + label(k, v): v + for k, v in inputs.items() + if not is_excluded_prompt_input_name(k) + } checked = 0 - for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:12]: - situation = hs.situation_for(scenario) + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:30]: couple = scenario.marital_couple() or () head = couple[0] if couple else scenario.adults[0].name - people = [ - { - "name": name, - "role": "head" - if name == head - else ("spouse" if name in couple else "dependent"), - "facts": labeled(facts), - } - for name, facts in situation["people"].items() - ] + people = [] + for person in scenario.all_people: + facts = {"age": person.age} + if person.name.startswith("adult") or abs(person.employment_income) > 1e-6: + facts[label("employment_income", 1.0)] = person.employment_income + facts |= prompt_facts(person.inputs) + facts |= DEFAULT_TAKEUP_INPUTS["person"] + role = ( + "head" + if person.name == head + else ("spouse" if person.name in couple else "dependent") + ) + people.append({"name": person.name, "role": role, "facts": facts}) household_facts = {} - for group in ("tax_units", "spm_units", "families", "households"): - for unit in situation[group].values(): - household_facts.update( - labeled( - { - k: v - for k, v in unit.items() - if k not in ("members", "state_code") - } - ) - ) + for inputs in ( + scenario.tax_unit_inputs, + scenario.spm_unit_inputs, + scenario.household_inputs, + ): + household_facts |= prompt_facts(inputs) + for entity in ("tax_unit", "spm_unit", "household"): + household_facts |= DEFAULT_TAKEUP_INPUTS[entity] household = { "state": scenario.state, "people": people, @@ -595,10 +791,12 @@ def labeled(facts): request, error = F.normalize_request( {"household": json.dumps(household), "variables": ["snap"]}, registry, - features=frozenset({F.FLAT, F.LABELS, F.NAMES}), + features=frozenset({F.FLAT, F.LABELS}), ) assert error is None, (scenario.id, error) - own = Simulation(tax_benefit_system=system, situation=situation) + own = Simulation( + tax_benefit_system=system, situation=hs.situation_for(scenario) + ) flat = Simulation(tax_benefit_system=system, situation=request["situation"]) for variable in variables[scenario.id]: name = _pe_variable_for_output(variable, "us") @@ -612,4 +810,12 @@ def labeled(facts): ] assert values[0] == values[1], (scenario.id, variable, values) checked += 1 - assert checked > 120 and len(worded) > 10, (checked, worded) + worded = {w for f, w in shown if w != f.replace("_", " ")} + assert checked > 300 and len(worded) > 10, (checked, sorted(worded)) + # No label a US prompt can show names two of this engine's inputs. + ambiguous = { + key: targets + for key, targets in F.prompt_labels().items() + if len({e for e, dropped in targets if not dropped and e in registry}) > 1 + } + assert not ambiguous, ambiguous From d570b8e7af46a99f0dd4110be2ec8ea9147de835 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 20:09:47 -0400 Subject: [PATCH 7/9] Tool v3 review r2: report malformed repeats, resumable harness errors - flat: a household fact repeated with a malformed value was dropped without a word. _same_value now reports the repeat's validation problems, and only a repeat validated as equal is dropped quietly. A malformed first value is reported once, by normal validation, and never as a conflict. - The fallback record for a household the harness failed on carried no treatment, so one failure blocked resuming any custom arm. The record is built by harness_error_record, which now includes the notice hash and profile_treatment. resume_treatment_problems checks only the records a resume keeps, since a rerunnable record is replaced under the current profile. pe, web+pe and pe2 still behave exactly as at 5bbd571a: 4,000 calls and 2,000 pool messages match. 5 mutants of these fixes were all caught. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 102 +++++++++++++++++++++++++----------- tests/test_tool_v3.py | 54 +++++++++++++++++++ 2 files changed, 126 insertions(+), 30 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 5f586522..b7b66b72 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1280,14 +1280,29 @@ def flat_names(people) -> list: return names -def _same_value(registry, name: str, a, b, period: str) -> bool | None: - """Whether two values of an input are the same once normalized (None when - either is malformed; normalization reports it).""" - scratch: list[str] = [] +def _same_value( + registry, + name: str, + a, + b, + period: str, + problems: list[str] | None = None, + where: str = "", +) -> bool | None: + """Whether two values of an input are the same once normalized, or None + when either is malformed. ``a`` stays in the situation, so normalization + reports it later; a caller that drops ``b`` passes ``problems`` so a + malformed ``b`` is reported, not lost.""" info = registry[name] - first = _normalize_value("", name, a, info, period, scratch) - second = _normalize_value("", name, b, info, period, scratch) - return None if scratch else first == second + first_problems: list[str] = [] + second_problems: list[str] = [] + first = _normalize_value(where, name, a, info, period, first_problems) + second = _normalize_value(where, name, b, info, period, second_problems) + if problems is not None: + problems.extend(second_problems) + if first_problems or second_problems: + return None + return first == second def flat_situation( @@ -1460,7 +1475,12 @@ def engine_name(name): else: target = unit_of[entity] if name in target: - if _same_value(registry, resolved, target[name], value, period) is False: + # The repeat is dropped, so only one validated as equal is kept quiet. + same = _same_value( + registry, resolved, target[name], value, period, problems, + "household_facts", + ) + if same is False: problems.append( f"household_facts: {resolved} is set twice with different values" ) @@ -3792,6 +3812,46 @@ def record_treatment_problem(record: dict, profile: "ModelProfile") -> str | Non return None +def resume_treatment_problems( + previous: Mapping[str, dict], done: set[str], profile: "ModelProfile" +) -> list[str]: + """record_treatment_problem for each record a resume keeps. A rerunnable + record (RERUNNABLE_STATUSES) is replaced under this profile, so it cannot + mix treatments and does not block the resume.""" + return [ + problem + for sid, record in previous.items() + if sid in done + and (problem := record_treatment_problem(record, profile)) is not None + ] + + +def harness_error_record( + sid: str, variables: list[str], profile: "ModelProfile", error: Exception +) -> dict: + """The record of a household the harness itself failed on: rerunnable, + with every setup and treatment key a resume and score check, so a folder + holding it still resumes and scores.""" + record = { + "scenario_id": sid, + "model": profile.name, + "arm": profile.arm, + "prompt_version": profile.prompt_version, + "prompt_wording_sha256": prompt_wording_sha256(profile), + "benchmark_search_rule": profile.benchmark_search_rule, + "tool_notice_sha256": hashlib.sha256(profile.notice.encode()).hexdigest(), + **profile_treatment(profile), + "variables": variables, + "status": "harness_error", + "error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS], + "turns": [], + "predictions": {v: None for v in variables}, + "explanations": {v: None for v in variables}, + } + record["totals"] = household_totals(record) + return record + + def records_arm(records: Mapping[str, dict]) -> str: """The one arm a run directory's records belong to; records written before the web arms are the PE tool's.""" @@ -3850,11 +3910,7 @@ def command_run(args) -> None: f"{households_dir} holds records sent under other v2-pilot wording " f"({records_wording(previous)}); use a fresh --out-dir" ) - mixed = [ - problem - for record in previous.values() - if (problem := record_treatment_problem(record, profile)) is not None - ] + mixed = resume_treatment_problems(previous, done, profile) if mixed: raise SystemExit( f"{households_dir} holds records from another treatment (" @@ -3959,23 +4015,9 @@ def one(sid: str) -> dict: profile=profile, ) except Exception as error: # one household must not end the run - record = { - "scenario_id": sid, - "model": profile.name, - # The setup keys a resume and score check, so a folder with - # this record still resumes and scores. - "arm": profile.arm, - "prompt_version": profile.prompt_version, - "prompt_wording_sha256": prompt_wording_sha256(profile), - "benchmark_search_rule": profile.benchmark_search_rule, - "variables": board["variables"][sid], - "status": "harness_error", - "error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS], - "turns": [], - "predictions": {v: None for v in board["variables"][sid]}, - "explanations": {v: None for v in board["variables"][sid]}, - } - record["totals"] = household_totals(record) + record = harness_error_record( + sid, board["variables"][sid], profile, error + ) if sid in previous: prior = previous[sid] record["prior_attempts"] = prior.get("prior_attempts", []) + [ diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index 4b1e9c99..dcf6317a 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -819,3 +819,57 @@ def prompt_facts(inputs): if len({e for e, dropped in targets if not dropped and e in registry}) > 1 } assert not ambiguous, ambiguous + + +def test_flat_reports_a_malformed_repeat_of_a_fact_it_drops(): + # Review r2: a repeat that fails validation was dropped without a word. + one = [{"name": "Solo", "role": "head", "facts": {"employment_income": 30000}}] + for bad in ("40000", [40000], {"26": 40000}, {}): + args = _flat(people=one, household_facts={"employment_income": bad}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None, bad + assert any("employment_income" in m for m in error["problems"]), (bad, error) + # An equal repeat, scalar or by period, still merges quietly; a different + # one is still refused. + for same in (30000, {"2026": 30000}): + args = _flat(people=one, household_facts={"employment_income": same}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, (same, error) + args = _flat(people=one, household_facts={"employment_income": 40000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None + assert any("set twice with different values" in m for m in error["problems"]) + + +def test_a_harness_error_record_resumes_under_its_own_arm(): + # Review r2: the fallback record carried no treatment, so one harness + # failure blocked every custom-arm resume. + for arm in ("pe3", "pe:labels+flat", "pe", "pe2", "web+pe", "none", "web"): + profile = _profile(arm) + record = F.harness_error_record("mt_000", ["snap"], profile, RuntimeError("x")) + assert record["status"] in F.RERUNNABLE_STATUSES + assert F.record_treatment_problem(record, profile) is None, arm + assert record["predictions"] == {"snap": None} + assert record["totals"]["cost_usd"] == 0 + + +def test_a_resume_checks_the_treatment_only_of_records_it_keeps(): + pe3 = _profile("pe3") + stale = _record(pe3, tool_features=["names"]) + previous = { + "mt_000": dict(stale, status="complete"), + "mt_001": dict(stale, scenario_id="mt_001", status="harness_error"), + } + # Kept, so refused. + assert len(F.resume_treatment_problems(previous, {"mt_000"}, pe3)) == 1 + # Rerun under this profile, so no mixing and no block. + assert F.resume_treatment_problems(previous, set(), pe3) == [] + + +def test_flat_reports_a_malformed_first_value_once_and_not_as_a_conflict(): + one = [{"name": "Solo", "role": "head", "facts": {"employment_income": "30000"}}] + args = _flat(people=one, household_facts={"employment_income": 40000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None + assert not any("set twice" in m for m in error["problems"]), error + assert sum("employment_income" in m for m in error["problems"]) == 1, error From fb35d8273b0d1f6fd5df1aea76b6e22d15af0e8f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 20:45:58 -0400 Subject: [PATCH 8/9] Format the r2 fix Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index b7b66b72..abab108f 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1477,7 +1477,12 @@ def engine_name(name): if name in target: # The repeat is dropped, so only one validated as equal is kept quiet. same = _same_value( - registry, resolved, target[name], value, period, problems, + registry, + resolved, + target[name], + value, + period, + problems, "household_facts", ) if same is False: From 6b00a6f0c4997e73b9968b412b57646957cad631 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 01:29:28 -0400 Subject: [PATCH 9/9] Tool v3 review r3: check every record's treatment, at resume and at score A rerun is not sure to replace a rerunnable record: a budget stop, lost workers or a narrower --scenarios leave it in place, and scoring read whatever was there. treatment_problems now checks every record in the folder, rerunnable ones included, and command_score refuses a folder that mixes treatments, rebuilding the notices on the board's engine. Runs written before 6e1578fc (test batch 0's frontier calculator arms, b6d835a8) were told policyengine-us 2.41.0 while computing on 2.60.4. A folder written wholly so is one treatment and scores with a note; a mix of the two texts is refused, and a resume into it stays refused, since its new records would state the board's engine. All 39 run folders score (628 records); the six frontier calculator folders take the note. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 63 ++++++++++++--- tests/test_tool_v3.py | 154 ++++++++++++++++++++++++++++++++++-- 2 files changed, 201 insertions(+), 16 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index abab108f..f09e0b6f 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -48,6 +48,7 @@ import argparse import concurrent.futures +import contextlib import copy import decimal import difflib @@ -343,6 +344,18 @@ def set_stated_engine(version: str) -> None: _STATED_ENGINE["version"] = version +@contextlib.contextmanager +def stating_engine(version: str): + """The texts as a run on `version` sent them, for a check that rebuilds a + notice; the stated engine is restored afterwards.""" + before = _STATED_ENGINE["version"] + set_stated_engine(version) + try: + yield + finally: + set_stated_engine(before) + + def engine_text(text: str) -> str: return text.replace( f"policyengine-us {REFERENCE_ENGINE}", @@ -3817,20 +3830,43 @@ def record_treatment_problem(record: dict, profile: "ModelProfile") -> str | Non return None -def resume_treatment_problems( - previous: Mapping[str, dict], done: set[str], profile: "ModelProfile" +def treatment_problems( + records: Mapping[str, dict], profile: "ModelProfile" ) -> list[str]: - """record_treatment_problem for each record a resume keeps. A rerunnable - record (RERUNNABLE_STATUSES) is replaced under this profile, so it cannot - mix treatments and does not block the resume.""" + """record_treatment_problem for every record in a run directory, a + rerunnable one (RERUNNABLE_STATUSES) too. A rerun is not sure to replace + it: a budget stop, lost workers or a narrower --scenarios leave it in + place, and scoring reads whatever is there (review r3). So a resume and + a score both refuse a directory holding any record from another treatment.""" return [ problem - for sid, record in previous.items() - if sid in done - and (problem := record_treatment_problem(record, profile)) is not None + for record in records.values() + if (problem := record_treatment_problem(record, profile)) is not None ] +def scored_treatment_problems( + records: Mapping[str, dict], profile: "ModelProfile", engine: str +) -> tuple[list[str], str | None]: + """treatment_problems as score applies it, with the note score prints. + The records' notices state the board's engine, as command_run set it. + Runs before the texts stated the board's engine (6e1578fc) named the + reference engine whatever the board's: a directory written wholly so is + one treatment and scores as it ran. A mix of the two texts does not. A + resume stays strict, since its new records would state the board's.""" + with stating_engine(engine): + problems = treatment_problems(records, profile) + if not problems or engine == REFERENCE_ENGINE: + return problems, None + with stating_engine(REFERENCE_ENGINE): + if treatment_problems(records, profile): + return problems, None + return [], ( + f"NOTE: these records' notices name policyengine-us {REFERENCE_ENGINE}, " + f"as runs before 6e1578fc stated it; the board's engine is {engine}" + ) + + def harness_error_record( sid: str, variables: list[str], profile: "ModelProfile", error: Exception ) -> dict: @@ -3915,7 +3951,7 @@ def command_run(args) -> None: f"{households_dir} holds records sent under other v2-pilot wording " f"({records_wording(previous)}); use a fresh --out-dir" ) - mixed = resume_treatment_problems(previous, done, profile) + mixed = treatment_problems(previous, profile) if mixed: raise SystemExit( f"{households_dir} holds records from another treatment (" @@ -4148,6 +4184,15 @@ def command_score(args) -> None: profile = model_profile( records_model(records), records_arm(records), *records_prompt(records) ) + mixed, note = scored_treatment_problems(records, profile, board_engine(snapshot)) + if note: + print(note) + if mixed: + raise SystemExit( + f"{out / 'households'} holds records from another treatment (" + + "; ".join(mixed[:3]) + + "); scoring would mix them" + ) arm1_label, arm2_label = profile.arm1_label, profile.arm2_label complete = sorted( sid for sid, r in records.items() if r.get("status") == "complete" diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index dcf6317a..71bc73c3 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -853,17 +853,157 @@ def test_a_harness_error_record_resumes_under_its_own_arm(): assert record["totals"]["cost_usd"] == 0 -def test_a_resume_checks_the_treatment_only_of_records_it_keeps(): +def test_a_rerunnable_record_from_another_treatment_is_refused_too(): + # Review r3: a rerun is not sure to replace a rerunnable record (a budget + # stop, lost workers, a narrower --scenarios), so every record is checked. pe3 = _profile("pe3") stale = _record(pe3, tool_features=["names"]) - previous = { + records = { "mt_000": dict(stale, status="complete"), - "mt_001": dict(stale, scenario_id="mt_001", status="harness_error"), + "mt_001": dict(stale, scenario_id="mt_001", status="provider_error"), + "mt_002": F.harness_error_record("mt_002", ["snap"], pe3, RuntimeError("x")), } - # Kept, so refused. - assert len(F.resume_treatment_problems(previous, {"mt_000"}, pe3)) == 1 - # Rerun under this profile, so no mixing and no block. - assert F.resume_treatment_problems(previous, set(), pe3) == [] + problems = F.treatment_problems(records, pe3) + assert [p.split()[0] for p in problems] == ["mt_000", "mt_001"] + + +SIX = ("concise", "explain", "flat", "names", "overrides", "placement") + + +def _scored_record(sid, profile, treatment_of=None, status="complete", **changes): + """A household record as command_run writes it, with the treatment of + `treatment_of` (another arm's profile) when given.""" + import hashlib + + board = pilot.load_board() + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + variables = board["variables"][sid] + source = treatment_of or profile + record = { + "scenario_id": sid, + "model": profile.name, + "arm": profile.arm, + "prompt_version": profile.prompt_version, + "benchmark_search_rule": profile.benchmark_search_rule, + "tool_notice_sha256": hashlib.sha256(source.notice.encode()).hexdigest(), + **F.profile_treatment(source), + "variables": variables, + "predictions": {v: float(reference[(sid, v)]) for v in variables}, + "explanations": {v: "x" for v in variables}, + "status": status, + "turns": [{"total_cost_usd": 1.25, "calls": []}], + "cap_hit": False, + "wall_seconds": 1.0, + } + record.update(changes) + record["totals"] = F.household_totals(record) + return record + + +def _write(directory, *records): + (directory / "households").mkdir(parents=True, exist_ok=True) + for record in records: + path = directory / "households" / f"{record['scenario_id']}.json" + path.write_text(json.dumps(record)) + + +def test_a_resume_refuses_an_old_treatments_provider_error_before_running( + tmp_path, monkeypatch +): + # The review's case: a pe3 folder holding a provider_error written under + # six features, with the budget already spent. Before the fix the resume + # passed, `one` stopped on the budget without replacing it, and scoring + # read its predictions as pe3's. + import argparse + + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + six = F.model_profile(F.MODEL_NAME, "pe:" + "+".join(SIX)) + _write(tmp_path, _scored_record("scenario_000", pe3, six, status="provider_error")) + monkeypatch.setattr(F, "_require_engine", lambda board: None) + monkeypatch.setenv(F.provider_key(pe3.model_id), "test-key") + + def no_pool(**kwargs): + raise AssertionError("the run got as far as starting the calculator") + + monkeypatch.setattr(F, "ToolPool", no_pool) + args = argparse.Namespace( + out_dir=str(tmp_path), + board=None, + model=F.MODEL_NAME, + arm="pe3", + prompt="v1", + forbid_benchmark_search=False, + scenarios=None, + limit=None, + workers=1, + budget_usd=1.0, + reserve_usd=0.0, + ) + with pytest.raises(SystemExit, match="another treatment"): + F.command_run(args) + assert not list(tmp_path.glob("manifest-*.json")) + + +def test_score_refuses_a_folder_mixing_treatments(tmp_path): + import argparse + + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + six = F.model_profile(F.MODEL_NAME, "pe:" + "+".join(SIX)) + _write( + tmp_path, + _scored_record("scenario_000", pe3), + _scored_record("scenario_001", pe3, six, status="provider_error"), + ) + args = argparse.Namespace(out_dir=str(tmp_path), complete_only=False) + with pytest.raises(SystemExit, match="another treatment"): + F.command_score(args) + assert not (tmp_path / "score.json").exists() + # The same folder without the stray record scores. + (tmp_path / "households" / "scenario_001.json").unlink() + F.command_score(args) + assert json.loads((tmp_path / "score.json").read_text())["households"] == 1 + + +def test_score_rebuilds_notices_on_the_boards_engine_and_restores_it( + tmp_path, monkeypatch +): + # A household set's records state its engine (2.60.4) in their notices and + # schema, so a check on the reference engine's text would refuse them all. + import argparse + + monkeypatch.setattr(F, "board_engine", lambda snapshot=None: "2.60.4") + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + with F.stating_engine("2.60.4"): + current = _scored_record("scenario_000", pe3) + _write(tmp_path, current) + args = argparse.Namespace(out_dir=str(tmp_path), complete_only=False) + F.command_score(args) + assert F._STATED_ENGINE["version"] == F.REFERENCE_ENGINE + # Beside it, a record sent under the reference engine's text mixes two. + _write(tmp_path, _scored_record("scenario_001", pe3)) + with pytest.raises(SystemExit, match="another tool notice"): + F.command_score(args) + assert F._STATED_ENGINE["version"] == F.REFERENCE_ENGINE + + +def test_score_takes_a_folder_written_wholly_before_the_texts_stated_the_engine( + tmp_path, monkeypatch, capsys +): + # Test batch 0's frontier calculator arms (b6d835a8) were told 2.41.0 while + # the board and calculator ran 2.60.4: one treatment, scored as it ran. A + # resume into it is still refused, since new records would say 2.60.4. + import argparse + + monkeypatch.setattr(F, "board_engine", lambda snapshot=None: "2.60.4") + pe = F.model_profile(F.MODEL_NAME, "pe") + _write( + tmp_path, _scored_record("scenario_000", pe), _scored_record("scenario_001", pe) + ) + F.command_score(argparse.Namespace(out_dir=str(tmp_path), complete_only=False)) + assert "name policyengine-us 2.41.0" in capsys.readouterr().out + assert json.loads((tmp_path / "score.json").read_text())["households"] == 2 + with F.stating_engine("2.60.4"): + assert len(F.treatment_problems(F.load_records(tmp_path), pe)) == 2 def test_flat_reports_a_malformed_first_value_once_and_not_as_a_conflict():