diff --git a/policybench/reference_adversary.py b/policybench/reference_adversary.py index 5dfd5740..301a2d7e 100644 --- a/policybench/reference_adversary.py +++ b/policybench/reference_adversary.py @@ -73,6 +73,11 @@ REFERENCE_FREEZES = {"us": REFERENCE_FREEZE, "uk": "2026-10-10"} UK_FISCAL_YEAR = f"{TAX_YEAR}-{str(TAX_YEAR + 1)[-2:]}" +# How the stage headers say the consensus models answered. A payload's +# ``answerCondition`` replaces it, for clusters of tool-using arms (the tool +# pilot's web arms); the default renders the headers byte for byte as before. +DEFAULT_ANSWER_CONDITION = "answering from memory, without tools" + VERDICTS = ( "reference_holds", "reference_wrong", @@ -457,6 +462,7 @@ class AdversaryCase: metric_type: str = "amount" state: str = "" trigger: tuple[str, ...] = () + answer_condition: str = DEFAULT_ANSWER_CONDITION def to_manifest_row(self) -> dict: """The case as a ``cases.jsonl`` row, without the derivation.""" @@ -465,6 +471,8 @@ def to_manifest_row(self) -> dict: row["consensus"] = [dict(cluster) for cluster in self.consensus] row["trigger"] = list(self.trigger) row["derivation_sha256"] = sha256_text(derivation) if derivation else None + if self.answer_condition == DEFAULT_ANSWER_CONDITION: + row.pop("answer_condition") return row @classmethod @@ -486,6 +494,9 @@ def from_manifest_row(cls, row: dict, derivation: str) -> AdversaryCase: metric_type=str(row.get("metric_type", "amount")), state=str(row.get("state", "") or ""), trigger=tuple(row.get("trigger", ())), + answer_condition=str( + row.get("answer_condition") or DEFAULT_ANSWER_CONDITION + ), ) @@ -600,6 +611,10 @@ def build_adversary_cases( country = str(payload.get("country") or "us") scenarios = payload.get("scenarios") or {} predictions = payload.get("scenarioPredictions") or {} + # A payload may state the definitions its models saw (the tool pilot's + # v2-pilot wording) and the condition they answered under. + definitions = payload.get("outputDefinitions") or {} + condition = answer_condition_of(payload) cases: list[AdversaryCase] = [] seen: set[tuple[str, str]] = set() for flag in flags: @@ -673,7 +688,8 @@ def build_adversary_cases( country=country, scenario_id=scenario_id, variable=variable, - definition=get_variable_description(variable, country=country), + definition=_clean(definitions.get(variable)) + or get_variable_description(variable, country=country), household_prompt=household_prompt, reference_value=reference, consensus=tuple(clusters), @@ -684,11 +700,33 @@ def build_adversary_cases( metric_type=metric_type_for_output(variable), state=_clean(flag.get("state") or scenario.get("state")), trigger=tuple(str(t) for t in flag.get("trigger", ())), + answer_condition=condition, ) ) return cases +def answer_condition_of(payload: dict) -> str: + """The payload's ``answerCondition``, checked, or the default.""" + condition = payload.get("answerCondition") + if condition is None: + return DEFAULT_ANSWER_CONDITION + if not isinstance(condition, str) or not condition.startswith("answering "): + raise ValueError( + f"answerCondition must be a phrase starting 'answering ': {condition!r}" + ) + return condition + + +def _with_condition(header: str, case: AdversaryCase) -> str: + """A stage header stating the case's answer condition.""" + if case.answer_condition == DEFAULT_ANSWER_CONDITION: + return header + if header.count(DEFAULT_ANSWER_CONDITION) != 1: + raise ValueError(f"{case.country} header does not state the answer condition") + return header.replace(DEFAULT_ANSWER_CONDITION, case.answer_condition) + + # --- Prompts ------------------------------------------------------------------- _DOMAINS = ", ".join(BLOCKED_DOMAINS) @@ -1022,7 +1060,7 @@ def _case_sections(case: AdversaryCase) -> list[str]: def render_stage1_prompt(case: AdversaryCase) -> str: """The law-first prompt: everything about the case except the derivation.""" - header = _country_template(_STAGE1_HEADERS, case.country) + header = _with_condition(_country_template(_STAGE1_HEADERS, case.country), case) return "\n".join([header, "", *_case_sections(case)]) @@ -1036,7 +1074,7 @@ def render_stage2_prompt(case: AdversaryCase, stage1: dict, stage1_sha256: str) derivation = case.derivation or "(no derivation was recorded for this case)" return "\n".join( [ - _country_template(_STAGE2_HEADERS, case.country), + _with_condition(_country_template(_STAGE2_HEADERS, case.country), case), "", *_case_sections(case), "", @@ -1646,6 +1684,24 @@ def _markers_in(text: str) -> list[str]: return [marker for marker in DERIVATION_MARKERS if marker in lowered] +# An exclusion of a blocked domain in a search query ("-site:policyengine.org"): +# how a judge whose search tool takes no blocked_domains (Codex) follows the +# prompt's rule to exclude those domains on every search. It keeps a domain +# out; it does not search for it. +_EXCLUDED_DOMAIN = re.compile( + r"(? str: + """A search query without its exclusions of blocked domains: what the + judge searched for, for the contamination check.""" + return _EXCLUDED_DOMAIN.sub(" ", str(query or "")) + + def codex_events_audit(events_text: str) -> tuple[list[str], dict]: """Audit a ``codex exec --json`` event log for contamination. @@ -1714,7 +1770,7 @@ def codex_events_audit(events_text: str) -> tuple[list[str], dict]: ] url = str(action.get("url") or "") for text in [str(item.get("query") or ""), *searched]: - for marker in _markers_in(text): + for marker in _markers_in(searched_for(text)): problems.append(f"line {number}: a web search names {marker!r}") for marker in _markers_in(str(action.get("pattern") or "")): problems.append(f"line {number}: an in-page search names {marker!r}") @@ -1819,7 +1875,7 @@ def claude_transcript_audit( query = str(payload.get("query") or "") activity["searches"].append(query) searches.append((str(part.get("id")), query)) - if any(n in query.lower() for n in BLOCKED_NAMES): + if any(n in searched_for(query).lower() for n in BLOCKED_NAMES): problems.append(f"the judge searched for the engine: {query}") else: problems.append(f"the judge called a tool it was not given: {name}") @@ -2064,6 +2120,7 @@ def check_login( declared: str, config_dir: str, desktop_opt_in: bool, + api_key_secret: str = "", ) -> dict: """The Claude login a runner may judge with, or ``ValueError`` saying why not. @@ -2073,6 +2130,12 @@ def check_login( in which case it must be exactly the desktop's own claude.ai login. Unlike that script, it also refuses a token lane whose declared account (AUDIT_ACCOUNT) is the desktop login's. + + With ``api_key_secret`` (the runner's JUDGE_API_KEY_SECRET, Max + 2026-10-22) the judges bill that API key instead: the CLI must report a + first-party ``api_key`` login from ``ANTHROPIC_API_KEY`` in exactly + ``config_dir`` (the runner's fresh, empty one), with no lane token and no + desktop opt-in. """ def parse(text: str) -> dict: @@ -2092,6 +2155,31 @@ def parse(text: str) -> dict: f"the login is through {status.get('apiProvider')!r}, not a " "first-party subscription" ) + if api_key_secret: + if token or desktop_opt_in: + raise ValueError( + "an API key judge takes no lane token and no desktop opt-in" + ) + if method != "api_key" or status.get("apiKeySource") != "ANTHROPIC_API_KEY": + raise ValueError( + f"JUDGE_API_KEY_SECRET is set but the CLI logs in by {method!r} " + f"from {status.get('apiKeySource')!r}, not ANTHROPIC_API_KEY" + ) + if not ( + isinstance(reported, str) + and os.path.isdir(reported) + and os.path.samefile(reported, config_dir) + ): + raise ValueError( + f"the CLI reports config directory {reported!r}, not the run's " + f"own {config_dir}" + ) + return { + "method": "api_key", + "account": None, + "org": None, + "secret": api_key_secret, + } if method not in ("oauth_token", "claude.ai"): raise ValueError(f"the login method {method!r} is not a subscription login") same_account = bool( @@ -2188,6 +2276,7 @@ def main(argv: list[str] | None = None) -> int: login.add_argument("--declared", default="") login.add_argument("--config-dir", required=True) login.add_argument("--desktop-opt-in", choices=["0", "1"], default="0") + login.add_argument("--api-key-secret", default="") args = parser.parse_args(argv) if args.command == "render-stage2": @@ -2226,6 +2315,7 @@ def main(argv: list[str] | None = None) -> int: declared=args.declared, config_dir=args.config_dir, desktop_opt_in=args.desktop_opt_in == "1", + api_key_secret=args.api_key_secret, ) except ValueError as exc: print(str(exc), file=sys.stderr) diff --git a/policybench/scenarios.py b/policybench/scenarios.py index 35e88d00..4e8ae379 100644 --- a/policybench/scenarios.py +++ b/policybench/scenarios.py @@ -918,19 +918,29 @@ def scenario_from_dict(data: dict[str, Any]) -> Scenario: ) -def load_certified_us_person_frame() -> tuple[pd.DataFrame, int, str]: +def load_certified_us_person_frame( + dataset: str | None = None, dataset_label: str | None = None +) -> tuple[pd.DataFrame, int, str]: """Load a person-level frame from the certified US microsimulation dataset. Returns ``(person_df, dataset_year, dataset_label)``. The label is read from the live simulation's PolicyEngine bundle (``default_dataset``), so it reflects the dataset the run actually materialized -- the certified populace build (``populace_us_2024``) -- rather than a hardcoded name. + + ``dataset`` (a local H5 path) replaces the certified dataset, for test + household sets only; ``dataset_label`` then names it. """ from policybench.policyengine_runtime import make_us_microsimulation - sim = make_us_microsimulation() + if dataset is None: + sim = make_us_microsimulation() + dataset_label = str(sim.policyengine_bundle["default_dataset"]) + else: + if not dataset_label: + raise ValueError("an explicit dataset needs a dataset_label") + sim = make_us_microsimulation(dataset=dataset) dataset_year = sim.default_input_period - dataset_label = str(sim.policyengine_bundle["default_dataset"]) input_specs = get_promptable_input_specs() values = {} @@ -1822,10 +1832,22 @@ def generate_scenarios( seed: int = SEED, excluded_household_ids: set[int] | None = None, country: str = DEFAULT_COUNTRY, + dataset: str | None = None, + dataset_label: str | None = None, ) -> list[Scenario]: - """Generate benchmark scenarios for a country.""" + """Generate benchmark scenarios for a country. + + ``dataset`` and ``dataset_label`` replace the certified US dataset (see + :func:`load_certified_us_person_frame`); test household sets only. + """ + if dataset is not None and country != "us": + raise ValueError("an explicit dataset is supported for the US only") if country == "us": - person_df, dataset_year, dataset_label = load_certified_us_person_frame() + person_df, dataset_year, dataset_label = ( + load_certified_us_person_frame() + if dataset is None + else load_certified_us_person_frame(dataset, dataset_label) + ) return scenarios_from_cps_frame( person_df, n=n, diff --git a/scripts/batch_judge.py b/scripts/batch_judge.py new file mode 100644 index 00000000..17e33045 --- /dev/null +++ b/scripts/batch_judge.py @@ -0,0 +1,328 @@ +"""AI judges on each batch of a household set, before the next batch runs. + +Max, 2026-10-22: "we should go slowly since the new microdata could reveal +different pe bugs - we should do ai judges in small batches or real time". + +Each batch of a household set (scripts/household_set.py) runs every model in +every arm of the tool pilot (scripts/tool_pilot_argon.py --board). This script +then hands the batch to PolicyBench's reference adversary +(policybench/reference_adversary.py), arm by arm: + + batch_judge.py payload --board DIR --runs ROOT --stage S --arm none \\ + --batch K --out payload.json + policybench consensus-flags --payload payload.json --output flags.json \\ + --min-models 2 --top-k 3 --min-top 2 --zero-cluster-min-models 2 + batch_judge.py derivations --board DIR --payload payload.json \\ + --flags flags.json --out-dir annotations/ + policybench adversary-prepare --payload payload.json --flags flags.json \\ + --adversary-dir judge/ --annotations-dir annotations/ + scripts/run_reference_adversary_claude.sh judge/ + policybench adversary-collect --adversary-dir claude=judge/ --output-dir out/ + +``payload`` writes one arm's answers as a dashboard payload, each model a +"model" of the board, so the consensus trigger and the adversary run on it +unchanged. Only the no-tools and web arms are judged: a calculator arm answers +with the reference engine itself, and its explanations name it, which the +adversary's blinding refuses. The payload carries what the models saw: the +v2-pilot household prompt (the tool contract's, checked against each record's +sent prompt hash), the v2-pilot output definitions, and the condition they +answered under (from memory, or with web search). + +The trigger's thresholds are the board's, scaled to three models: a cell is +flagged when two of the three models agree on one answer other than the +reference (``min_models`` 2, a zero cluster too), or when two of the three +(all are "top") do. + +``derivations`` writes the engine derivation of each flagged reference: the +pruned computation trace on the household set's own reference system (its +engine and fix modules), in the frozen annotations' CSV format. A trace is the +engine's own record; no model narrates it. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +import pandas as pd + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +sys.path.insert(0, str(ROOT / "scripts")) + +import tool_pilot_argon as pilot # noqa: E402 + +MODELS = ("gpt-6-sol", "claude-opus-5.5", "gemini-4-argon") +JUDGED_ARMS = { + "none": "answering from memory, without tools", + "web": "answering with web search", +} +# The tool contract's household prompt is the dashboard payload's +# ``prompt.tool``; this model is on the tool contract. +PROMPT_MODEL = "gpt-6-sol" +TRACE_MAX_LINES = 400 +YEAR = 2026 + + +def arm_folder(arm: str) -> str: + return arm.replace("+", "-") + + +def load_batch(board: Path, batch: int) -> list[str]: + batches = json.loads((board / "batches.json").read_text())["batches"] + if not 0 <= batch < len(batches): + raise SystemExit(f"{board} has batches 0 to {len(batches) - 1}") + return list(batches[batch]) + + +def prompt_model_for(models) -> str: + """The model whose household prompt a payload shows: the first on the tool + contract (the dashboard's ``prompt.tool``), else the first model.""" + for model in models: + if pilot.model_profile(model, "none").contract == "tool": + return model + return models[0] + + +def household_prompt( + scenario, variables: list[str], prompt_version: str, model: str = PROMPT_MODEL +) -> str: + """The household prompt ``model`` saw, without the arm's notice: the + board's prompt in the run's prompt version.""" + profile = pilot.model_profile(model, "none", prompt_version, False) + prompt, kwargs = pilot.build_initial_request(scenario, variables, profile) + return prompt[: len(prompt) - len(profile.notice)] if profile.notice else prompt + + +def output_definitions(variables, prompt_version: str) -> dict[str, str]: + from policybench.prompts import get_variable_description + + definitions = {} + for variable in sorted(set(variables)): + text = get_variable_description(variable, country="us") + if prompt_version == "v2-pilot": + text = pilot.v2_pilot_text(text) + definitions[variable] = text + return definitions + + +def build_payload( + board_dir: Path, + runs: Path, + stage: str, + arm: str, + batch: int | None = None, + scenarios: list[str] | None = None, + models: tuple[str, ...] = MODELS, +) -> dict: + """One arm's answers on one batch (or on named households), as a + dashboard payload, from ``models`` (the frontier three by default; cheap + models for development, Max 2026-10-23).""" + models = tuple(models) + if len(models) < 2: + raise SystemExit("a payload needs at least two models") + prompt_model = prompt_model_for(models) + if arm not in JUDGED_ARMS: + raise SystemExit(f"only {sorted(JUDGED_ARMS)} are judged, not {arm}") + if (batch is None) == (not scenarios): + raise SystemExit("name a batch or households, not both") + board = pilot.load_board(board_dir) + ids = list(scenarios) if scenarios else load_batch(board_dir, batch) + unknown = sorted(set(ids) - set(board["scenarios"])) + if unknown: + raise SystemExit(f"{board_dir} has no households {unknown}") + scored = set(zip(board["scored"].scenario_id, board["scored"].variable)) + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + records = {} + for model in models: + folder = runs / model / arm_folder(arm) / stage / "households" + found = {} + for sid in ids: + path = folder / f"{sid}.json" + if not path.exists(): + raise SystemExit(f"{path} is missing: the batch has not finished") + found[sid] = json.loads(path.read_text()) + if pilot.records_model(found) != model or pilot.records_arm(found) != arm: + raise SystemExit(f"{folder} does not hold {model} {arm} records") + records[model] = found + versions = {pilot.records_prompt(r)[0] for r in records.values()} + if len(versions) != 1: + raise SystemExit(f"the models ran different prompt versions: {versions}") + prompt_version = versions.pop() + + labels = {model: f"{model}+{arm}" for model in models} + scenarios, predictions, variables_seen = {}, {}, set() + for sid in ids: + scenario = board["scenarios"][sid] + variables = board["variables"][sid] + prompt = household_prompt(scenario, variables, prompt_version, prompt_model) + sent = records[prompt_model][sid].get("sent_board_prompt_sha256") + if sent != hashlib.sha256(prompt.encode()).hexdigest(): + raise SystemExit( + f"{sid}: the rebuilt prompt is not the one {prompt_model} saw" + ) + scenarios[sid] = {"prompt": {"tool": prompt}, "state": scenario.state} + cells = {} + for variable in variables: + variables_seen.add(variable) + cell = {} + for model, label in labels.items(): + record = records[model][sid] + cell[label] = { + "prediction": record["predictions"].get(variable), + "explanation": record["explanations"].get(variable), + "groundTruth": float(reference[(sid, variable)]), + "scored": (sid, variable) in scored, + } + cells[variable] = cell + predictions[sid] = cells + return { + "country": "us", + "source": { + "board": str(board_dir), + "board_reference_sha256": board["meta"].get("reference_csv_sha256"), + "runs": str(runs), + "stage": stage, + "arm": arm, + "batch": batch, + "households": ids, + "prompt_version": prompt_version, + }, + "answerCondition": JUDGED_ARMS[arm], + "models": list(models), + "outputDefinitions": output_definitions(variables_seen, prompt_version), + "modelStats": [{"model": labels[model]} for model in models], + "scenarios": scenarios, + "scenarioPredictions": predictions, + } + + +def render_trace(system, scenario, variable: str) -> str: + """The pruned computation trace of one output on ``system``.""" + import household_set + from policyengine_us import Simulation + + from policybench.case_reference_explanations import ( + _find_target_tree, + _render_trace, + ) + from policybench.ground_truth import ( + _extract_person_value, + _pe_variable_for_output, + _person_index_for_output, + ) + + pe_variable = _pe_variable_for_output(variable, "us") + simulation = Simulation( + tax_benefit_system=system, situation=household_set.situation_for(scenario) + ) + simulation.trace = True + result = simulation.calculate(pe_variable, YEAR) + tree = _find_target_tree(simulation.tracer.trees, pe_variable) + if tree is None: + raise RuntimeError(f"no trace for {pe_variable}") + value = float(_extract_person_value(result, scenario, variable)) + lines = _render_trace(tree) + cut = len(lines) > TRACE_MAX_LINES + person = _person_index_for_output(scenario, variable) + header = [ + f"Output {variable} = {value:g}, computed as the engine variable " + f"{pe_variable} for {YEAR}" + + (f", read for person index {person}." if person is not None else "."), + "Each line is a variable and its value; a variable's inputs are indented " + "beneath it. Zero-valued branches below the first level are omitted.", + "", + ] + footer = [f"... ({len(lines) - TRACE_MAX_LINES} more lines cut)"] if cut else [] + return "\n".join(header + lines[:TRACE_MAX_LINES] + footer) + + +def command_payload(args) -> None: + payload = build_payload( + Path(args.board), + Path(args.runs), + args.stage, + args.arm, + batch=args.batch, + scenarios=args.scenarios, + models=tuple(args.models), + ) + Path(args.out).write_text(json.dumps(payload, indent=1) + "\n") + cells = sum(len(v) for v in payload["scenarioPredictions"].values()) + print(f"{len(payload['scenarios'])} households, {cells} cells, arm {args.arm}") + + +def command_derivations(args) -> None: + board_dir = Path(args.board) + pilot._require_engine(board_dir) + payload = json.loads(Path(args.payload).read_text()) + if payload["source"]["board"] != str(board_dir): + raise SystemExit("the payload is from another board") + flags = json.loads(Path(args.flags).read_text())["flags"] + board = pilot.load_board(board_dir) + system = pilot.build_tool_system(reform=True, snapshot=board_dir) + rows = [] + for flag in flags: + sid, variable = flag["scenario_id"], flag["variable"] + rows.append( + { + "country": "us", + "scenario_id": sid, + "variable": variable, + "reference_value": flag["reference"], + "trace_lines": None, + "explanation": render_trace(system, board["scenarios"][sid], variable), + "error": None, + } + ) + for row in rows: + row["trace_lines"] = row["explanation"].count("\n") + 1 + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + pd.DataFrame( + rows, + columns=[ + "country", + "scenario_id", + "variable", + "reference_value", + "trace_lines", + "explanation", + "error", + ], + ).to_csv(out / "us_case_reference_explanations.csv", index=False) + print(f"{len(rows)} derivations in {out}") + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + sub = parser.add_subparsers(dest="command", required=True) + payload = sub.add_parser("payload", help="one arm's batch as a dashboard payload") + payload.add_argument("--board", required=True) + payload.add_argument("--runs", required=True) + payload.add_argument("--stage", required=True) + payload.add_argument("--arm", required=True, choices=sorted(JUDGED_ARMS)) + payload.add_argument("--batch", type=int) + payload.add_argument( + "--scenarios", nargs="*", help="households, for a board without batches" + ) + payload.add_argument("--out", required=True) + payload.add_argument( + "--models", nargs="+", default=list(MODELS), help="the models to judge" + ) + payload.set_defaults(func=command_payload) + derivations = sub.add_parser("derivations", help="traces of flagged references") + derivations.add_argument("--board", required=True) + derivations.add_argument("--payload", required=True) + derivations.add_argument("--flags", required=True) + derivations.add_argument("--out-dir", required=True) + derivations.set_defaults(func=command_derivations) + args = parser.parse_args(argv) + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/household_set.py b/scripts/household_set.py new file mode 100644 index 00000000..ff385040 --- /dev/null +++ b/scripts/household_set.py @@ -0,0 +1,560 @@ +"""Build a fresh US household set for PolicyBench, in the frozen board's layout. + +Max, 2026-10-22: v2's head-to-head runs on fresh households from the latest +Microcosm build, with references from the latest policyengine-us, served as +one certified pair by the installed policyengine.py bundle. The run goes +slowly, in small batches with AI judges between them, because new microdata +can reach engine paths the frozen 100 households never did. + + household_set.py build --n 100 --seed 20261022 --batch-size 10 --out DIR + household_set.py references --board DIR --out FILE + household_set.py check-frozen --out FILE + +``build`` samples households from the dataset the installed policyengine.py +bundle certifies (:func:`policybench.scenarios.generate_scenarios`, the board's +own sampler and eligibility rules), renames them ``mc_000`` onward so they +never share an id with the frozen board, computes every output on the +installed policyengine-us with the reference conventions +(reference_audit/2026-09-28/fixes/latest_final.py), and writes a board +directory the tool pilot reads with ``--board``: + + scenarios.csv (+ .meta.json) the households, as the frozen board + reference_outputs.csv (+ .meta.json) value and impact_weight per output + reference_exclusions.json an empty exclusion record + batches.json the households in batches, in draw order + +The reference sidecar has the frozen board's shape: one ``engine_upgrade`` +revision naming the engine and pinning every fix module by sha256, so the +pilot's pin check reads either board the same way. + +``references`` recomputes a board's references and reports any output that +differs (the determinism check). ``check-frozen`` is the differential check +of the reference path: on the frozen board's own households it must reproduce +every scored reference, which holds only on the engine those references were +built with (policyengine-us 2.41.0). + +The set is not public. It stays outside the repository until a release +publishes it, so a web-search arm cannot find its answers. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import datetime +import hashlib +import importlib.util +import json +import sys +from pathlib import Path + +import pandas as pd + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +FROZEN_BOARD = ( + ROOT + / "paper/snapshot/20260501/runs/us_full_run_20260612_policyengine_4_16_1_populace" +) +FIX_MODULE = ROOT / "reference_audit/2026-09-28/fixes/latest_final.py" +YEAR = 2026 +# The v1.1 reference run renamed this input for policyengine-us 1.755.4; the +# frozen households still carry the old name (reference_audit/2026-09-28/ +# scripts/sweep.py build_situation). Harmless for a household without it. +RENAME = {"partnership_se_income": "partnership_self_employment_net_earnings"} +EXCLUSIONS_RULE = ( + "An output is excluded from scoring for every model when review finds its " + "reference wrong in law or undetermined by the stated facts. A fresh household " + "set starts with no exclusions; its batch judges propose them." +) + + +def sha256_bytes(raw: bytes) -> str: + return hashlib.sha256(raw).hexdigest() + + +def sha256_file(path: Path) -> str: + return sha256_bytes(Path(path).read_bytes()) + + +def installed_engine() -> str: + from importlib import metadata + + return f"policyengine-us {metadata.version('policyengine-us')}" + + +def frozen_programs() -> list[str]: + meta = json.loads((FROZEN_BOARD / "reference_outputs.csv.meta.json").read_text()) + return list(meta["programs"]) + + +def frozen_fix_pins() -> list[dict]: + """The fix modules the frozen 2.41.0 references were built with.""" + meta = json.loads((FROZEN_BOARD / "reference_outputs.csv.meta.json").read_text()) + revisions = [ + r + for r in meta["revisions"] + if r.get("kind") == "engine_upgrade" + and r.get("engine_version") == "policyengine-us 2.41.0" + ] + if len(revisions) != 1: + raise SystemExit("the frozen board has no single 2.41.0 engine_upgrade") + return revisions[0]["fix_modules"] + + +def fix_module_pins(fix_module: Path = FIX_MODULE) -> list[dict]: + """Each module the reform composes, at its current bytes. + + The set of modules is the frozen board's; each must exist and match the + frozen pin, so a household set's references follow exactly the frozen + conventions. Changing a convention for a new engine is an engine-upgrade + audit, which pins its own modules. + """ + pins, problems = [], [] + for pin in frozen_fix_pins(): + path = ROOT / pin["path"] + if not path.exists(): + problems.append(f"{pin['path']} is missing") + continue + digest = sha256_file(path) + if digest != pin["sha256"]: + problems.append(f"{pin['path']} is not the frozen {pin['sha256'][:12]}") + pins.append({"module": pin["module"], "path": pin["path"], "sha256": digest}) + if problems: + raise SystemExit("; ".join(problems)) + return pins + + +def load_reform(fix_module: Path = FIX_MODULE): + spec = importlib.util.spec_from_file_location("household_set_fixes", fix_module) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module.reform + + +def reference_system(fix_module: Path = FIX_MODULE): + from policyengine_us import CountryTaxBenefitSystem + + return CountryTaxBenefitSystem(reform=load_reform(fix_module)) + + +def situation_for(scenario) -> dict: + situation = scenario.to_pe_household() + for person in situation["people"].values(): + for old, new in RENAME.items(): + if old in person: + person[new] = person.pop(old) + return situation + + +def compute_references(scenarios, programs: list[str], system) -> pd.DataFrame: + """Every output of every household on ``system``: the ground-truth path + (policybench.ground_truth._calculate_ground_truth_us_scalar) with the + reference conventions.""" + from policyengine_us import Simulation + + from policybench.ground_truth import ( + _extract_person_impact_weight, + _extract_person_value, + _impact_weight_variable_for_output, + _pe_variable_for_output, + ) + from policybench.spec import expand_programs_for_scenario + + rows = [] + for scenario in scenarios: + simulation = Simulation( + tax_benefit_system=system, situation=situation_for(scenario) + ) + for variable in expand_programs_for_scenario(programs, scenario): + result = simulation.calculate(_pe_variable_for_output(variable, "us"), YEAR) + weight_variable = _impact_weight_variable_for_output(variable, "us") + weight = None + if weight_variable is not None: + weight = _extract_person_impact_weight( + result, + simulation.calculate(weight_variable, YEAR), + scenario, + variable, + ) + rows.append( + { + "scenario_id": scenario.id, + "variable": variable, + "value": float(_extract_person_value(result, scenario, variable)), + "impact_weight": weight, + } + ) + return pd.DataFrame( + rows, columns=["scenario_id", "variable", "value", "impact_weight"] + ) + + +def rename_scenarios(scenarios, prefix: str) -> list: + """The sampled households under fresh ids, in draw order.""" + return [ + dataclasses.replace(scenario, id=f"{prefix}_{index:03d}") + for index, scenario in enumerate(scenarios) + ] + + +def split_batches(ids: list[str], size: int) -> list[list[str]]: + if size < 1: + raise ValueError("batch size must be at least 1") + return [ids[start : start + size] for start in range(0, len(ids), size)] + + +def load_scenarios(board: Path) -> list: + from policybench.scenarios import scenario_from_dict + + frame = pd.read_csv(Path(board) / "scenarios.csv") + return [scenario_from_dict(json.loads(text)) for text in frame["scenario_json"]] + + +def _utc_now() -> str: + return datetime.datetime.now(datetime.timezone.utc).isoformat() + + +def explicit_dataset(repo: str, filename: str, revision: str) -> tuple[Path, dict]: + """A published dataset named by repo, file and tag (not the bundle's), and + what its release manifest says about the engine it was built for.""" + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError + + path = Path(hf_hub_download(repo, filename, revision=revision, repo_type="dataset")) + record = { + "uri": f"hf://datasets/{repo}/{filename}@{revision}", + "release_manifest": None, + } + try: + manifest_path = hf_hub_download( + repo, + f"releases/{revision}/release_manifest.json", + revision=revision, + repo_type="dataset", + ) + except EntryNotFoundError: + return path, record + manifest = json.loads(Path(manifest_path).read_text()) + build = manifest.get("build") or {} + record["release_manifest"] = { + "sha256": sha256_file(Path(manifest_path)), + "built_with_model_package": build.get("built_with_model_package"), + "built_with_core_package": build.get("built_with_core_package"), + "compatible_model_packages": manifest.get("compatible_model_packages"), + } + return path, record + + +# How a dataset's stated filing status reads as a scenario's. +STATED_FILING_STATUS = { + "SINGLE": "single", + "JOINT": "joint", + "HEAD_OF_HOUSEHOLD": "head_of_household", + "SEPARATE": "separate", + "SURVIVING_SPOUSE": "surviving_spouse", +} + + +def stated_filing_status(path: Path) -> dict[int, str] | None: + """Each tax unit's filing_status_input from a Populace H5 (pandas table + format), or None when the file states none. Engines after 2.2.1 no longer + read the column; they derive filing status from the household.""" + import h5py + + if not h5py.is_hdf5(path): + return None + with h5py.File(path, "r") as handle: + if "tax_unit" not in handle or "table" not in handle["tax_unit"]: + return None + table = handle["tax_unit"]["table"][:] + columns = {} + for field in table.dtype.names: + sub = table.dtype[field] + for name in sub.names or (field,): + columns[name] = table[field][name] if sub.names else table[field] + if "filing_status_input" not in columns or "tax_unit_id" not in columns: + return None + return { + int(unit): (value.decode() if isinstance(value, bytes) else str(value)) + for unit, value in zip(columns["tax_unit_id"], columns["filing_status_input"]) + } + + +def filing_status_mismatches(scenarios, stated: dict[int, str] | None) -> list[dict]: + """Households whose scenario filing status (the engine's) is not the one + the dataset states.""" + if stated is None: + return [] + found = [] + for scenario in scenarios: + unit = scenario.metadata.get("tax_unit_id") + said = stated.get(int(unit)) if unit is not None else None + if said is None: + continue + if STATED_FILING_STATUS.get(said, said.lower()) != scenario.filing_status: + found.append( + { + "scenario_id": scenario.id, + "tax_unit_id": int(unit), + "stated": said, + "engine": scenario.filing_status, + } + ) + return found + + +def command_build(args) -> None: + from policybench.policyengine_runtime import ( + materialize_certified_us_dataset, + policyengine_bundles_for_countries, + ) + from policybench.scenarios import generate_scenarios, scenario_manifest + + out = Path(args.out) + if out.exists() and any(out.iterdir()): + raise SystemExit(f"{out} is not empty; a household set is written once") + named = [args.dataset_repo, args.dataset_file, args.dataset_revision] + if any(named) and not all(named): + raise SystemExit( + "name a dataset by --dataset-repo, --dataset-file and --dataset-revision" + ) + engine = installed_engine() + pins = fix_module_pins(Path(args.fix_module)) + programs = frozen_programs() + bundles = policyengine_bundles_for_countries(["us"]) + bundle = bundles["us"] + if all(named): + # Max, 2026-10-22: "we can still test before that but i'd like v2 to be + # clean". A dataset other than the bundle's makes a test set only. + dataset_path, dataset_record = explicit_dataset(*named) + dataset_label = f"{Path(args.dataset_file).stem}@{args.dataset_revision}" + certified_pair = False + else: + dataset_path = Path(materialize_certified_us_dataset()) + dataset_record = {"uri": bundle.get("default_dataset_uri")} + dataset_label = None + certified = bundle.get("certified_data_artifact_sha256") + if certified and certified != sha256_file(dataset_path): + raise SystemExit( + f"{dataset_path} hashes to {sha256_file(dataset_path)}, not the " + f"bundle's {certified}" + ) + # The bundle certifies its data for its own model pin; the pair is + # certified only when that pin is the installed engine. + certified_pair = bundle.get("bundled_model_version") == bundle.get( + "model_version" + ) + if not certified_pair and args.purpose != "test": + raise SystemExit( + "this dataset and engine are not a certified pair: build it with " + "--purpose test, or install a bundle that certifies the engine" + ) + dataset_sha256 = sha256_file(dataset_path) + + sampled = generate_scenarios( + n=args.n, + seed=args.seed, + country="us", + dataset=str(dataset_path) if all(named) else None, + dataset_label=dataset_label, + ) + scenarios = rename_scenarios(sampled, args.id_prefix) + mismatches = filing_status_mismatches(scenarios, stated_filing_status(dataset_path)) + references = compute_references( + scenarios, programs, reference_system(Path(args.fix_module)) + ) + + out.mkdir(parents=True, exist_ok=True) + manifest = scenario_manifest(scenarios) + manifest.to_csv(out / "scenarios.csv", index=False) + references.to_csv(out / "reference_outputs.csv", index=False) + batches = split_batches([s.id for s in scenarios], args.batch_size) + (out / "batches.json").write_text( + json.dumps({"batch_size": args.batch_size, "batches": batches}, indent=1) + "\n" + ) + (out / "reference_exclusions.json").write_text( + json.dumps( + { + "schema_version": 1, + "rule": EXCLUSIONS_RULE, + "derivation": f"Built {_utc_now()[:10]} with no exclusions.", + "exclusions": [], + }, + indent=1, + ) + + "\n" + ) + generated = _utc_now() + common = { + "country": "us", + "metadata_version": 1, + "board_kind": "household_set", + "purpose": args.purpose, + "certified_pair": certified_pair, + "generated_at_utc": generated, + "num_scenarios": len(scenarios), + "requested_num_scenarios": args.n, + "seed": args.seed, + "id_prefix": args.id_prefix, + "program_set": "headline", + "programs": programs, + "policyengine_bundles": bundles, + "dataset": { + **dataset_record, + "path_name": dataset_path.name, + "sha256": dataset_sha256, + "bundle_default_dataset_uri": bundle.get("default_dataset_uri"), + }, + # The dataset states a filing status the engine no longer reads; these + # households' scenarios carry the engine's, as their references do. + "filing_status_mismatches": mismatches, + "source_households": [ + {"scenario_id": s.id, **{k: s.metadata[k] for k in sorted(s.metadata)}} + for s in scenarios + ], + } + (out / "scenarios.csv.meta.json").write_text( + json.dumps( + { + **common, + "task": "scenario_manifest", + "scenarios_csv_sha256": sha256_file(out / "scenarios.csv"), + }, + indent=1, + ) + + "\n" + ) + (out / "reference_outputs.csv.meta.json").write_text( + json.dumps( + { + **common, + "task": "reference_outputs", + "reference_csv_sha256": sha256_file(out / "reference_outputs.csv"), + "scenarios_csv_sha256": sha256_file(out / "scenarios.csv"), + "revisions": [ + { + "kind": "engine_upgrade", + "engine_version": engine, + "date": generated[:10], + "builder": "scripts/household_set.py build", + "fix_module": { + "path": str( + Path(args.fix_module).resolve().relative_to(ROOT) + ), + "sha256": sha256_file(Path(args.fix_module)), + }, + "fix_modules": pins, + "outputs": len(references), + } + ], + }, + indent=1, + ) + + "\n" + ) + print( + f"{len(scenarios)} households, {len(references)} outputs on {engine}; " + f"{len(batches)} batches of up to {args.batch_size} in {out}; purpose " + f"{args.purpose}, certified pair {certified_pair}; " + f"{len(mismatches)} filing status mismatches" + ) + + +def differences( + expected: pd.DataFrame, computed: pd.DataFrame, tolerance: float +) -> pd.DataFrame: + """Outputs whose computed value is missing or differs beyond ``tolerance``.""" + merged = expected.merge( + computed[["scenario_id", "variable", "value"]], + on=["scenario_id", "variable"], + how="left", + suffixes=("", "_computed"), + ) + merged["delta"] = (merged["value_computed"] - merged["value"]).abs() + return merged[merged["value_computed"].isna() | (merged["delta"] > tolerance)] + + +def command_references(args) -> None: + board = Path(args.board) + expected = pd.read_csv(board / "reference_outputs.csv") + meta = json.loads((board / "reference_outputs.csv.meta.json").read_text()) + computed = compute_references( + load_scenarios(board), meta["programs"], reference_system(Path(args.fix_module)) + ) + diff = differences(expected, computed, args.tolerance) + computed.to_csv(args.out, index=False) + print( + f"{len(computed)} outputs recomputed on {installed_engine()}; " + f"{len(diff)} differ" + ) + if len(diff): + print(diff.to_string(index=False)) + raise SystemExit(1) + + +def command_check_frozen(args) -> None: + """The reference path on the frozen board must reproduce every scored output.""" + from policybench.reference_exclusions import ( + load_reference_exclusions, + split_reference, + ) + + engine = installed_engine() + if engine != "policyengine-us 2.41.0": + raise SystemExit(f"{engine} is installed; the frozen references are 2.41.0's") + fix_module_pins(FIX_MODULE) + reference = pd.read_csv(FROZEN_BOARD / "reference_outputs.csv") + scored, _ = split_reference(reference, load_reference_exclusions(FROZEN_BOARD)) + computed = compute_references( + load_scenarios(FROZEN_BOARD), frozen_programs(), reference_system(FIX_MODULE) + ) + computed.to_csv(args.out, index=False) + diff = differences(scored, computed, args.tolerance) + print( + f"{len(scored)} scored outputs; {len(scored) - len(diff)} reproduced within " + f"{args.tolerance} on {engine}" + ) + if len(diff): + print(diff.to_string(index=False)) + raise SystemExit(1) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + sub = parser.add_subparsers(dest="command", required=True) + build = sub.add_parser("build", help="sample a household set and its references") + build.add_argument("--n", type=int, required=True) + build.add_argument("--seed", type=int, required=True) + build.add_argument("--batch-size", type=int, default=10) + build.add_argument("--id-prefix", default="mc") + build.add_argument("--fix-module", default=str(FIX_MODULE)) + build.add_argument("--out", required=True) + build.add_argument( + "--purpose", + choices=["v2", "test"], + default="v2", + help="v2 needs a certified dataset-engine pair; test takes any", + ) + build.add_argument("--dataset-repo", help="a published dataset, for a test set") + build.add_argument("--dataset-file") + build.add_argument("--dataset-revision") + build.set_defaults(func=command_build) + refs = sub.add_parser("references", help="recompute a board's references") + refs.add_argument("--board", required=True) + refs.add_argument("--fix-module", default=str(FIX_MODULE)) + refs.add_argument("--tolerance", type=float, default=1e-6) + refs.add_argument("--out", required=True) + refs.set_defaults(func=command_references) + frozen = sub.add_parser("check-frozen", help="reproduce the frozen references") + frozen.add_argument("--tolerance", type=float, default=1e-6) + frozen.add_argument("--out", required=True) + frozen.set_defaults(func=command_check_frozen) + args = parser.parse_args(argv) + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_reference_adversary_claude.sh b/scripts/run_reference_adversary_claude.sh index 69818c67..5c680040 100755 --- a/scripts/run_reference_adversary_claude.sh +++ b/scripts/run_reference_adversary_claude.sh @@ -58,6 +58,17 @@ # judge runs. A call the API refuses with 401, 403 or 429 stops the run: no # further judge starts, and the run resumes where it stopped. # +# API key (Max, 2026-10-22: "in case youre running any claude api stuff please +# use the key on my clipboard"): JUDGE_API_KEY_SECRET names an agent-secret +# entry (account JUDGE_API_KEY_ACCOUNT, default $USER; AUDIT_SECRET_BIN +# replaces agent-secret) holding an Anthropic API key. The judges then bill +# that key and nothing else: it cannot be combined with a lane, a lane token +# or the desktop login; every claude call runs in a fresh empty config +# directory, removed when the run ends; the key joins only the claude calls' +# allowlisted environment; and `claude auth status` must report a first-party +# api_key login in that directory. The sidecars record the mode and the +# secret's name, never the key. +# # Provenance: beside stage1.json and verdict.json the runner writes # stage1.meta.json and verdict.meta.json with the runner, the stage, the model # requested and the models the CLI reports, the CLI version, the effort, the @@ -165,7 +176,34 @@ for base in {home, pwd.getpwuid(os.getuid()).pw_dir}: print("other") PY } -if [ "$DESKTOP_OPT_IN" = 1 ]; then +API_SECRET="${JUDGE_API_KEY_SECRET:-}" +API_KEY="" +if [ -n "$API_SECRET" ]; then + [ "$DESKTOP_OPT_IN" = 0 ] || { + echo "JUDGE_API_KEY_SECRET and JUDGE_ALLOW_DESKTOP_LOGIN=1 are exclusive; refusing to start" >&2 + exit 1 + } + [ -z "${CLAUDE_CODE_OAUTH_TOKEN:-}" ] || { + echo "JUDGE_API_KEY_SECRET takes no lane token: unset CLAUDE_CODE_OAUTH_TOKEN; refusing to start" >&2 + exit 1 + } + [ -z "${AUDIT_ACCOUNT:-}" ] || { + echo "JUDGE_API_KEY_SECRET bills a key, not an account: unset AUDIT_ACCOUNT; refusing to start" >&2 + exit 1 + } + [ -z "${CLAUDE_CONFIG_DIR:-}" ] || { + echo "JUDGE_API_KEY_SECRET runs in its own empty config directory: unset CLAUDE_CONFIG_DIR; refusing to start" >&2 + exit 1 + } + API_KEY=$("${AUDIT_SECRET_BIN:-agent-secret}" get "$API_SECRET" "${JUDGE_API_KEY_ACCOUNT:-$USER}" 2>/dev/null) + [ -n "$API_KEY" ] || { + echo "could not read the API key $API_SECRET; refusing to start" >&2 + exit 1 + } + CONFIG_DIR=$(mktemp -d "${TMPDIR:-/tmp}/pb-adversary-config.XXXXXX") || exit 1 + CONFIG_DIR="$(cd "$CONFIG_DIR" && pwd -P)" + trap 'rm -rf "$CONFIG_DIR"' EXIT +elif [ "$DESKTOP_OPT_IN" = 1 ]; then [ -z "${CLAUDE_CODE_OAUTH_TOKEN:-}" ] || { echo "JUDGE_ALLOW_DESKTOP_LOGIN=1 takes no lane token: unset CLAUDE_CODE_OAUTH_TOKEN; refusing to start" >&2 exit 1 @@ -218,6 +256,7 @@ while IFS= read -r line; do [ -n "$line" ] && CHILD_ENV+=("$line"); done <&2 + --desktop-opt-in "$DESKTOP_OPT_IN" --api-key-secret "$API_SECRET") || { + echo "refusing to start: the judges must bill a lane's subscription login, the opted-in desktop login or the declared API key" >&2 exit 1 } DECLARED="${AUDIT_ACCOUNT:-}" [ "$DESKTOP_OPT_IN" = 1 ] && DECLARED="$DESKTOP_DECLARED" +[ -n "$API_SECRET" ] && DECLARED="API key $API_SECRET (JUDGE_API_KEY_SECRET, Max 2026-10-22)" CLI_VERSION=$(cd / && claude_child --version /dev/null | head -n 1) valid() { [ -s "$2" ] && adversary validate --schema "$1" --file "$2" >/dev/null 2>&1; } diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 97d880e7..f09e0b6f 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -48,9 +48,11 @@ import argparse import concurrent.futures +import contextlib import copy import decimal import difflib +import functools import hashlib import importlib.metadata import importlib.util @@ -152,13 +154,215 @@ # 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") +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", } +# 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; +# - 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 +# 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 (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, LABELS, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( + "names", + "labels", + "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(), + "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 " + "(policyengine-us 2.41.0), an open-source rules engine for US federal and " + "state taxes and benefits, on a household situation you build from the facts " + "above. 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. Its result says which names it read, which " + "inputs it moved to the entity they belong to, and which inputs replace a " + "formula; ask it to `explain` up to three variables to see how it computed " + "them, with the parameters and laws it used. You may call it in up to 8 " + "turns. Then call `submit_outputs` exactly once with your final answers, as " + "instructed above." +) +JSON_TOOL_NOTICE_V2 = TOOL_NOTICE_V2.replace( + "Then call `submit_outputs` exactly once with your final answers, as " + "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 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 + 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 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 += ( + " 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 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") + 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 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" + ) + 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. +_STATED_ENGINE = {"version": REFERENCE_ENGINE} + + +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}", + f"policyengine-us {_STATED_ENGINE['version']}", + ) + + # A web arm's request timeout. A search-heavy turn can outlast the board's # (300 s for Gemini 4 Argon, the harness default for its card; its web-only arm # never finished scenario_005 within it in six attempts on 2026-10-21), so the @@ -452,12 +656,130 @@ def refuse_policyengine(label: str, value) -> None: ) -def calculate_tool_schema() -> dict: +CALCULATE_DESCRIPTION_V2 = CALCULATE_DESCRIPTION.replace( + "An unknown name, a variable set on the wrong entity, or a malformed " + "situation returns an error with suggestions instead of results.", + "A name may also be written with spaces (as the facts are labeled) or be one " + "of the output keys you are asked for; the result's `notes` say how each was " + "read. An input set on the wrong entity is moved to the one unit it belongs " + "to, and the notes say so; where there is no single such unit, or a name is " + "unknown, or the situation is malformed, it returns an error with " + "suggestions instead of results. Setting a variable the engine computes " + "replaces its formula, and the notes say so; a contribution set as the " + "amount contributed is read as the amount asked for before limits " + "(`..._desired`). `explain` lists up to 3 variables to trace: each comes " + "back with the steps that produced it and the parameters, with their " + "values and sources, that it used.", +) +MAX_EXPLAIN = 3 + + +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 `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 " + '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. 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." +) + + +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"] + 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 + ) + 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( + "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 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 " + "holding its people; where no single unit holds them, it " + "returns an error." + ) + 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" ' + f"(the trace, up to {EXPLAIN_MAX_LINES} steps and " + f"{EXPLAIN_MAX_PARAMETERS} parameters)." + ), + } + return schema return { "type": "function", "function": { "name": TOOL_NAME, - "description": CALCULATE_DESCRIPTION, + "description": engine_text(CALCULATE_DESCRIPTION), "parameters": { "type": "object", "properties": { @@ -499,9 +821,49 @@ class VariableInfo: label: str | None = None definition_period: str | None = None possible_values: tuple[str, ...] | None = None + # Whether the engine computes the variable (a formula or a sum of other + # variables) and no household fact the benchmark states sets it: setting + # it replaces what the engine would compute. + computed: bool = False + +# Inputs the benchmark's reference situations set besides its promptable +# facts (policybench.scenarios.Scenario.to_pe_household), and the convention +# switches (take-up, filing) any household may set. +BASE_STATED_INPUTS = frozenset( + { + "age", + "employment_income", + "self_employment_income", + "long_term_capital_gains", + "partnership_s_corp_income", + "is_tax_unit_head", + "is_tax_unit_spouse", + "state_code", + "weekly_hours_worked_before_lsr", + } +) +CONVENTION_PREFIXES = ("takes_up_", "would_") -def build_registry(system) -> dict[str, 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 + + 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) + + +def build_registry( + system, stated: frozenset[str] | None = None +) -> dict[str, VariableInfo]: + """The engine's variables. ``stated`` (default :func:`stated_inputs`) marks + which ones a stated fact sets; the rest that the engine computes are + ``computed``.""" + if stated is None: + stated = stated_inputs() registry = {} for name, variable in system.variables.items(): value_type = getattr(variable.value_type, "__name__", str(variable.value_type)) @@ -514,6 +876,13 @@ def build_registry(system) -> dict[str, VariableInfo]: label=variable.label, definition_period=str(variable.definition_period), possible_values=possible, + computed=bool( + getattr(variable, "formulas", None) + or getattr(variable, "adds", None) + or getattr(variable, "subtracts", None) + ) + and name not in stated + and not name.startswith(CONVENTION_PREFIXES), ) return registry @@ -555,6 +924,44 @@ def _type_problem(value, info: VariableInfo) -> str | None: return None +def _normalize_value( + where: str, + name: str, + value, + info: VariableInfo, + period: str, + problems: list[str], +) -> dict | None: + """One input's {period: value}, or None after recording why not.""" + if isinstance(value, dict): + if not value: + problems.append( + f"{where}: {name} has no values; set it as {{period: value}}." + ) + return None + periods = value + else: + periods = {period: value} + normalized = {} + for key, item in periods.items(): + if not isinstance(key, str) or not PERIOD_RE.match(key): + problems.append( + f"{where}: {name} has period {key!r}; use a year such as " + '"2026" or a month such as "2026-01".' + ) + continue + scalar = _scalar_problem(item) + if scalar is not None: + problems.append(f"{where}: {name} for {key} {scalar}.") + continue + typed = _type_problem(item, info) + if typed is not None: + problems.append(f"{where}: {name} {typed}.") + continue + normalized[key] = item + return normalized + + def _normalize_inputs( where: str, entity: str, @@ -563,11 +970,23 @@ def _normalize_inputs( period: str, problems: list[str], allow_members: bool, + v2: dict | None = None, ) -> dict: + """Each input under its engine name. ``v2`` (tool version 2) collects the + names it read, the inputs it redirected or flagged, and the inputs set on + another entity, which normalize_request then places.""" out: dict[str, Any] = {} 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: + resolved, note = read_fact_name(name, registry, v2["features"]) + 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}." @@ -575,52 +994,62 @@ def _normalize_inputs( ) continue info = registry[name] + if v2 is not None and OVERRIDES in v2["features"]: + desired = f"{name}{DESIRED_SUFFIX}" + if ( + info.computed + and desired in registry + and registry[desired].entity == info.entity + ): + v2["notes"].append( + f"{where}: set {desired} (the amount asked for, before " + f"limits) for {name}, which the engine computes from it" + ) + name, info = desired, registry[desired] + elif info.computed: + v2["notes"].append( + f"{where}: {name} is computed by the engine; the value set " + "replaces its formula" + ) if info.entity != entity: + if v2 is not None and PLACEMENT in v2["features"]: + v2["misplaced"].append((where, entity, name, value)) + continue problems.append( f"{where}: {name} is a {info.entity} variable; set it under " f"{ENTITY_GROUP.get(info.entity, info.entity)}." ) continue - if isinstance(value, dict): - if not value: - problems.append( - f"{where}: {name} has no values; set it as {{period: value}}." - ) - continue - periods = value - else: - periods = {period: value} - normalized = {} - for key, item in periods.items(): - if not isinstance(key, str) or not PERIOD_RE.match(key): - problems.append( - f"{where}: {name} has period {key!r}; use a year such as " - '"2026" or a month such as "2026-01".' - ) - continue - scalar = _scalar_problem(item) - if scalar is not None: - problems.append(f"{where}: {name} for {key} {scalar}.") - continue - typed = _type_problem(item, info) - if typed is not None: - problems.append(f"{where}: {name} {typed}.") - continue - normalized[key] = item + normalized = _normalize_value(where, name, value, info, period, problems) + if normalized is None: + continue + if v2 is not None and name in out and out[name] != normalized: + problems.append(f"{where}: {name} is set twice with different values") + continue out[name] = normalized return out def normalize_request( - args, registry: Mapping[str, VariableInfo] + 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). Accepts only plain JSON data. Every input and output name must be a variable of the engine, on its own entity; nothing else in the situation - (axes, reforms, parameters) is accepted. + (axes, reforms, parameters) is accepted. Tool version 2 also reads a name + written with spaces or a benchmark output key, redirects a contribution to + its ``_desired`` input, flags an input that replaces a formula, places an + input set on another entity where exactly one unit holds it, and takes + ``explain``; the request's ``notes`` say what it did. """ problems: list[str] = [] + 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, { @@ -643,7 +1072,11 @@ def fail() -> tuple[None, dict]: if not isinstance(args, dict): problems.append("arguments must be an object") return fail() - unknown = sorted(set(args) - {"situation", "variables", "period"}) + 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))}") @@ -658,6 +1091,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, period + ) + if args["situation"] is None: + return fail() situation = args.get("situation") if isinstance(situation, str): if len(situation) > MAX_SITUATION_CHARS: @@ -702,7 +1142,7 @@ def fail() -> tuple[None, dict]: problems.append(f"people.{name} must be an object of inputs") continue normalized_people[name] = _normalize_inputs( - f"people.{name}", "person", data, registry, period, problems, False + f"people.{name}", "person", data, registry, period, problems, False, v2 ) normalized["people"] = normalized_people @@ -746,7 +1186,9 @@ def fail() -> tuple[None, dict]: seen[member] = unit_name normalized_unit = {"members": list(members)} normalized_unit.update( - _normalize_inputs(where, entity, unit, registry, period, problems, True) + _normalize_inputs( + where, entity, unit, registry, period, problems, True, v2 + ) ) normalized_units[unit_name] = normalized_unit normalized[group] = normalized_units @@ -763,18 +1205,556 @@ 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 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 ( + 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}") + name = resolved if name not in registry: problems.append( f"variables: unknown variable {name!r}." + _suggest(name, registry) ) elif name not in requested: requested.append(name) + explain = [] + if EXPLAIN in features and args.get("explain") is not None: + names = args["explain"] + if isinstance(names, str): + names = [names] + if not isinstance(names, list) or not all(isinstance(n, str) for n in names): + problems.append("explain must be a list of variable names") + names = [] + elif len(names) > MAX_EXPLAIN: + problems.append(f"explain at most {MAX_EXPLAIN} variables per call") + for name in names: + 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}") + name = resolved + if name not in registry: + problems.append( + f"explain: unknown variable {name!r}." + _suggest(name, registry) + ) + elif name not in explain: + explain.append(name) + + detail = None + 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() - return {"situation": normalized, "variables": requested, "period": period}, None + request = {"situation": normalized, "variables": requested, "period": period} + if v2 is not None: + request["explain"] = explain + request["notes"] = v2["notes"] + if detail is not None: + request["detail"] = detail + return request, None + + +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 _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_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( + 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. 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") + 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, 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 + 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 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 + if name in roles: + problems.append(f"people[{index}]: {name!r} is listed twice") + 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] + 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] + 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 = 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 "") + ) + 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 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 + 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: + # 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" + ) + continue + target[name] = value + return None if problems else situation + + +DESIRED_SUFFIX = "_desired" +_LABEL_SEPARATORS = re.compile(r"[\s\-]+") + + +def output_reading(name: str) -> tuple[str, str] | None: + """A benchmark output key's engine variable and how the output reads it, + or None when the name is not an output key.""" + from policybench.ground_truth import ( + _aggregation_for_output, + _pe_variable_for_output, + ) + from policybench.spec import find_output_spec, parse_person_output + + person = parse_person_output(name) + if person is None and find_output_spec(name, country="us") is None: + return None + engine = _pe_variable_for_output(name, "us") + if engine == name: + return None + aggregation = _aggregation_for_output(name, "us") + if person is not None: + how = f"{engine} for the person the key names ({person[0]})" + elif aggregation == "any_positive": + how = f"1 when {engine} is above 0, else 0" + else: + how = f"{engine}, summed over the household's units" + 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()) + + +# 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, 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_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 (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: + 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 + 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]: + """Tool version 2's reading of a name the engine does not know: the name + with spaces or hyphens as underscores, or (for outputs) a benchmark output + key. Returns (engine name, note) or (None, None).""" + if name in registry: + return name, None + written = _LABEL_SEPARATORS.sub("_", name.strip().lower()) + if written in registry: + return written, f"read {name!r} as {written}" + if output_names: + reading = output_reading(name) + if reading is not None and reading[0] in registry: + return reading[0], f"{name} is {reading[1]}; computed {reading[0]}" + return None, None + + +def _place_misplaced( + v2: dict, + normalized: dict, + situation: dict, + registry: Mapping[str, VariableInfo], + period: str, + problems: list[str], +) -> None: + """Move each input set on another entity to the one unit that holds the + people it was set for: a group variable set on a person goes to that + person's unit (a group left out of the situation becomes one unit of + everyone, as the engine reads it), a person variable set on a one-person + unit goes to that person, and a group variable set on another group goes + to the unit holding all of its members. Anything else stays an error.""" + for where, source_entity, name, value in v2["misplaced"]: + info = registry[name] + if source_entity == "person": + persons = [where.split(".", 1)[1]] + else: + group, unit = where.split(".", 1) + persons = list(normalized.get(group, {}).get(unit, {}).get("members", [])) + target_group = ENTITY_GROUP.get(info.entity, info.entity) + if info.entity == "person": + if len(persons) != 1: + problems.append( + f"{where}: {name} is a person variable; set it on one person " + "under people." + ) + continue + destination = normalized["people"].get(persons[0]) + to = f"people.{persons[0]}" + else: + if target_group not in situation and target_group not in normalized: + normalized[target_group] = { + "all": {"members": list(normalized["people"])} + } + units = normalized.get(target_group, {}) + holding = [ + unit + for unit, data in units.items() + if set(persons) <= set(data.get("members", [])) + ] + if len(holding) != 1: + problems.append( + f"{where}: {name} is a {info.entity} variable; set it under " + f"{target_group} (no single unit there holds " + f"{', '.join(persons) or 'its members'})." + ) + continue + destination = units[holding[0]] + to = f"{target_group}.{holding[0]}" + if destination is None: + problems.append(f"{where}: {name} has no place to go") + continue + normalized_value = _normalize_value(to, name, value, info, period, problems) + if normalized_value is None: + continue + if name in destination and destination[name] != normalized_value: + problems.append(f"{to}: {name} is set twice with different values") + continue + destination[name] = normalized_value + v2["notes"].append(f"moved {name} from {where} to {to}") def _jsonable(value): @@ -823,47 +1803,212 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic "period": request["period"], "values": {i: _jsonable(v) for i, v in zip(ids, list(array))}, } - return {"results": results} + out = {"results": results} + if "notes" in request: + out["notes"] = list(request["notes"]) + if request.get("explain"): + out["explanations"] = { + 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, the first BRIEF_INPUTS), 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 the trace', + } + + +EXPLAIN_MAX_LINES = 60 +EXPLAIN_MAX_DEPTH = 6 +EXPLAIN_MAX_PARAMETERS = 15 +# Switches every variable has; never part of a reason. +_QUIET_PARAMETERS = ("gov.abolitions.",) + + +def _trace_value(value) -> Any: + if hasattr(value, "decode_to_str"): + value = value.decode_to_str() + try: + items = list(value) + except TypeError: + return _jsonable(value) + items = [_jsonable(v) for v in items] + return items[0] if len(items) == 1 else items + + +def _all_zero(value) -> bool: + try: + return all((not v) if not isinstance(v, str) else False for v in list(value)) + except TypeError: + return not value + + +def parameter_sources(system, name: str, limit: int = 2) -> list[dict]: + """A parameter's cited sources: its own ``reference`` metadata, or the + nearest ancestor's.""" + node = system.parameters + chain = [node] + for part in name.split("."): + children = getattr(node, "children", None) + if not isinstance(children, dict) or part not in children: + break + node = children[part] + chain.append(node) + for node in reversed(chain): + refs = (getattr(node, "metadata", None) or {}).get("reference") + if refs: + refs = refs if isinstance(refs, list) else [refs] + out = [] + for ref in refs[:limit]: + if isinstance(ref, dict): + out.append({k: ref[k] for k in ("title", "href") if ref.get(k)}) + else: + out.append({"href": str(ref)}) + return out + return [] + + +def explain_variable(system, request: dict, name: str, simulation_cls=None) -> dict: + """How the engine computed ``name`` for the request's household: the + computation steps (each variable, its period and value, with its inputs + indented beneath it; zero-valued branches below the first level left out), + then the parameters it read, each with its value and cited sources.""" + if simulation_cls is None: + from policyengine_us import Simulation as simulation_cls + try: + simulation = simulation_cls( + tax_benefit_system=system, situation=copy.deepcopy(request["situation"]) + ) + simulation.trace = True + simulation.calculate(name, request["period"]) + except Exception as error: + return {"error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS]} + + def find(nodes): + for node in nodes: + if node.name == name: + return node + hit = find(node.children) + if hit is not None: + return hit + return None + + root = find(simulation.tracer.trees) + if root is None: + return {"error": f"no trace for {name}"} + steps: list[str] = [] + parameters: dict[str, Any] = {} + + def walk(node, depth): + if len(steps) >= EXPLAIN_MAX_LINES or depth > EXPLAIN_MAX_DEPTH: + return + if depth > 1 and _all_zero(node.value): + return + steps.append( + " " * depth + f"{node.name} ({node.period}) = {_trace_value(node.value)}" + ) + for parameter in node.parameters: + if parameter.name.startswith(_QUIET_PARAMETERS): + continue + parameters.setdefault(parameter.name, _trace_value(parameter.value)) + for child in node.children: + walk(child, depth + 1) + + walk(root, 0) + return { + "steps": steps, + "parameters": [ + {"name": key, "value": value, "sources": parameter_sources(system, key)} + for key, value in list(parameters.items())[:EXPLAIN_MAX_PARAMETERS] + ], + } + + +def _board_upgrade(snapshot: Path) -> dict: + """A board's last engine_upgrade revision: the engine its references are + on and the fix modules it pins. The frozen board's is 2.41.0's; a household + set (scripts/household_set.py) records one.""" + meta = json.loads((snapshot / "reference_outputs.csv.meta.json").read_text()) + revisions = [r for r in meta["revisions"] if r.get("kind") == "engine_upgrade"] + if not revisions: + raise SystemExit(f"{snapshot} records no engine_upgrade revision") + return revisions[-1] + + +def board_engine(snapshot: Path = SNAPSHOT) -> str: + """The policyengine-us version a board's references were computed on.""" + return _board_upgrade(snapshot)["engine_version"].removeprefix("policyengine-us ") + + +def is_household_set(snapshot: Path) -> bool: + meta = json.loads((snapshot / "reference_outputs.csv.meta.json").read_text()) + return meta.get("board_kind") == "household_set" def pinned_fix_modules(snapshot: Path = SNAPSHOT) -> list[dict]: - """The fix modules the 2.41.0 references were built with, as the sidecar's + """The fix modules a board's references were built with, as its sidecar's engine_upgrade revision pins them (path and sha256).""" - meta = json.loads((snapshot / "reference_outputs.csv.meta.json").read_text()) - revisions = [ - r - for r in meta["revisions"] - if r.get("kind") == "engine_upgrade" - and r.get("engine_version") == f"policyengine-us {REFERENCE_ENGINE}" - ] - if len(revisions) != 1: - raise SystemExit(f"expected one {REFERENCE_ENGINE} engine_upgrade revision") - return revisions[0]["fix_modules"] + return _board_upgrade(snapshot)["fix_modules"] + + +def board_fix_module_sha256(snapshot: Path = SNAPSHOT) -> str: + """The pinned latest_final.py: the constant for the frozen board, the + sidecar's record for a household set.""" + if not is_household_set(snapshot): + return FIX_MODULE_SHA256 + return _board_upgrade(snapshot)["fix_module"]["sha256"] -def verify_fix_pins(fixes_dir: Path = FIX_MODULE.parent) -> None: +def verify_fix_pins( + fixes_dir: Path = FIX_MODULE.parent, snapshot: Path = SNAPSHOT +) -> None: """Every module latest_final.py composes must be the pinned bytes. Each pin is checked at its recorded path; a module the reform reads from its own directory (the sales tax tables) is also checked there. """ problems = [] - for pin in pinned_fix_modules(): + for pin in pinned_fix_modules(snapshot): candidates = {ROOT / pin["path"], fixes_dir / pin["module"]} for path in sorted(candidates): if path.exists() and _sha256(path) != pin["sha256"]: problems.append(f"{path} does not match the pinned {pin['sha256']}") if not any(path.exists() for path in candidates): problems.append(f"{pin['module']} is missing") - if _sha256(FIX_MODULE) != FIX_MODULE_SHA256: - problems.append(f"{FIX_MODULE} does not match {FIX_MODULE_SHA256}") + pinned = board_fix_module_sha256(snapshot) + if _sha256(FIX_MODULE) != pinned: + problems.append(f"{FIX_MODULE} does not match {pinned}") if problems: raise SystemExit("; ".join(problems)) -def load_fix_reform(path: Path = FIX_MODULE): +def load_fix_reform(path: Path = FIX_MODULE, snapshot: Path = SNAPSHOT): if path == FIX_MODULE: - verify_fix_pins() + verify_fix_pins(snapshot=snapshot) spec = importlib.util.spec_from_file_location("tool_pilot_latest_final", path) module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module @@ -871,18 +2016,61 @@ def load_fix_reform(path: Path = FIX_MODULE): return module.reform -def build_tool_system(reform: bool = True, path: Path = FIX_MODULE): +def build_tool_system( + reform: bool = True, path: Path = FIX_MODULE, snapshot: Path = SNAPSHOT +): from policyengine_us import CountryTaxBenefitSystem if not reform: return CountryTaxBenefitSystem() - return CountryTaxBenefitSystem(reform=load_fix_reform(path)) + 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 _worker_main(conn, fix_path: str, reform: bool) -> None: + +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" - system = build_tool_system(reform=reform, path=Path(fix_path)) + snapshot = Path(board) if board else SNAPSHOT + system = build_tool_system(reform=reform, path=Path(fix_path), snapshot=snapshot) registry = build_registry(system) conn.send(("ready", None)) while True: @@ -892,7 +2080,7 @@ def _worker_main(conn, fix_path: str, reform: bool) -> None: return if args is None: return - request, error = normalize_request(args, registry) + request, error = worker_request(args, registry) if error is not None: conn.send(("ok", error)) continue @@ -910,9 +2098,21 @@ def __init__( reform: bool = True, worker_target: Callable | None = None, 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. + self._board = None if board is None or board == SNAPSHOT else str(board) self._reform = reform self._target = worker_target or _worker_main self.timeout = timeout @@ -932,11 +2132,10 @@ def __init__( def _start(self): parent, child = self._ctx.Pipe() - process = self._ctx.Process( - target=self._target, - args=(child, self._fix_path, self._reform), - daemon=True, - ) + args = (child, self._fix_path, self._reform) + if self._board is not None: + args += (self._board,) + process = self._ctx.Process(target=self._target, args=args, daemon=True) process.start() return process, parent @@ -976,6 +2175,8 @@ def call(self, args) -> tuple[dict, float]: time.perf_counter() - started, ) process, conn = worker + if self._tool_features: + args = tool_call_frame(self._tool_features, args) try: conn.send(args) if conn.poll(self.timeout): @@ -1131,13 +2332,27 @@ def notice(self) -> str: notice = WEB_NOTICE if tool else JSON_WEB_NOTICE elif self.arm == "web+pe": 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 TOOL_NOTICE if tool else JSON_TOOL_NOTICE - return notice + (BENCHMARK_SEARCH_RULE if self.benchmark_search_rule else "") + return engine_text(TOOL_NOTICE if tool else JSON_TOOL_NOTICE) + return engine_text(notice) + ( + BENCHMARK_SEARCH_RULE if self.benchmark_search_rule else "" + ) @property def uses_calculator(self) -> bool: - return self.arm in ("pe", "web+pe") + 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 2 if self.tool_features == V2_FEATURES else 1 @property def uses_web(self) -> bool: @@ -1155,7 +2370,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: @@ -1181,8 +2396,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 " @@ -1816,7 +3034,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() + 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 [] ) @@ -1857,6 +3075,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", } @@ -2366,28 +3585,38 @@ def _versions() -> dict: return out -def _require_engine() -> None: +def board_path(args) -> Path: + """The board a command reads: --board, or the frozen snapshot.""" + return Path(getattr(args, "board", None) or SNAPSHOT) + + +def _require_engine(board: Path = SNAPSHOT) -> None: version = _versions()["policyengine-us"] - if version != REFERENCE_ENGINE: + pinned = board_engine(board) + if version != pinned: + where = ( + f"the reference engine of household set {board}" + if is_household_set(board) + else "the reference engine of release dashboard-data-20261014" + ) raise SystemExit( - f"policyengine-us {version} is installed; the pilot pins " - f"{REFERENCE_ENGINE}, the reference engine of release " - "dashboard-data-20261014." + f"policyengine-us {version} is installed; the pilot pins {pinned}, {where}." ) def command_oracle(args) -> None: """Reference situations through the tool engine: every scored reference must come back within the board's tolerance.""" - _require_engine() + snapshot = board_path(args) + _require_engine(snapshot) from policybench.ground_truth import _extract_person_value, _pe_variable_for_output - board = load_board() + board = load_board(snapshot) scored_keys = set(zip(board["scored"].scenario_id, board["scored"].variable)) values = board["reference"].set_index(["scenario_id", "variable"])["value"] rows = [] for label, reform in (("reference_engine", True), ("plain_engine", False)): - system = build_tool_system(reform=reform) + system = build_tool_system(reform=reform, snapshot=snapshot) registry = build_registry(system) for sid, scenario in board["scenarios"].items(): situation = scenario.to_pe_household() @@ -2555,6 +3784,115 @@ 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 treatment_problems( + records: Mapping[str, dict], profile: "ModelProfile" +) -> list[str]: + """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 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: + """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.""" @@ -2565,7 +3903,10 @@ def records_arm(records: Mapping[str, dict]) -> str: def command_run(args) -> None: - _require_engine() + snapshot = board_path(args) + _require_engine(snapshot) + # The notices and the tool description state the board's engine. + set_stated_engine(board_engine(snapshot)) profile = model_profile( args.model, args.arm, args.prompt, args.forbid_benchmark_search ) @@ -2579,7 +3920,7 @@ def command_run(args) -> None: from litellm import completion, responses completion_fn = responses if profile.api == "responses" else completion - board = load_board() + board = load_board(snapshot) out = Path(args.out_dir) households_dir = out / "households" households_dir.mkdir(parents=True, exist_ok=True) @@ -2610,6 +3951,13 @@ 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 = treatment_problems(previous, profile) + 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 = { @@ -2620,6 +3968,9 @@ def command_run(args) -> None: "versions": _versions(), "fix_module": str(FIX_MODULE.relative_to(ROOT)), "fix_module_sha256": _sha256(FIX_MODULE), + "board": str(snapshot), + "board_engine": f"policyengine-us {board_engine(snapshot)}", + "board_reference_sha256": board["meta"].get("reference_csv_sha256"), "model": profile.name, "model_id": profile.model_id, "arm": profile.arm, @@ -2635,7 +3986,13 @@ def command_run(args) -> None: "api": profile.api, "answer_contract": profile.contract, "tool_notice": profile.notice, - "calculate_tool": calculate_tool_schema() if profile.uses_calculator else None, + "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, "web_search_tool": web_search_tool(profile) if profile.uses_web else None, # The timeout actually sent: the board's for the model, or the web arms'. "request_timeout_seconds": build_initial_request( @@ -2662,7 +4019,16 @@ def command_run(args) -> None: if not todo: return # Arms without the calculator start no engine workers. - pool = ToolPool(size=args.workers) if profile.uses_calculator else None + pool = ( + ToolPool( + size=args.workers, + board=snapshot, + tool_version=profile.tool_version, + tool_features=profile.tool_features, + ) + if profile.uses_calculator + else None + ) def no_calculator(arguments) -> tuple[dict, float]: return {"error": "not_run", "message": "no calculator in this arm"}, 0.0 @@ -2690,23 +4056,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", []) + [ @@ -2797,16 +4149,50 @@ def arm1_predictions( return rows +def baseline_predictions( + records: Mapping[str, dict], label: str = ARM1_LABEL +) -> pd.DataFrame: + """A no-tools run's records as the board's prediction frame, with each + household's usage on its first row so per-household sums are its totals.""" + frame = arm2_predictions(records, label) + usage = { + "prompt_tokens": "prompt_tokens", + "completion_tokens": "completion_tokens", + "reasoning_tokens": "reasoning_tokens", + "total_cost_usd": "cost_usd", + "elapsed_seconds": "model_seconds", + "cached_prompt_tokens": "cached_prompt_tokens", + } + first = ~frame.duplicated("scenario_id") + for column, total in usage.items(): + frame[column] = [ + float(records[sid]["totals"].get(total) or 0.0) if is_first else 0.0 + for sid, is_first in zip(frame.scenario_id, first) + ] + frame["call_id"] = frame.scenario_id + return frame + + def command_score(args) -> None: from policybench.analysis import analyze_no_tools from policybench.spec import output_group_id - board = load_board() + snapshot = board_path(args) + board = load_board(snapshot) out = Path(args.out_dir) records = load_records(out) 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" @@ -2825,7 +4211,28 @@ def command_score(args) -> None: f"failure and score as misses until a resume reruns them: {leftover}" ) scored = board["scored"][board["scored"].scenario_id.isin(households)] - arm1 = arm1_predictions(label=arm1_label) + if is_household_set(snapshot): + # No board run exists for a fresh household set: the baseline is the + # same model's no-tools arm on it, run with the same prompt setup. + baseline_dir = getattr(args, "baseline_dir", None) + if not baseline_dir: + raise SystemExit("a household set is scored against --baseline-dir") + baseline = load_records(Path(baseline_dir)) + if records_model(baseline) != profile.name or records_arm(baseline) != "none": + raise SystemExit( + f"--baseline-dir holds {records_model(baseline)} " + f"{records_arm(baseline)} records, not {profile.name} none" + ) + if records_prompt(baseline)[0] != profile.prompt_version: + raise SystemExit("--baseline-dir was run on another prompt version") + missing = sorted(set(households) - set(baseline)) + if missing: + raise SystemExit(f"--baseline-dir lacks households {missing}") + arm1 = baseline_predictions( + {sid: baseline[sid] for sid in households}, arm1_label + ) + else: + arm1 = arm1_predictions(label=arm1_label) arm1 = arm1[arm1.scenario_id.isin(households)] arm2 = arm2_predictions({sid: records[sid] for sid in households}, arm2_label) scenarios = board["scenarios_frame"][ @@ -3328,6 +4735,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 [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"]) @@ -3934,7 +5347,9 @@ def main_call(record: dict, scenario, outputs: list[str]) -> dict | None: return main_call_and_count(record, scenario, outputs)[0] -def situation_fidelity(records, board, system, registry, defaults, out: Path) -> None: +def situation_fidelity( + records, board, system, registry, defaults, out: Path, features=frozenset() +) -> None: """How each household's main situation differs from the reference's. The main situation is ``main_call``'s: the successful call that computed @@ -3955,7 +5370,7 @@ def situation_fidelity(records, board, system, registry, defaults, out: Path) -> if call is None: rows.append({"scenario_id": sid, "differences": None}) continue - request, _ = normalize_request(call["arguments"], registry) + request, _ = normalize_request(call["arguments"], registry, features=features) if request is None: rows.append({"scenario_id": sid, "differences": None}) continue @@ -4003,9 +5418,12 @@ def situation_fidelity(records, board, system, registry, defaults, out: Path) -> def command_diagnose(args) -> None: from policybench.spec import metric_type_for_output - board = load_board() + snapshot = board_path(args) + board = load_board(snapshot) out = Path(args.out_dir) records = load_records(out) + # A run's calls are read as its own tool version read them. + features = arm_features(records_arm(records)) or frozenset() rows_path = out / "scored_rows.csv" if not rows_path.exists(): raise SystemExit("run score first") @@ -4014,14 +5432,14 @@ def command_diagnose(args) -> None: registry = None defaults = {} if not args.no_engine_defaults: - _require_engine() - system = build_tool_system(reform=True) + _require_engine(snapshot) + system = build_tool_system(reform=True, snapshot=snapshot) registry = build_registry(system) defaults = { name: system.variables[name].default_value for name in system.variables } if registry is not None: - situation_fidelity(records, board, system, registry, defaults, out) + situation_fidelity(records, board, system, registry, defaults, out, features) out_rows = [] if misses.empty: print("no arm-2 misses") @@ -4039,7 +5457,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) + request, _ = normalize_request(arguments, registry, features=features) if request is None: return None return output_value(system, request["situation"], scenario, output) @@ -4063,7 +5481,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) + request, _ = normalize_request(arguments, registry, features=features) if request is not None: reference_situation = scenario.to_pe_household() diff = situation_diff( @@ -4146,6 +5564,12 @@ def main(argv: list[str] | None = None) -> None: for name in ("oracle", "run", "score", "diagnose"): p = sub.add_parser(name) p.add_argument("--out-dir", required=True) + p.add_argument( + "--board", + default=str(SNAPSHOT), + help="the frozen board (default) or a household set from " + "scripts/household_set.py", + ) if name == "run": p.add_argument( "--model", @@ -4154,9 +5578,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", @@ -4186,6 +5611,11 @@ def main(argv: list[str] | None = None) -> None: action="store_true", help="score only households whose conversation completed", ) + p.add_argument( + "--baseline-dir", + help="on a household set, the same model's no-tools run: the " + "baseline each arm is compared with", + ) if name == "diagnose": p.add_argument("--no-engine-defaults", action="store_true") args = parser.parse_args(argv) diff --git a/tests/conftest.py b/tests/conftest.py index 0802969c..27797a13 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -57,3 +57,58 @@ def low_income_scenario(): def sample_scenarios(simple_single_scenario, family_scenario, low_income_scenario): """A small set of test scenarios.""" return [simple_single_scenario, family_scenario, low_income_scenario] + + +# A three-household set built from the frozen board's first households, with +# sampling and the dataset download faked (tests/test_household_set.py, +# tests/test_batch_judge.py). Its references must be the frozen ones on the +# frozen engine, policyengine-us 2.41.0; on any other engine it is skipped. +@pytest.fixture(scope="session") +def household_board(tmp_path_factory): + """A three-household set built from the frozen board's first households.""" + import importlib.metadata + import sys + from pathlib import Path + + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) + import household_set as hs + + if importlib.metadata.version("policyengine-us") != "2.41.0": + pytest.skip("the build test reproduces the 2.41.0 references") + import policybench.policyengine_runtime as runtime + import policybench.scenarios as scenarios_module + + root = tmp_path_factory.mktemp("set") + dataset = root / "populace_us_2024.h5" + dataset.write_bytes(b"not a real dataset") + digest = hs.sha256_file(dataset) + sampled = hs.load_scenarios(hs.FROZEN_BOARD)[:3] + patch = pytest.MonkeyPatch() + patch.setattr( + scenarios_module, + "generate_scenarios", + lambda n, seed, country, **kwargs: sampled[:n], + ) + patch.setattr(runtime, "materialize_certified_us_dataset", lambda: str(dataset)) + patch.setattr( + runtime, + "policyengine_bundles_for_countries", + lambda countries: { + "us": { + "bundle_id": "us-test", + "certified_data_artifact_sha256": digest, + "default_dataset_uri": "hf://policyengine/populace-us/x.h5@test", + # The bundle certifies its data for the installed engine. + "bundled_model_version": "2.41.0", + "model_version": "2.41.0", + } + }, + ) + out = root / "board" + try: + hs.main( + ["build", "--n", "3", "--seed", "7", "--batch-size", "2", "--out", str(out)] + ) + finally: + patch.undo() + return out, sampled diff --git a/tests/test_batch_judge.py b/tests/test_batch_judge.py new file mode 100644 index 00000000..f9e4f7d2 --- /dev/null +++ b/tests/test_batch_judge.py @@ -0,0 +1,284 @@ +"""The batch judge's payload and derivations, and the adversary hooks it uses. + +The payload tests run on the three-household set from tests/conftest.py, so +they need the frozen engine; the adversary hook tests need no engine. +""" + +from __future__ import annotations + +import hashlib +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) + +import batch_judge as bj # noqa: E402 +import tool_pilot_argon as pilot # noqa: E402 + +from policybench import reference_adversary as ra # noqa: E402 +from policybench.consensus import ConsensusParams, consensus_flags # noqa: E402 + +THREE = ConsensusParams(min_models=2, top_k=3, min_top=2, zero_cluster_min_models=2) + + +def _case(**overrides): + case = dict( + case_id="us__s__snap", + country="us", + scenario_id="s", + variable="snap", + definition="annual SNAP benefit", + household_prompt="A household.", + reference_value=100.0, + consensus=( + { + "answer": 0.0, + "n_models": 2, + "n_top": 2, + "top_models": ["a", "b"], + "members": [ + {"model": "a", "prediction": 0.0, "explanation": "x"}, + {"model": "b", "prediction": 0.0, "explanation": "y"}, + ], + }, + ), + models_answered=3, + models_exact=1, + derivation="snap = 100", + ) + case.update(overrides) + return ra.AdversaryCase(**case) + + +def test_the_default_condition_leaves_both_headers_unchanged(): + case = _case() + assert ra._STAGE1_HEADER in ra.render_stage1_prompt(case) + stage2 = ra.render_stage2_prompt(case, {"answer": 1}, "0" * 64) + assert stage2.startswith(ra._STAGE2_HEADER) + assert "answer_condition" not in case.to_manifest_row() + + +def test_a_web_condition_restates_both_headers_and_survives_the_manifest(): + case = _case(answer_condition="answering with web search") + stage1 = ra.render_stage1_prompt(case) + stage2 = ra.render_stage2_prompt(case, {"answer": 1}, "0" * 64) + for text in (stage1, stage2): + assert "AI models answering with web search, agreed" in text + assert ra.DEFAULT_ANSWER_CONDITION not in text + row = case.to_manifest_row() + assert row["answer_condition"] == "answering with web search" + again = ra.AdversaryCase.from_manifest_row(row, case.derivation) + assert again == case + default = _case() + assert ( + ra.AdversaryCase.from_manifest_row( + default.to_manifest_row(), default.derivation + ) + == default + ) + + +@pytest.mark.parametrize("condition", ["with web search", 3, ""]) +def test_a_malformed_condition_is_refused(condition): + with pytest.raises(ValueError, match="answerCondition"): + ra.answer_condition_of({"answerCondition": condition}) + + +def _payload(predictions, reference=100.0, definitions=None, condition=None): + payload = { + "country": "us", + "modelStats": [{"model": m} for m in ("a", "b", "c")], + "scenarios": {"s": {"prompt": {"tool": "A household."}, "state": "CA"}}, + "scenarioPredictions": { + "s": { + "snap": { + model: { + "prediction": value, + "explanation": f"{model} says {value}", + "groundTruth": reference, + "scored": True, + } + for model, value in predictions.items() + } + } + }, + } + if definitions is not None: + payload["outputDefinitions"] = definitions + if condition is not None: + payload["answerCondition"] = condition + return payload + + +def test_two_of_three_agreeing_off_the_reference_flag_a_cell_and_one_does_not(): + flagged = consensus_flags(_payload({"a": 0.0, "b": 0.0, "c": 100.0}), THREE) + assert [(f["scenario_id"], f["variable"]) for f in flagged] == [("s", "snap")] + assert flagged[0]["clusters"][0]["models"] == ["a", "b"] + alone = consensus_flags(_payload({"a": 0.0, "b": 50.0, "c": 100.0}), THREE) + assert alone == [] + exact = consensus_flags(_payload({"a": 100.4, "b": 99.6, "c": 100.0}), THREE) + assert exact == [] + + +def test_cases_take_the_payload_definition_and_condition(): + payload = _payload( + {"a": 0.0, "b": 0.0, "c": 100.0}, + definitions={"snap": "v2 SNAP definition"}, + condition="answering with web search", + ) + flags = consensus_flags(payload, THREE) + (case,) = ra.build_adversary_cases(payload, flags) + assert case.definition == "v2 SNAP definition" + assert case.answer_condition == "answering with web search" + payload.pop("outputDefinitions") + payload.pop("answerCondition") + (case,) = ra.build_adversary_cases(payload, flags) + assert case.answer_condition == ra.DEFAULT_ANSWER_CONDITION + assert case.definition != "v2 SNAP definition" + + +# --- On a household set ---------------------------------------------------- + + +def _write_arm(runs, board, model, arm, stage, wrong=None): + """Records for one model and arm whose sent prompt is the v2-pilot one.""" + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + folder = runs / model / bj.arm_folder(arm) / stage / "households" + folder.mkdir(parents=True) + for sid in sorted(board["scenarios"]): + scenario, variables = board["scenarios"][sid], board["variables"][sid] + prompt = bj.household_prompt(scenario, variables, "v2-pilot", model) + predictions = {v: float(reference[(sid, v)]) for v in variables} + if wrong and sid == wrong[0]: + predictions[wrong[1]] = wrong[2] + record = { + "scenario_id": sid, + "model": model, + "arm": arm, + "prompt_version": "v2-pilot", + "benchmark_search_rule": arm == "web", + "variables": variables, + "predictions": predictions, + "explanations": {v: f"{model} on {v}" for v in variables}, + "status": "complete", + "sent_board_prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(), + } + (folder / f"{sid}.json").write_text(json.dumps(record)) + + +@pytest.mark.slow +def test_a_batch_payload_flags_what_two_models_get_wrong_and_judges_it_blind( + household_board, tmp_path +): + out, _ = household_board + board = pilot.load_board(out) + variable = next( + v + for v in board["variables"]["mc_000"] + if v == "federal_income_tax_before_refundable_credits" + ) + runs = tmp_path / "runs" + for model in bj.MODELS: + wrong = ("mc_000", variable, 123456.0) if model != "gemini-4-argon" else None + _write_arm(runs, board, model, "web", "b0", wrong=wrong) + payload = bj.build_payload(out, runs, "b0", "web", batch=0) + assert sorted(payload["scenarios"]) == ["mc_000", "mc_001"] + assert payload["answerCondition"] == "answering with web search" + assert [m["model"] for m in payload["modelStats"]] == [ + f"{m}+web" for m in bj.MODELS + ] + medicaid = payload["outputDefinitions"].get("head_medicaid_eligible") + if medicaid is not None: + assert "PolicyEngine" not in medicaid + flags = consensus_flags(payload, THREE) + assert [(f["scenario_id"], f["variable"]) for f in flags] == [("mc_000", variable)] + (case,) = ra.build_adversary_cases(payload, flags) + assert case.household_prompt == payload["scenarios"]["mc_000"]["prompt"]["tool"] + assert "with web search" in ra.render_stage1_prompt(case) + + flags_path = tmp_path / "flags.json" + flags_path.write_text(json.dumps({"flags": flags})) + payload_path = tmp_path / "payload.json" + payload_path.write_text(json.dumps(payload)) + bj.main( + [ + "derivations", + "--board", + str(out), + "--payload", + str(payload_path), + "--flags", + str(flags_path), + "--out-dir", + str(tmp_path / "annotations"), + ] + ) + derivations = ra.load_derivations(tmp_path / "annotations") + text = derivations[("mc_000", variable)] + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + assert text.startswith( + f"Output {variable} = {float(reference[('mc_000', variable)]):g}" + ) + assert "income_tax_before_refundable_credits" in text + + +@pytest.mark.slow +def test_a_payload_refuses_an_unfinished_batch_a_calculator_arm_and_another_prompt( + household_board, tmp_path +): + out, _ = household_board + board = pilot.load_board(out) + runs = tmp_path / "runs" + with pytest.raises(SystemExit, match="only"): + bj.build_payload(out, runs, "b0", "pe", batch=0) + for model in bj.MODELS[:2]: + _write_arm(runs, board, model, "none", "b0") + with pytest.raises(SystemExit, match="has not finished"): + bj.build_payload(out, runs, "b0", "none", batch=0) + _write_arm(runs, board, bj.MODELS[2], "none", "b0") + record_path = runs / bj.PROMPT_MODEL / "none" / "b0" / "households" / "mc_000.json" + record = json.loads(record_path.read_text()) + record["sent_board_prompt_sha256"] = "0" * 64 + record_path.write_text(json.dumps(record)) + with pytest.raises(SystemExit, match="rebuilt prompt"): + bj.build_payload(out, runs, "b0", "none", batch=0) + with pytest.raises(SystemExit, match="batches 0 to 1"): + bj.build_payload(out, runs, "b0", "none", batch=2) + with pytest.raises(SystemExit, match="not both"): + bj.build_payload(out, runs, "b0", "none", batch=0, scenarios=["mc_000"]) + + +CHEAP = ("claude-haiku-5.5", "gpt-6-luna", "gemini-3.1-flash-lite-preview") + + +def test_the_prompt_model_is_the_first_on_the_tool_contract(): + assert bj.prompt_model_for(bj.MODELS) == "gpt-6-sol" + assert bj.prompt_model_for(CHEAP) == "claude-haiku-5.5" + assert bj.prompt_model_for(("gemini-3.1-flash-lite-preview", "gpt-6-luna")) == ( + "gpt-6-luna" + ) + + +@pytest.mark.slow +def test_a_payload_judges_any_model_set(household_board, tmp_path): + out, _ = household_board + board = pilot.load_board(out) + runs = tmp_path / "runs" + variable = "federal_income_tax_before_refundable_credits" + for model in CHEAP: + wrong = ("mc_000", variable, 999.0) if model != "gpt-6-luna" else None + _write_arm(runs, board, model, "none", "b0", wrong=wrong) + payload = bj.build_payload(out, runs, "b0", "none", batch=0, models=CHEAP) + assert payload["models"] == list(CHEAP) + assert [m["model"] for m in payload["modelStats"]] == [f"{m}+none" for m in CHEAP] + flags = consensus_flags(payload, THREE) + assert [(f["scenario_id"], f["variable"]) for f in flags] == [("mc_000", variable)] + assert flags[0]["clusters"][0]["models"] == [ + "claude-haiku-5.5+none", + "gemini-3.1-flash-lite-preview+none", + ] + with pytest.raises(SystemExit, match="at least two models"): + bj.build_payload(out, runs, "b0", "none", batch=0, models=CHEAP[:1]) diff --git a/tests/test_household_set.py b/tests/test_household_set.py new file mode 100644 index 00000000..99d42f9b --- /dev/null +++ b/tests/test_household_set.py @@ -0,0 +1,335 @@ +"""The household-set builder and the tool pilot's --board path. + +The build test runs the real engine on three of the frozen board's own +households, with sampling and the dataset download faked, so the references it +writes must equal the frozen ones on the frozen engine (policyengine-us +2.41.0). It is marked slow. +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import sys +from pathlib import Path + +import pandas as pd +import pytest +from hypothesis import given +from hypothesis import strategies as st + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) + +import household_set as hs # noqa: E402 +import tool_pilot_argon as pilot # noqa: E402 + +ON_FROZEN_ENGINE = importlib.metadata.version("policyengine-us") == "2.41.0" +SOL = "gpt-6-sol" + + +def test_rename_gives_fresh_ids_in_draw_order(): + scenarios = hs.load_scenarios(hs.FROZEN_BOARD)[:3] + renamed = hs.rename_scenarios(list(reversed(scenarios)), "mc") + assert [s.id for s in renamed] == ["mc_000", "mc_001", "mc_002"] + assert renamed[0].state == scenarios[2].state + assert renamed[0].metadata == scenarios[2].metadata + assert [s.id for s in scenarios] == ["scenario_000", "scenario_001", "scenario_002"] + + +@given(st.lists(st.text(min_size=1), unique=True, max_size=60), st.integers(1, 15)) +def test_batches_partition_the_draw_in_order(ids, size): + batches = hs.split_batches(ids, size) + assert [sid for batch in batches for sid in batch] == ids + assert all(len(batch) == size for batch in batches[:-1]) + assert all(1 <= len(batch) <= size for batch in batches) + + +def test_batches_refuse_a_size_below_one(): + with pytest.raises(ValueError): + hs.split_batches(["a"], 0) + + +def test_differences_report_missing_and_moved_outputs(): + expected = pd.DataFrame( + { + "scenario_id": ["a", "a", "b"], + "variable": ["x", "y", "x"], + "value": [1.0, 2.0, 3.0], + } + ) + computed = pd.DataFrame( + {"scenario_id": ["a", "a"], "variable": ["x", "y"], "value": [1.0, 2.5]} + ) + diff = hs.differences(expected, computed, tolerance=1e-6) + assert sorted(zip(diff.scenario_id, diff.variable)) == [("a", "y"), ("b", "x")] + assert hs.differences(expected, computed, tolerance=1.0).scenario_id.tolist() == [ + "b" + ] + + +def test_fix_module_pins_are_the_frozen_revision(): + pins = hs.fix_module_pins() + assert pins == [ + {"module": p["module"], "path": p["path"], "sha256": p["sha256"]} + for p in hs.frozen_fix_pins() + ] + + +def test_the_frozen_board_reads_as_before(): + assert pilot.board_engine() == pilot.REFERENCE_ENGINE + assert not pilot.is_household_set(pilot.SNAPSHOT) + assert pilot.board_fix_module_sha256() == pilot.FIX_MODULE_SHA256 + assert pilot.board_path(argparse.Namespace()) == pilot.SNAPSHOT + + +@pytest.mark.slow +def test_build_writes_the_frozen_references_for_frozen_households(household_board): + out, sampled = household_board + frozen = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") + built = pd.read_csv(out / "reference_outputs.csv") + names = {f"mc_{i:03d}": s.id for i, s in enumerate(sampled)} + built["frozen_id"] = built.scenario_id.map(names) + merged = built.merge( + frozen, + left_on=["frozen_id", "variable"], + right_on=["scenario_id", "variable"], + suffixes=("", "_frozen"), + ) + assert ( + len(merged) + == len(built) + == len(frozen[frozen.scenario_id.isin(names.values())]) + ) + from policybench.reference_exclusions import load_reference_exclusions + + excluded = { + (e["scenario_id"], e["variable"]) + for e in load_reference_exclusions(hs.FROZEN_BOARD) + } + scored = merged[ + [(s, v) not in excluded for s, v in zip(merged.frozen_id, merged.variable)] + ] + assert (scored.value - scored.value_frozen).abs().max() < 1e-6 + + +@pytest.mark.slow +def test_build_writes_a_board_the_pilot_reads(household_board): + out, _ = household_board + meta = json.loads((out / "reference_outputs.csv.meta.json").read_text()) + assert meta["board_kind"] == "household_set" + assert meta["seed"] == 7 and meta["num_scenarios"] == 3 + assert meta["reference_csv_sha256"] == hs.sha256_file(out / "reference_outputs.csv") + assert json.loads((out / "batches.json").read_text())["batches"] == [ + ["mc_000", "mc_001"], + ["mc_002"], + ] + board = pilot.load_board(out) + assert sorted(board["scenarios"]) == ["mc_000", "mc_001", "mc_002"] + assert board["excluded"].empty + assert pilot.is_household_set(out) + assert pilot.board_engine(out) == "2.41.0" + pilot.verify_fix_pins(snapshot=out) + pilot._require_engine(out) + + +@pytest.mark.slow +def test_build_refuses_a_written_directory_and_a_dataset_off_the_bundle( + household_board, tmp_path, monkeypatch +): + out, _ = household_board + with pytest.raises(SystemExit, match="not empty"): + hs.main(["build", "--n", "3", "--seed", "7", "--out", str(out)]) + import policybench.policyengine_runtime as runtime + + dataset = tmp_path / "x.h5" + dataset.write_bytes(b"other bytes") + monkeypatch.setattr( + runtime, "materialize_certified_us_dataset", lambda: str(dataset) + ) + monkeypatch.setattr( + runtime, + "policyengine_bundles_for_countries", + lambda countries: {"us": {"certified_data_artifact_sha256": "0" * 64}}, + ) + with pytest.raises(SystemExit, match="not the bundle's"): + hs.main(["build", "--n", "1", "--seed", "1", "--out", str(tmp_path / "new")]) + + +def _record(sid, variables, predictions, model, arm, prompt="v2-pilot"): + record = { + "scenario_id": sid, + "model": model, + "arm": arm, + "prompt_version": prompt, + "benchmark_search_rule": arm in ("web", "web+pe"), + "variables": variables, + "predictions": predictions, + "explanations": {v: "x" for v in variables}, + "status": "complete", + "turns": [{"total_cost_usd": 0.02, "calls": []}], + "cap_hit": False, + "wall_seconds": 1.0, + } + record["totals"] = pilot.household_totals(record) + return record + + +def _write_run(directory, board, model, arm, miss=None, prompt="v2-pilot"): + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + (directory / "households").mkdir(parents=True) + for sid in sorted(board["scenarios"]): + variables = board["variables"][sid] + predictions = {v: float(reference[(sid, v)]) for v in variables} + if miss is not None and sid == miss[0]: + predictions[miss[1]] = float(reference[(sid, miss[1])]) + 100 + record = _record(sid, variables, predictions, model, arm, prompt) + (directory / "households" / f"{sid}.json").write_text(json.dumps(record)) + + +@pytest.mark.slow +def test_a_household_set_scores_each_arm_against_its_no_tools_run( + household_board, tmp_path +): + out, _ = household_board + board = pilot.load_board(out) + amount = next( + v + for v in board["variables"]["mc_000"] + if not v.endswith("_eligible") and v in set(board["scored"].variable) + ) + none_dir, web_dir = tmp_path / "none", tmp_path / "web" + _write_run(none_dir, board, SOL, "none", miss=("mc_000", amount)) + _write_run(web_dir, board, SOL, "web") + pilot.command_score( + argparse.Namespace( + out_dir=str(web_dir), + complete_only=False, + board=str(out), + baseline_dir=str(none_dir), + ) + ) + score = json.loads((web_dir / "score.json").read_text()) + assert score["unweighted_hits"]["arm2"] == score["scored_outputs"] + assert score["unweighted_hits"]["arm1"] == score["scored_outputs"] - 1 + usage = pd.read_csv(web_dir / "usage_by_household.csv") + assert usage.arm1_cost_usd.round(6).tolist() == [0.02, 0.02, 0.02] + + with pytest.raises(SystemExit, match="--baseline-dir"): + pilot.command_score( + argparse.Namespace( + out_dir=str(web_dir), complete_only=False, board=str(out) + ) + ) + with pytest.raises(SystemExit, match="not gpt-6-sol none"): + pilot.command_score( + argparse.Namespace( + out_dir=str(web_dir), + complete_only=False, + board=str(out), + baseline_dir=str(web_dir), + ) + ) + v1_dir = tmp_path / "none-v1" + _write_run(v1_dir, board, SOL, "none", prompt="v1") + with pytest.raises(SystemExit, match="another prompt version"): + pilot.command_score( + argparse.Namespace( + out_dir=str(web_dir), + complete_only=False, + board=str(out), + baseline_dir=str(v1_dir), + ) + ) + + +def _scenario(sid, filing_status, tax_unit): + from policybench.scenarios import Person, Scenario + + return Scenario( + id=sid, + state="CA", + filing_status=filing_status, + adults=[Person(name="head", age=40, employment_income=0.0)], + metadata={"household_id": 1, "tax_unit_id": tax_unit}, + ) + + +def test_stated_filing_status_is_read_from_the_tax_unit_table(tmp_path): + import h5py + import numpy as np + + path = tmp_path / "set.h5" + rows = np.array( + [(0, 11, b"JOINT"), (1, 12, b"SEPARATE"), (2, 13, b"SURVIVING_SPOUSE")], + dtype=[ + ("index", " AdversaryCase: print("9.9.9 (Claude Code fake)") sys.exit(0) if args[:1] == ["auth"]: + if os.environ.get("ANTHROPIC_API_KEY"): + print(json.dumps({ + "loggedIn": True, "authMethod": "api_key", "apiProvider": "firstParty", + "apiKeySource": "ANTHROPIC_API_KEY", + "configDirectory": os.environ.get("CLAUDE_CONFIG_DIR"), + })) + sys.exit(0) print(json.dumps( {"loggedIn": True, "authMethod": "oauth_token", "apiProvider": "firstParty"} )) @@ -538,6 +545,74 @@ def test_claude_runner_refuses_without_a_lane_login(tmp_path: Path): assert _calls(bin_dir) == [] +def _api_key_env(tmp_path: Path, env: dict, key: str = "sk-ant-judge") -> dict: + """The runner's environment for an API key judge, with a fake secret store.""" + store = tmp_path / "bin" / "secret" + store.write_text( + f"#!{sys.executable}\n" + "import sys\n" + f"print({key!r}) if sys.argv[1:3] == ['get', 'JUDGE_KEY'] else sys.exit(1)\n" + ) + store.chmod(0o755) + env = {**env, "JUDGE_API_KEY_SECRET": "JUDGE_KEY", "AUDIT_SECRET_BIN": str(store)} + for name in ("CLAUDE_CONFIG_DIR", "AUDIT_ACCOUNT", "CLAUDE_CODE_OAUTH_TOKEN"): + env.pop(name, None) + return env + + +def test_claude_runner_judges_on_a_declared_api_key_in_its_own_config(tmp_path: Path): + adversary, bin_dir, env = _setup(tmp_path, "claude") + env = _api_key_env(tmp_path, env) + env["TMPDIR"] = str(tmp_path / "tmp") + (tmp_path / "tmp").mkdir() + result = _run("claude", adversary, env, tmp_path) + assert result.returncode == 0, result.stderr + result.stdout + calls = _calls(bin_dir) + assert calls + configs = {call["env"]["CLAUDE_CONFIG_DIR"] for call in calls} + assert len(configs) == 1 + config = configs.pop() + assert Path(config).parent == (tmp_path / "tmp").resolve() + assert not Path(config).exists() # removed when the run ended + for call in calls: + # The declared key, never the parent's ANTHROPIC_API_KEY. + assert call["env"]["ANTHROPIC_API_KEY"] == "sk-ant-judge" + assert "CLAUDE_CODE_OAUTH_TOKEN" not in call["env"] + for case in CASES: + meta = json.loads( + (adversary / "cases" / case.case_id / "verdict.meta.json").read_text() + ) + assert meta["judge_auth"] == { + "method": "api_key", + "account": None, + "org": None, + "secret": "JUDGE_KEY", + } + assert "JUDGE_KEY" in meta["judge_account_declared"] + assert "sk-ant-judge" not in json.dumps(meta) + + +def test_claude_runner_refuses_an_api_key_mixed_with_a_login(tmp_path: Path): + adversary, bin_dir, env = _setup(tmp_path, "claude") + base = _api_key_env(tmp_path, env) + mixes = { + "CLAUDE_CONFIG_DIR": (str(tmp_path / "lane-config"), "own empty config"), + "AUDIT_ACCOUNT": ("claude:lane@example.org", "unset AUDIT_ACCOUNT"), + "CLAUDE_CODE_OAUTH_TOKEN": ("token", "takes no lane token"), + "JUDGE_ALLOW_DESKTOP_LOGIN": ("1", "are exclusive"), + } + for name, (value, message) in mixes.items(): + result = _run("claude", adversary, {**base, name: value}, tmp_path) + assert result.returncode == 1 + assert message in result.stderr, (name, result.stderr) + result = _run( + "claude", adversary, {**base, "JUDGE_API_KEY_SECRET": "MISSING"}, tmp_path + ) + assert result.returncode == 1 + assert "could not read the API key MISSING" in result.stderr + assert _calls(bin_dir) == [] + + # --- Codex ------------------------------------------------------------------------- diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index a5d531b0..4e586374 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -2293,6 +2293,8 @@ def test_web_notices_name_search_and_keep_the_answer_instruction(): "pe": "gpt-6-sol+pe-tool", "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_v2.py b/tests/test_tool_v2.py new file mode 100644 index 00000000..68248090 --- /dev/null +++ b/tests/test_tool_v2.py @@ -0,0 +1,326 @@ +"""The calculator's tool version 2 (the pe2 arm). + +Version 2 reads the names the prompt uses (labels written with spaces, and the +benchmark's output keys), redirects a contribution set as the realized amount +to its ``_desired`` input, flags an input that replaces a formula, places an +input set on another entity where exactly one unit holds it, and explains a +variable on request. Version 1 must behave exactly as before. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +from hypothesis import given +from hypothesis import strategies as st + +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"), + "rent": V("person", "float", "rent", "year"), + "traditional_401k_contributions": V( + "person", "float", "401(k)", "year", computed=True + ), + "traditional_401k_contributions_desired": V( + "person", "float", "401(k) asked", "year" + ), + "is_medicaid_eligible": V("person", "bool", "Medicaid", "year", computed=True), + "snap": V("spm_unit", "float", "SNAP", "month", computed=True), + "free_school_meals": V("spm_unit", "float", "free meals", "year", computed=True), + "qualified_passenger_vehicle_loan_interest": V( + "household", "float", "car loan interest", "year" + ), + "state_code": V("household", "Enum", "state", "year", ("CA", "TX")), + "filing_status": V( + "tax_unit", "Enum", "filing status", "year", ("SINGLE", "JOINT"), True + ), +} + + +def _situation(**overrides): + base = { + "people": { + "head": {"age": {"2026": 40}, "employment_income": 30000}, + "kid": {"age": {"2026": 8}}, + }, + "spm_units": {"spm": {"members": ["head", "kid"]}}, + "households": {"h": {"members": ["head", "kid"], "state_code": "CA"}}, + } + base.update(overrides) + return base + + +def _call(situation, variables=("snap",), **extra): + return {"situation": json.dumps(situation), "variables": list(variables), **extra} + + +def _v2(args): + return pilot.normalize_request(args, REGISTRY, 2) + + +def test_version_1_still_refuses_what_version_2_reads(): + situation = _situation() + situation["people"]["head"]["traditional 401k contributions desired"] = 1000 + situation["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 500 + request, error = pilot.normalize_request( + _call(situation, ("free_school_meals_eligible",), explain=["snap"]), REGISTRY + ) + assert request is None + problems = " ".join(error["problems"]) + assert "unknown arguments: explain" in problems + assert "unknown variable 'traditional 401k contributions desired'" in problems + assert "is a household variable; set it under households" in problems + assert "unknown variable 'free_school_meals_eligible'" in problems + + +def test_version_2_reads_a_label_written_with_spaces(): + situation = _situation() + situation["people"]["head"]["Traditional 401k contributions desired"] = 1000 + request, error = _v2(_call(situation)) + assert error is None + head = request["situation"]["people"]["head"] + assert head["traditional_401k_contributions_desired"] == {"2026": 1000} + assert any( + "as traditional_401k_contributions_desired" in n for n in request["notes"] + ) + + +def test_version_2_sets_the_asked_amount_for_a_computed_contribution(): + situation = _situation() + situation["people"]["head"]["traditional_401k_contributions"] = 18523 + request, error = _v2(_call(situation)) + assert error is None + head = request["situation"]["people"]["head"] + assert "traditional_401k_contributions" not in head + assert head["traditional_401k_contributions_desired"] == {"2026": 18523} + assert any("before limits" in n for n in request["notes"]) + + +def test_version_2_flags_an_input_that_replaces_a_formula(): + situation = _situation() + situation["spm_units"]["spm"]["snap"] = 100 + request, error = _v2(_call(situation)) + assert error is None + assert request["situation"]["spm_units"]["spm"]["snap"] == {"2026": 100} + assert any("snap is computed by the engine" in n for n in request["notes"]) + + +def test_version_2_moves_a_household_input_to_the_persons_household(): + situation = _situation() + situation["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 500 + request, error = _v2(_call(situation)) + assert error is None + assert request["situation"]["households"]["h"][ + "qualified_passenger_vehicle_loan_interest" + ] == {"2026": 500} + assert ( + "qualified_passenger_vehicle_loan_interest" + not in (request["situation"]["people"]["head"]) + ) + assert any( + "moved qualified_passenger_vehicle_loan_interest" in n for n in request["notes"] + ) + + +def test_version_2_makes_the_one_unit_the_engine_would_when_a_group_is_left_out(): + situation = _situation() + del situation["households"] + situation["people"]["kid"]["state_code"] = "TX" + request, error = _v2(_call(situation)) + assert error is None + households = request["situation"]["households"] + assert households == { + "all": {"members": ["head", "kid"], "state_code": {"2026": "TX"}} + } + + +def test_version_2_moves_a_person_input_off_a_one_person_unit_only(): + one = _situation( + spm_units={"a": {"members": ["head"], "rent": 9000}, "b": {"members": ["kid"]}} + ) + request, error = _v2(_call(one)) + assert error is None + assert request["situation"]["people"]["head"]["rent"] == {"2026": 9000} + both = _situation() + both["spm_units"]["spm"]["rent"] = 9000 + request, error = _v2(_call(both)) + assert request is None + assert any("set it on one person" in p for p in error["problems"]) + + +def test_version_2_refuses_two_different_values_for_one_input(): + situation = _situation() + situation["households"]["h"]["qualified_passenger_vehicle_loan_interest"] = 400 + situation["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 500 + request, error = _v2(_call(situation)) + assert request is None + assert any("set twice with different values" in p for p in error["problems"]) + + +def test_version_2_refuses_a_group_input_no_single_unit_holds(): + situation = _situation( + households={"h1": {"members": ["head"]}, "h2": {"members": ["kid"]}}, + tax_units={"t": {"members": ["head", "kid"], "state_code": "CA"}}, + ) + request, error = _v2(_call(situation)) + assert request is None + assert any("no single unit there holds" in p for p in error["problems"]) + + +def test_version_2_reads_output_keys_as_their_engine_variables(): + request, error = _v2( + _call(_situation(), ("free_school_meals_eligible", "head_medicaid_eligible")) + ) + assert error is None + assert request["variables"] == ["free_school_meals", "is_medicaid_eligible"] + notes = " ".join(request["notes"]) + assert "1 when free_school_meals is above 0" in notes + assert "is_medicaid_eligible for the person the key names (head)" in notes + + +def test_version_2_explains_at_most_three_known_variables(): + request, error = _v2(_call(_situation(), explain=["snap", "free school meals"])) + assert error is None + assert request["explain"] == ["snap", "free_school_meals"] + request, error = _v2(_call(_situation(), explain=["snap", "a", "b", "c"])) + assert request is None + assert any("explain at most 3" in p for p in error["problems"]) + request, error = _v2(_call(_situation(), explain=["no_such_thing"])) + assert request is None + assert any("explain: unknown variable" in p for p in error["problems"]) + + +@given(st.sampled_from(sorted(REGISTRY)), st.sampled_from([" ", " ", "-", " - "])) +def test_a_label_written_with_any_separator_reads_back_to_its_name(name, separator): + written = name.replace("_", separator).upper() + resolved, note = pilot.resolve_name(written, REGISTRY, output_names=False) + assert resolved == name + assert (note is None) == (written == name) + + +def test_the_stated_engine_follows_the_board_and_defaults_to_the_frozen_one(): + profile = pilot.ModelProfile("m", "gpt-x", "chat", "tool", arm="pe") + frozen_notice = profile.notice + frozen_schema = pilot.calculate_tool_schema(1) + assert "policyengine-us 2.41.0" in frozen_notice + try: + pilot.set_stated_engine("2.60.4") + assert "policyengine-us 2.60.4" in profile.notice + assert "2.41.0" not in profile.notice + assert ( + "policyengine-us 2.60.4" + in pilot.calculate_tool_schema(2)["function"]["description"] + ) + finally: + pilot.set_stated_engine(pilot.REFERENCE_ENGINE) + assert profile.notice == frozen_notice + assert pilot.calculate_tool_schema(1) == frozen_schema + + +def test_the_pe2_arm_offers_tool_version_2(): + profile = pilot.ModelProfile("m", "gpt-x", "chat", "tool", arm="pe2") + assert profile.uses_calculator and not profile.uses_web + assert profile.tool_version == 2 + assert profile.notice == pilot.TOOL_NOTICE_V2 + assert "explain" in profile.notice + assert profile.arm2_label == "m+pe-tool-v2" + json_profile = pilot.ModelProfile("m", "claude-x", "chat", "json", arm="pe2") + assert json_profile.notice.endswith("single JSON object described above.") + for arm in ("pe", "web+pe"): + assert ( + pilot.ModelProfile("m", "gpt-x", "chat", "tool", arm=arm).tool_version == 1 + ) + v2 = pilot.calculate_tool_schema(2)["function"] + assert "explain" in v2["parameters"]["properties"] + assert ( + "explain" + not in pilot.calculate_tool_schema(1)["function"]["parameters"]["properties"] + ) + + +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_tool_version(): + pool = pilot.ToolPool( + size=1, timeout=5.0, worker_target=_echo_worker, tool_version=2 + ) + try: + result, _ = pool.call({"situation": "{}"}) + # A frame no model's JSON arguments can be (tool v3 review r1). + assert result == { + "received": ( + "calculate", + ("explain", "names", "overrides", "placement"), + {"situation": "{}"}, + ) + } + finally: + pool.close() + pool = pilot.ToolPool(size=1, timeout=5.0, worker_target=_echo_worker) + try: + result, _ = pool.call({"situation": "{}"}) + assert result == {"received": {"situation": "{}"}} + finally: + pool.close() + + +@pytest.mark.slow +def test_an_explanation_traces_the_engine_and_cites_its_parameters(): + system = pilot.build_tool_system(reform=True) + registry = pilot.build_registry(system) + # Computed: the engine derives it and no stated fact sets it. + assert registry["snap"].computed and registry["is_medicaid_eligible"].computed + assert not registry["employment_income"].computed # summed, but stated + assert not registry["age"].computed + assert not registry["takes_up_snap_if_eligible"].computed + situation = { + "people": { + "head": { + "age": {"2026": 40}, + "employment_income": {"2026": 20000}, + "traditional_401k_contributions": {"2026": 1000}, + } + }, + "households": {"h": {"members": ["head"], "state_code": "TX"}}, + } + request, error = pilot.normalize_request( + { + "situation": json.dumps(situation), + "variables": ["federal_income_tax_before_refundable_credits"], + "explain": ["income_tax_before_refundable_credits"], + }, + registry, + 2, + ) + assert error is None, error + assert request["variables"] == ["income_tax_before_refundable_credits"] + result = pilot.execute_request(system, registry, request) + if "traditional_401k_contributions_desired" in registry: + assert any("before limits" in n for n in result["notes"]) + explanation = result["explanations"]["income_tax_before_refundable_credits"] + assert explanation["steps"][0].startswith( + "income_tax_before_refundable_credits (2026)" + ) + assert len(explanation["steps"]) <= pilot.EXPLAIN_MAX_LINES + assert explanation["parameters"], explanation + assert not any( + p["name"].startswith("gov.abolitions.") for p in explanation["parameters"] + ) + assert any(p["sources"] for p in explanation["parameters"]) + value = result["results"]["income_tax_before_refundable_credits"]["values"] + assert set(value) == {"0"} or len(value) == 1 diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py new file mode 100644 index 00000000..71bc73c3 --- /dev/null +++ b/tests/test_tool_v3.py @@ -0,0 +1,1015 @@ +"""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 itertools +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("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" + ), + "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 + + +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 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]" + 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 "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(): + 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 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_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() + 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} + } + 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): + 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": { + "spm_unit_energy_subsidy": 900, + "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"]["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" + ] == {"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_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_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_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": {}}, + {"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={"spm unit energy subsidy": 900}), F.FLAT, F.NAMES + ) + assert error is None + 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(): + 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_in_a_frame(): + 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": ("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 + 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 + + +@pytest.mark.slow +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.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) + shown = set() + + 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)[:30]: + couple = scenario.marital_couple() or () + head = couple[0] if couple else scenario.adults[0].name + 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 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, + "household_facts": household_facts, + } + request, error = F.normalize_request( + {"household": json.dumps(household), "variables": ["snap"]}, + registry, + features=frozenset({F.FLAT, F.LABELS}), + ) + assert error is None, (scenario.id, error) + 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") + 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 + 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 + + +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_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"]) + records = { + "mt_000": dict(stale, status="complete"), + "mt_001": dict(stale, scenario_id="mt_001", status="provider_error"), + "mt_002": F.harness_error_record("mt_002", ["snap"], pe3, RuntimeError("x")), + } + 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(): + 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