From cee4ee455926c10bad2a0937a35f120fdea72987 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 08:37:00 -0400 Subject: [PATCH 01/16] Build fresh household sets and run the tool pilot on them scripts/household_set.py samples households from the dataset the installed policyengine.py bundle certifies, renames them mc_000 onward, computes every output on the installed policyengine-us with the reference conventions, and writes a board in the frozen layout with batches. On the frozen board's own households it reproduces all 1,935 scored 2.41.0 references exactly. The tool pilot takes --board: the engine and fix pins come from the board's last engine_upgrade revision, and a household set's arms are scored against the same model's no-tools run (--baseline-dir). Co-Authored-By: Claude Opus 5.5 --- scripts/household_set.py | 421 ++++++++++++++++++++++++++++++++++++ scripts/tool_pilot_argon.py | 187 ++++++++++++---- tests/test_household_set.py | 282 ++++++++++++++++++++++++ 3 files changed, 848 insertions(+), 42 deletions(-) create mode 100644 scripts/household_set.py create mode 100644 tests/test_household_set.py diff --git a/scripts/household_set.py b/scripts/household_set.py new file mode 100644 index 00000000..5b6a01ce --- /dev/null +++ b/scripts/household_set.py @@ -0,0 +1,421 @@ +"""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 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") + engine = installed_engine() + pins = fix_module_pins(Path(args.fix_module)) + programs = frozen_programs() + bundles = policyengine_bundles_for_countries(["us"]) + dataset_path = Path(materialize_certified_us_dataset()) + dataset_sha256 = sha256_file(dataset_path) + certified = bundles["us"].get("certified_data_artifact_sha256") + if certified and certified != dataset_sha256: + raise SystemExit( + f"{dataset_path} hashes to {dataset_sha256}, not the bundle's {certified}" + ) + + sampled = generate_scenarios(n=args.n, seed=args.seed, country="us") + scenarios = rename_scenarios(sampled, args.id_prefix) + 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", + "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": { + "path_name": dataset_path.name, + "sha256": dataset_sha256, + "bundle_default_dataset_uri": bundles["us"].get("default_dataset_uri"), + }, + "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}" + ) + + +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.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/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 97d880e7..ea8f85f4 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -826,44 +826,67 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic return {"results": results} +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 +894,21 @@ 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)) -def _worker_main(conn, fix_path: str, reform: bool) -> None: +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: @@ -910,9 +936,13 @@ def __init__( reform: bool = True, worker_target: Callable | None = None, acquire_timeout: float | None = None, + board: Path | None = None, ): self._ctx = mp.get_context("spawn") 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 +962,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 @@ -2366,28 +2395,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() @@ -2565,7 +2604,8 @@ def records_arm(records: Mapping[str, dict]) -> str: def command_run(args) -> None: - _require_engine() + snapshot = board_path(args) + _require_engine(snapshot) profile = model_profile( args.model, args.arm, args.prompt, args.forbid_benchmark_search ) @@ -2579,7 +2619,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) @@ -2620,6 +2660,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, @@ -2662,7 +2705,9 @@ 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) 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 @@ -2797,11 +2842,36 @@ 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( @@ -2825,7 +2895,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"][ @@ -4003,7 +4094,8 @@ 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) rows_path = out / "scored_rows.csv" @@ -4014,8 +4106,8 @@ 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 @@ -4146,6 +4238,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", @@ -4186,6 +4284,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/test_household_set.py b/tests/test_household_set.py new file mode 100644 index 00000000..ca51ea71 --- /dev/null +++ b/tests/test_household_set.py @@ -0,0 +1,282 @@ +"""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.fixture(scope="module") +def household_board(tmp_path_factory): + """A three-household set built from the frozen board's first households.""" + if not ON_FROZEN_ENGINE: + 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: 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", + } + }, + ) + out = root / "board" + try: + hs.main( + ["build", "--n", "3", "--seed", "7", "--batch-size", "2", "--out", str(out)] + ) + finally: + patch.undo() + return out, sampled + + +@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), + ) + ) From c62b48d893bb622c1a481e56ee7edad4a95ead40 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 08:44:34 -0400 Subject: [PATCH 02/16] Judge each household-set batch with the reference adversary scripts/batch_judge.py writes one arm's answers on a batch as a dashboard payload (each model a board model), so consensus-flags and the two-stage adversary run unchanged: two of three models agreeing off the reference flag a cell. The payload carries what the models saw: the v2-pilot household prompt (checked against each record's sent prompt hash), the v2-pilot definitions, and the condition they answered under. Derivations are pruned engine traces on the household set's own reference system. The adversary takes a payload's outputDefinitions and answerCondition; the default renders both stage headers byte for byte as before. Dry run on today's frozen-board no-tools answers (v2 wording): 23 of 1,935 scored cells flagged, 12 of them among the 2026-10-05 adversary's. Co-Authored-By: Claude Opus 5.5 --- policybench/reference_adversary.py | 44 ++++- scripts/batch_judge.py | 306 +++++++++++++++++++++++++++++ tests/conftest.py | 50 +++++ tests/test_batch_judge.py | 251 +++++++++++++++++++++++ tests/test_household_set.py | 39 ---- 5 files changed, 648 insertions(+), 42 deletions(-) create mode 100644 scripts/batch_judge.py create mode 100644 tests/test_batch_judge.py diff --git a/policybench/reference_adversary.py b/policybench/reference_adversary.py index 5dfd5740..eaba7ae7 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), "", diff --git a/scripts/batch_judge.py b/scripts/batch_judge.py new file mode 100644 index 00000000..77dc6ec2 --- /dev/null +++ b/scripts/batch_judge.py @@ -0,0 +1,306 @@ +"""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 household_prompt(scenario, variables: list[str], prompt_version: str) -> str: + """The household prompt the tool-contract models saw, without the arm's + notice: the board's prompt in the run's prompt version.""" + profile = pilot.model_profile(PROMPT_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, +) -> dict: + """One arm's answers on one batch (or on named households), as a + dashboard payload.""" + 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) + 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], + "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, + ) + 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.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/tests/conftest.py b/tests/conftest.py index 0802969c..79f1792b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -57,3 +57,53 @@ 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: 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", + } + }, + ) + 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..b34ae6d0 --- /dev/null +++ b/tests/test_batch_judge.py @@ -0,0 +1,251 @@ +"""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") + 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"]) diff --git a/tests/test_household_set.py b/tests/test_household_set.py index ca51ea71..aab84412 100644 --- a/tests/test_household_set.py +++ b/tests/test_household_set.py @@ -83,45 +83,6 @@ def test_the_frozen_board_reads_as_before(): assert pilot.board_path(argparse.Namespace()) == pilot.SNAPSHOT -@pytest.fixture(scope="module") -def household_board(tmp_path_factory): - """A three-household set built from the frozen board's first households.""" - if not ON_FROZEN_ENGINE: - 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: 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", - } - }, - ) - out = root / "board" - try: - hs.main( - ["build", "--n", "3", "--seed", "7", "--batch-size", "2", "--out", str(out)] - ) - finally: - patch.undo() - return out, sampled - - @pytest.mark.slow def test_build_writes_the_frozen_references_for_frozen_households(household_board): out, sampled = household_board From 078cafc48ac7fc60ccd3d1efd12606d426a838df Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 09:16:52 -0400 Subject: [PATCH 03/16] Let the Claude adversary runner bill a declared API key Max, 2026-10-22: "in case youre running any claude api stuff please use the key on my clipboard". With JUDGE_API_KEY_SECRET naming an agent-secret entry, the judges bill that key and nothing else: it refuses a lane config, a lane token, AUDIT_ACCOUNT or the desktop opt-in alongside it; every call runs in a fresh empty config directory removed at exit; the key joins only the claude calls' allowlisted environment; and check-login requires a first-party api_key login from ANTHROPIC_API_KEY in that directory. Sidecars record the mode and the secret's name, never the key. The lane and desktop paths are unchanged. Why now: the desktop login's CLI session had expired, and the Subfleet lanes were closed, so the batch judges had no login to run on. Co-Authored-By: Claude Opus 5.5 --- policybench/reference_adversary.py | 34 ++++++++++ scripts/run_reference_adversary_claude.sh | 46 +++++++++++++- tests/test_reference_adversary.py | 51 +++++++++++++++ tests/test_reference_adversary_runner.py | 75 +++++++++++++++++++++++ 4 files changed, 203 insertions(+), 3 deletions(-) diff --git a/policybench/reference_adversary.py b/policybench/reference_adversary.py index eaba7ae7..ebe7573a 100644 --- a/policybench/reference_adversary.py +++ b/policybench/reference_adversary.py @@ -2102,6 +2102,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. @@ -2111,6 +2112,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: @@ -2130,6 +2137,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( @@ -2226,6 +2258,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": @@ -2264,6 +2297,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/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/tests/test_reference_adversary.py b/tests/test_reference_adversary.py index 45c4352e..3496e381 100644 --- a/tests/test_reference_adversary.py +++ b/tests/test_reference_adversary.py @@ -1551,6 +1551,57 @@ def test_check_login_mirrors_the_audit_runner_rule(tmp_path: Path): assert check_login(lane, desktop, token=True, declared="claude:lane@x", **common) +def test_check_login_takes_an_api_key_only_as_declared_in_the_runs_own_directory( + tmp_path: Path, +): + own = tmp_path / "own" + own.mkdir() + other = tmp_path / "other" + other.mkdir() + + def status(**fields): + return json.dumps( + { + "loggedIn": True, + "authMethod": "api_key", + "apiProvider": "firstParty", + "apiKeySource": "ANTHROPIC_API_KEY", + "configDirectory": str(own), + **fields, + } + ) + + common = {"token": False, "declared": "", "config_dir": str(own)} + assert check_login( + status(), "", desktop_opt_in=False, api_key_secret="KEY", **common + ) == {"method": "api_key", "account": None, "org": None, "secret": "KEY"} + # Without the declaration an API key is still refused. + with pytest.raises(ValueError, match="not a subscription"): + check_login(status(), "", desktop_opt_in=False, **common) + refusals = [ + (status(authMethod="claude.ai"), "logs in by 'claude.ai'"), + (status(apiKeySource="/login managed key"), "not ANTHROPIC_API_KEY"), + (status(configDirectory=str(other)), "not the run's own"), + (status(apiProvider="bedrock"), "not a first-party"), + (status(loggedIn=False), "no login"), + ] + for text, message in refusals: + with pytest.raises(ValueError, match=message): + check_login(text, "", desktop_opt_in=False, api_key_secret="KEY", **common) + with pytest.raises(ValueError, match="no lane token and no desktop opt-in"): + check_login(status(), "", desktop_opt_in=True, api_key_secret="KEY", **common) + with pytest.raises(ValueError, match="no lane token and no desktop opt-in"): + check_login( + status(), + "", + token=True, + declared="", + config_dir=str(own), + desktop_opt_in=False, + api_key_secret="KEY", + ) + + # --- Command line ---------------------------------------------------------------- diff --git a/tests/test_reference_adversary_runner.py b/tests/test_reference_adversary_runner.py index 0a1d1e93..1dadf611 100644 --- a/tests/test_reference_adversary_runner.py +++ b/tests/test_reference_adversary_runner.py @@ -109,6 +109,13 @@ def _case(scenario_id: str) -> 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 ------------------------------------------------------------------------- From 879fd2bc3f1420cc404413de4ea6b9eee669f084 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 13:43:26 -0400 Subject: [PATCH 04/16] Let a judge exclude the blocked domains by operator The adversary prompt tells a judge to exclude policybench.org, policyengine.org and GitHub on every search. Codex's search tool takes no blocked_domains, so GPT-6.1 Sol did it with -site: operators, and the audit rejected those stage-1 outputs as searches naming PolicyEngine (rehearsal, 2026-10-22: 2 of 3). searched_for() drops a -site: exclusion of a blocked domain before the contamination check, in both runners' audits. A query that still names the engine or the benchmark (beside exclusions, under site:, -inurl:, or a look-alike domain) is still rejected. Co-Authored-By: Claude Opus 5.5 --- policybench/reference_adversary.py | 22 ++++++++++- tests/test_reference_adversary.py | 62 ++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/policybench/reference_adversary.py b/policybench/reference_adversary.py index ebe7573a..301a2d7e 100644 --- a/policybench/reference_adversary.py +++ b/policybench/reference_adversary.py @@ -1684,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. @@ -1752,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}") @@ -1857,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}") diff --git a/tests/test_reference_adversary.py b/tests/test_reference_adversary.py index 3496e381..89798920 100644 --- a/tests/test_reference_adversary.py +++ b/tests/test_reference_adversary.py @@ -47,6 +47,7 @@ prepare_adversary, render_stage1_prompt, render_stage2_prompt, + searched_for, sha256_text, write_stage2_prompt, ) @@ -1339,6 +1340,57 @@ def test_codex_audit_catches_contamination(item, problem): assert problems and any(problem in p for p in problems), problems +# The query GPT-6.1 Sol sent on 2026-10-22 (rehearsal case us__rh_000__snap): +# the prompt asks a judge to exclude the blocked domains on every search, and +# Codex's search tool takes no blocked_domains, so it excludes them by operator. +_EXCLUDING = ( + "site:fns.usda.gov SNAP fiscal year 2026 cost living adjustments " + + " ".join(f"-site:{domain}" for domain in BLOCKED_DOMAINS) +) + + +def test_codex_audit_accepts_a_search_that_excludes_the_blocked_domains(): + item = { + "id": "1", + "type": "web_search", + "query": _EXCLUDING, + "action": {"type": "search", "queries": [_EXCLUDING, '"7 CFR 273.10" minimum']}, + } + problems, activity = codex_events_audit(_events(item)) + assert problems == [] + assert _EXCLUDING in activity["searches"] + + +@pytest.mark.parametrize( + "query", + [ + "policyengine snap 2026 -site:policyengine.org", + "site:policyengine.org snap", + "snap -site:policyengine.org.example.com", + "snap -site:notpolicyengine.org", + "snap -inurl:policybench", + "snap -site:policybench.org policybench answers", + ], +) +def test_codex_audit_still_catches_a_search_for_the_engine_beside_exclusions(query): + problems, _ = codex_events_audit( + _events({"id": "1", "type": "web_search", "query": query}) + ) + assert any("a web search names" in p for p in problems), problems + + +@given( + st.lists(st.sampled_from(BLOCKED_DOMAINS), max_size=6), + st.text(alphabet="abcdefghijklmnopqrstuvwxyz0123456789 .:", max_size=40), +) +def test_searched_for_drops_only_blocked_domain_exclusions(domains, text): + query = " ".join([text, *(f"-site:{d}" for d in domains)]) + remaining = searched_for(query).lower() + assert all(f"-site:{d}" not in remaining.split() for d in domains) + # Everything else the judge typed survives. + assert all(word in remaining.split() for word in text.lower().split()) + + def test_codex_audit_greps_lines_that_are_not_events(): problems, _ = codex_events_audit("warning: read /x/data.json\n") assert problems == ["line 1: non-event output names 'data.json'"] @@ -1398,6 +1450,16 @@ def test_claude_audit_passes_web_research_and_returns_the_accepted_answer(): assert accepted == [_stage1()] +def test_claude_audit_takes_an_excluding_search_as_law_research(): + excluding = _use("1", "WebSearch", query=_EXCLUDING) + problems, activity, _ = claude_transcript_audit(_transcript("p", excluding), "p") + assert problems == [] + assert activity["searches"] == [_EXCLUDING] + seeking = _use("1", "WebSearch", query="policyengine snap -site:policyengine.org") + problems, _, _ = claude_transcript_audit(_transcript("p", seeking), "p") + assert any("searched for the engine" in p for p in problems), problems + + def test_claude_audit_records_a_fetch_the_deny_rule_refused(): """A blocked fetch whose result is an error returned nothing: recorded, not contamination. One that returned content is contamination.""" From b6d835a8a7324e6ac34ee2195c5d02a6fb67df48 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 07:13:06 -0400 Subject: [PATCH 05/16] Build test household sets on a named dataset; v2 needs a certified pair Max, 2026-10-23: "we can still test before that but i'd like v2 to be clean". household_set.py build takes --purpose v2 (the default) or test. A v2 set needs the installed policyengine.py bundle to certify its own dataset for the installed engine; anything else is refused unless --purpose test. A test set may name a published dataset (--dataset-repo/--dataset-file/ --dataset-revision); it records the release manifest's build engine and compatibility claim, purpose test and certified_pair false. Engines after policyengine-us 2.2.1 ignore Populace's filing_status_input and derive filing status from the household; the build records each sampled household whose stated status differs from the engine's (its scenario and reference carry the engine's). The sampler takes the dataset as an option; the default call is unchanged. Co-Authored-By: Claude Opus 5.5 --- policybench/scenarios.py | 32 ++++++-- scripts/household_set.py | 155 ++++++++++++++++++++++++++++++++++-- tests/conftest.py | 7 +- tests/test_household_set.py | 92 +++++++++++++++++++++ 4 files changed, 272 insertions(+), 14 deletions(-) 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/household_set.py b/scripts/household_set.py index 5b6a01ce..ff385040 100644 --- a/scripts/household_set.py +++ b/scripts/household_set.py @@ -217,6 +217,95 @@ 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, @@ -227,20 +316,53 @@ def command_build(args) -> None: 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"]) - dataset_path = Path(materialize_certified_us_dataset()) - dataset_sha256 = sha256_file(dataset_path) - certified = bundles["us"].get("certified_data_artifact_sha256") - if certified and certified != dataset_sha256: + 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( - f"{dataset_path} hashes to {dataset_sha256}, not the bundle's {certified}" + "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") + 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)) ) @@ -270,6 +392,8 @@ def command_build(args) -> None: "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, @@ -279,10 +403,14 @@ def command_build(args) -> None: "programs": programs, "policyengine_bundles": bundles, "dataset": { + **dataset_record, "path_name": dataset_path.name, "sha256": dataset_sha256, - "bundle_default_dataset_uri": bundles["us"].get("default_dataset_uri"), + "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 @@ -329,7 +457,9 @@ def command_build(args) -> None: ) print( f"{len(scenarios)} households, {len(references)} outputs on {engine}; " - f"{len(batches)} batches of up to {args.batch_size} in {out}" + 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" ) @@ -402,6 +532,15 @@ def main(argv: list[str] | None = None) -> None: 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) diff --git a/tests/conftest.py b/tests/conftest.py index 79f1792b..27797a13 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -85,7 +85,9 @@ def household_board(tmp_path_factory): sampled = hs.load_scenarios(hs.FROZEN_BOARD)[:3] patch = pytest.MonkeyPatch() patch.setattr( - scenarios_module, "generate_scenarios", lambda n, seed, country: sampled[:n] + scenarios_module, + "generate_scenarios", + lambda n, seed, country, **kwargs: sampled[:n], ) patch.setattr(runtime, "materialize_certified_us_dataset", lambda: str(dataset)) patch.setattr( @@ -96,6 +98,9 @@ def household_board(tmp_path_factory): "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", } }, ) diff --git a/tests/test_household_set.py b/tests/test_household_set.py index aab84412..99d42f9b 100644 --- a/tests/test_household_set.py +++ b/tests/test_household_set.py @@ -241,3 +241,95 @@ def test_a_household_set_scores_each_arm_against_its_no_tools_run( 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", " Date: Fri, 23 Oct 2026 12:41:39 -0400 Subject: [PATCH 06/16] Add calculator tool version 2 (the pe2 arm) Max, 2026-10-23, on what test batch 0 showed about the calculator ("yes go"): - Names: a name written with spaces (as the prompt labels its facts) or a benchmark output key (free_school_meals_eligible, head_medicaid_eligible) is read as its engine variable, and the result's notes say how. Every model had a call refused for a guessed name in all 10 households. - Placement: an input set on another entity moves to the one unit that holds the people it was set for (a group left out becomes one unit of everyone, as the engine reads it); otherwise it stays an error. - Overrides: a contribution set as the realized amount is read as its _desired input (Sol and Opus put "traditional 401k contributions desired" into traditional_401k_contributions in 6 households, moving tax), and an input the engine computes and no stated fact sets is flagged as replacing its formula (Sol set rent, not pre_subsidy_rent). - Explanations: explain (up to 3 variables) returns the computation steps and the parameters read, with values and cited sources. Opus overrode a correct free-meals answer in Louisiana, which pays the reduced-price copay. Tool version 1 (pe, web+pe) is byte for byte as before. The notices and tool description now state the board's engine (a household set on 2.60.4 had been told 2.41.0); the frozen board's text is unchanged. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 562 ++++++++++++++++++++++++++++++--- tests/test_tool_pilot_argon.py | 1 + tests/test_tool_v2.py | 321 +++++++++++++++++++ 3 files changed, 835 insertions(+), 49 deletions(-) create mode 100644 tests/test_tool_v2.py diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index ea8f85f4..c8a1559b 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -152,13 +152,54 @@ # 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") ARM_SUFFIX = { "none": "+no-tools", "pe": "+pe-tool", "web": "+web", "web+pe": "+web+pe-tool", + "pe2": "+pe-tool-v2", } +# "pe2" is the calculator, tool version 2 (Max, 2026-10-23: "yes go"), alone: +# it reads the names the prompt uses, places inputs set on the wrong entity, +# flags inputs that replace a formula, and explains a value on request. +TOOL_VERSION_BY_ARM = {"pe": 1, "web+pe": 1, "pe2": 2} + +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.", +) +# 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 + + +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 +493,43 @@ 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 + + +def calculate_tool_schema(version: int = 1) -> dict: + if version == 2: + schema = calculate_tool_schema(1) + function = schema["function"] + function["description"] = engine_text(CALCULATE_DESCRIPTION_V2) + function["parameters"]["properties"]["explain"] = { + "type": "array", + "items": {"type": "string"}, + "description": ( + f"Up to {MAX_EXPLAIN} variables to explain: how the engine " + "computed them." + ), + } + return schema return { "type": "function", "function": { "name": TOOL_NAME, - "description": CALCULATE_DESCRIPTION, + "description": engine_text(CALCULATE_DESCRIPTION), "parameters": { "type": "object", "properties": { @@ -499,9 +571,53 @@ 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 -def build_registry(system) -> dict[str, VariableInfo]: +# 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 stated_inputs() -> frozenset[str]: + """Every variable a household fact the benchmark states can set.""" + from policybench.scenarios import ( + PE_INPUT_ALIASES, + PE_INPUT_RENAMES, + get_promptable_input_specs, + ) + + names = {spec.source_name for spec in get_promptable_input_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 +630,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 +678,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 +724,20 @@ 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 = resolve_name(name, registry, output_names=False) + if resolved is not None: + v2["notes"].append(f"{where}: {note}") + name = resolved if not isinstance(name, str) or name not in registry: problems.append( f"{where}: unknown variable {name!r}." @@ -575,52 +745,57 @@ def _normalize_inputs( ) continue info = registry[name] + if v2 is not None: + 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: + 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 ) -> 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] = [] + v2 = {"notes": [], "misplaced": []} if version == 2 else None def fail() -> tuple[None, dict]: return None, { @@ -643,7 +818,10 @@ 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 = {"situation", "variables", "period"} | ( + {"explain"} if version == 2 else set() + ) + unknown = sorted(set(args) - allowed) if unknown: problems.append(f"unknown arguments: {', '.join(map(str, unknown))}") @@ -702,7 +880,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 +924,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 +943,162 @@ 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: + _place_misplaced(v2, normalized, situation, registry, period, problems) requested = [] for name in variables: + if v2 is not None and isinstance(name, str) and name not in registry: + 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 v2 is not None 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 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) 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"] + return request, None + + +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 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,7 +1147,121 @@ 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: explain_variable(system, request, name, simulation_cls) + for name in request["explain"] + } + return out + + +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: @@ -918,7 +1356,10 @@ def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> return if args is None: return - request, error = normalize_request(args, registry) + version = 1 + if isinstance(args, dict) and args.get("__tool_version__") == 2: + version, args = 2, args.get("arguments") + request, error = normalize_request(args, registry, version) if error is not None: conn.send(("ok", error)) continue @@ -937,8 +1378,10 @@ def __init__( worker_target: Callable | None = None, acquire_timeout: float | None = None, board: Path | None = None, + tool_version: int = 1, ): self._ctx = mp.get_context("spawn") + self._tool_version = tool_version 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. @@ -1005,6 +1448,8 @@ def call(self, args) -> tuple[dict, float]: time.perf_counter() - started, ) process, conn = worker + if self._tool_version == 2: + args = {"__tool_version__": 2, "arguments": args} try: conn.send(args) if conn.poll(self.timeout): @@ -1160,13 +1605,21 @@ 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) 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 self.arm in TOOL_VERSION_BY_ARM + + @property + def tool_version(self) -> int: + return TOOL_VERSION_BY_ARM.get(self.arm, 1) @property def uses_web(self) -> bool: @@ -1845,7 +2298,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(profile.tool_version) offered = ([web_search_tool(profile)] if profile.uses_web else []) + ( [calculate_tool] if profile.uses_calculator else [] ) @@ -2606,6 +3059,8 @@ def records_arm(records: Mapping[str, dict]) -> str: def command_run(args) -> None: 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 ) @@ -2678,7 +3133,10 @@ 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(profile.tool_version) + 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( @@ -2706,7 +3164,9 @@ def command_run(args) -> None: return # Arms without the calculator start no engine workers. pool = ( - ToolPool(size=args.workers, board=snapshot) if profile.uses_calculator else None + ToolPool(size=args.workers, board=snapshot, tool_version=profile.tool_version) + if profile.uses_calculator + else None ) def no_calculator(arguments) -> tuple[dict, float]: @@ -4025,7 +4485,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, version: int = 1 +) -> None: """How each household's main situation differs from the reference's. The main situation is ``main_call``'s: the successful call that computed @@ -4046,7 +4508,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, version) if request is None: rows.append({"scenario_id": sid, "differences": None}) continue @@ -4098,6 +4560,8 @@ def command_diagnose(args) -> None: 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. + version = TOOL_VERSION_BY_ARM.get(records_arm(records), 1) rows_path = out / "scored_rows.csv" if not rows_path.exists(): raise SystemExit("run score first") @@ -4113,7 +4577,7 @@ def command_diagnose(args) -> None: name: system.variables[name].default_value for name in system.variables } if registry is not None: - situation_fidelity(records, board, system, registry, defaults, out) + situation_fidelity(records, board, system, registry, defaults, out, version) out_rows = [] if misses.empty: print("no arm-2 misses") @@ -4131,7 +4595,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, version) if request is None: return None return output_value(system, request["situation"], scenario, output) @@ -4155,7 +4619,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, version) if request is not None: reference_situation = scenario.to_pe_household() diff = situation_diff( diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index a5d531b0..986a89d9 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -2293,6 +2293,7 @@ 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", } 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..d6e381b5 --- /dev/null +++ b/tests/test_tool_v2.py @@ -0,0 +1,321 @@ +"""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": "{}"}) + assert result == { + "received": {"__tool_version__": 2, "arguments": {"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 From 5bbd571a51d15ec8f5deb928fb0afacf3e72b3d8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 14:18:50 -0400 Subject: [PATCH 07/16] Judge a batch from any model set Max, 2026-10-23: "thinking we might want to do all these tests with a cheap model until we're super ready for v2 with the web and tool experiments". batch_judge.py payload takes --models (the frontier three by default). The payload's household prompt is the first tool-contract model's, checked against that model's records, and the payload records which models it holds. Co-Authored-By: Claude Opus 5.5 --- scripts/batch_judge.py | 44 +++++++++++++++++++++++++++++---------- tests/test_batch_judge.py | 35 ++++++++++++++++++++++++++++++- 2 files changed, 67 insertions(+), 12 deletions(-) diff --git a/scripts/batch_judge.py b/scripts/batch_judge.py index 77dc6ec2..17e33045 100644 --- a/scripts/batch_judge.py +++ b/scripts/batch_judge.py @@ -78,10 +78,21 @@ def load_batch(board: Path, batch: int) -> list[str]: return list(batches[batch]) -def household_prompt(scenario, variables: list[str], prompt_version: str) -> str: - """The household prompt the tool-contract models saw, without the arm's - notice: the board's prompt in the run's prompt version.""" - profile = pilot.model_profile(PROMPT_MODEL, "none", prompt_version, False) +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 @@ -105,9 +116,15 @@ def build_payload( 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.""" + 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): @@ -120,7 +137,7 @@ def build_payload( scored = set(zip(board["scored"].scenario_id, board["scored"].variable)) reference = board["reference"].set_index(["scenario_id", "variable"])["value"] records = {} - for model in MODELS: + for model in models: folder = runs / model / arm_folder(arm) / stage / "households" found = {} for sid in ids: @@ -136,16 +153,16 @@ def build_payload( raise SystemExit(f"the models ran different prompt versions: {versions}") prompt_version = versions.pop() - labels = {model: f"{model}+{arm}" for model in MODELS} + 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) - sent = records[PROMPT_MODEL][sid].get("sent_board_prompt_sha256") + 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" + f"{sid}: the rebuilt prompt is not the one {prompt_model} saw" ) scenarios[sid] = {"prompt": {"tool": prompt}, "state": scenario.state} cells = {} @@ -175,8 +192,9 @@ def build_payload( "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], + "modelStats": [{"model": labels[model]} for model in models], "scenarios": scenarios, "scenarioPredictions": predictions, } @@ -230,6 +248,7 @@ def command_payload(args) -> None: 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()) @@ -291,6 +310,9 @@ def main(argv: list[str] | None = None) -> None: "--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) diff --git a/tests/test_batch_judge.py b/tests/test_batch_judge.py index b34ae6d0..f9e4f7d2 100644 --- a/tests/test_batch_judge.py +++ b/tests/test_batch_judge.py @@ -150,7 +150,7 @@ def _write_arm(runs, board, model, arm, stage, wrong=None): 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") + 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] @@ -249,3 +249,36 @@ def test_a_payload_refuses_an_unfinished_batch_a_calculator_arm_and_another_prom 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]) From e881a6a418b78d6164cf062f35332ae21b830516 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:11:22 -0400 Subject: [PATCH 08/16] Calculator features one at a time: tool v3 and ablation arms The calculator's six features (names, placement, overrides, explain, concise, flat) are each on or off, so a study can change one at a time. "pe2" keeps tool version 2 (the first four, its notice byte for byte), "pe3" offers all six, and any other set runs as "pe:" plus its features, e.g. "pe:names+flat". - concise: an explanation gives the value, its direct inputs and up to five parameters unless `detail` is "full". - flat: the household is described as a state, people with a role (head, spouse or dependent) and facts, and household facts; the tool builds the units the benchmark's own situations have and places each household fact on its entity. Tests: the v2 notice is reproduced from its features; each feature adds only its own behavior; flat builds the benchmark's units and refuses what it cannot place; and on the installed engine eight frozen-board households described flat compute every board output exactly as Scenario.to_pe_household does. Five mutants of the flat and arm logic are all caught. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 522 +++++++++++++++++++++++++++++---- tests/test_tool_pilot_argon.py | 1 + tests/test_tool_v3.py | 377 ++++++++++++++++++++++++ 3 files changed, 850 insertions(+), 50 deletions(-) create mode 100644 tests/test_tool_v3.py diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index c8a1559b..88d8c4d6 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -152,18 +152,75 @@ # with no tools), so a prompt version other than the board's has its own # no-tools baseline. Repairs, retries and truncation then follow this pilot's # loop, not the board's; they matter only when a first answer does not parse. -ARMS = ("none", "pe", "web", "web+pe", "pe2") +ARMS = ("none", "pe", "web", "web+pe", "pe2", "pe3") ARM_SUFFIX = { "none": "+no-tools", "pe": "+pe-tool", "web": "+web", "web+pe": "+web+pe-tool", "pe2": "+pe-tool-v2", + "pe3": "+pe-tool-v3", } -# "pe2" is the calculator, tool version 2 (Max, 2026-10-23: "yes go"), alone: -# it reads the names the prompt uses, places inputs set on the wrong entity, -# flags inputs that replace a formula, and explains a value on request. -TOOL_VERSION_BY_ARM = {"pe": 1, "web+pe": 1, "pe2": 2} +# The calculator's features, each on or off, so a study can change one at a +# time (Max, 2026-10-23: "measuring this as an experiment could be +# publishable"; "do all these tests with a cheap model"): +# - names: a name written with spaces, or a benchmark output key, is read as +# its engine variable; +# - placement: an input set on another entity moves to the one unit holding +# its people; +# - overrides: a realized contribution is read as its _desired input, and an +# input that replaces a formula is flagged; +# - explain: `explain` traces up to three variables; +# - concise: an explanation is brief unless `detail` is "full"; +# - flat: the household is described as people with roles and facts +# (`household`), and the tool builds its units as the benchmark does. +# "pe2" is tool version 2 (the first four); "pe3" is all six. Any other set +# runs as an arm named "pe:" plus its features joined by "+", e.g. +# "pe:names+flat". +NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( + "names", + "placement", + "overrides", + "explain", + "concise", + "flat", +) +TOOL_FEATURES = (NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT) +V2_FEATURES = frozenset({NAMES, PLACEMENT, OVERRIDES, EXPLAIN}) +FEATURES_BY_ARM = { + "pe": frozenset(), + "web+pe": frozenset(), + "pe2": V2_FEATURES, + "pe3": frozenset(TOOL_FEATURES), +} +_FEATURE_ARM = re.compile(r"^pe:([a-z]+(?:\+[a-z]+)*)$") + + +def arm_features(arm: str) -> frozenset | None: + """The calculator features an arm offers, or None for an arm without the + calculator. Refuses an unknown feature, or concise without explain.""" + if arm in FEATURES_BY_ARM: + return FEATURES_BY_ARM[arm] + match = _FEATURE_ARM.match(arm or "") + if match is None: + return None + features = frozenset(match.group(1).split("+")) + unknown = sorted(features - set(TOOL_FEATURES)) + if unknown: + raise SystemExit( + f"unknown tool features {unknown}; the features are {TOOL_FEATURES}" + ) + if CONCISE in features and EXPLAIN not in features: + raise SystemExit("the concise feature needs explain") + return features + + +def arm_suffix(arm: str) -> str: + if arm in ARM_SUFFIX: + return ARM_SUFFIX[arm] + features = arm_features(arm) + return f"+pe-tool[{'+'.join(f for f in TOOL_FEATURES if f in features)}]" + TOOL_NOTICE_V2 = ( "\n\nYou also have a `calculate_household` function. It runs PolicyEngine US " @@ -183,6 +240,76 @@ "instructed above.", "Then give your final answers as the single JSON object described above.", ) + + +def _listing(items: list[str]) -> str: + if len(items) <= 2: + return " and ".join(items) + return ", ".join(items[:-1]) + ", and " + items[-1] + + +def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: + """The calculator notice for a set of features: the opening, a sentence + per feature, then the turn limit and the answer instruction. No feature + gives TOOL_NOTICE and the version 2 features give TOOL_NOTICE_V2, byte for + byte.""" + if not features: + return TOOL_NOTICE if tool_contract else JSON_TOOL_NOTICE + what = ( + "household you describe" + if FLAT in features + else "household situation you build" + ) + text = ( + "\n\nYou also have a `calculate_household` function. It runs PolicyEngine US " + "(policyengine-us 2.41.0), an open-source rules engine for US federal and " + f"state taxes and benefits, on a {what} from the facts above." + ) + if FLAT in features: + text += ( + " Describe it as `household`: each person with their role (head, spouse " + "or dependent) and facts, and the facts about the household as a " + "whole; it builds the household's units as the facts imply." + ) + if NAMES in features: + text += ( + " The facts above are labeled with its variable names, written with " + "spaces; it accepts them as written, and it accepts the output keys you " + "are asked for. Set every listed fact." + ) + notes = [] + if NAMES in features: + notes.append("which names it read") + if PLACEMENT in features: + notes.append("which inputs it moved to the entity they belong to") + if OVERRIDES in features: + notes.append("which inputs replace a formula") + explain = "" + if EXPLAIN in features: + explain = ( + "ask it to `explain` up to three variables to see in brief how it " + "computed them, with the parameters and laws it used, or every step " + 'with `detail` set to "full"' + if CONCISE in features + else "ask it to `explain` up to three variables to see how it computed " + "them, with the parameters and laws it used" + ) + if notes and explain: + text += f" Its result says {_listing(notes)}; {explain}." + elif notes: + text += f" Its result says {_listing(notes)}." + elif explain: + text += f" You may {explain}." + text += " You may call it in up to 8 turns. " + text += ( + "Then call `submit_outputs` exactly once with your final answers, as " + "instructed above." + if tool_contract + else "Then give your final answers as the single JSON object described above." + ) + return text + + # The engine version the notices and the tool description state: the board's # (set by command_run), so a household set on a newer engine says so. The # frozen board's 2.41.0 leaves every text byte for byte as before. @@ -511,19 +638,86 @@ def refuse_policyengine(label: str, value) -> None: MAX_EXPLAIN = 3 -def calculate_tool_schema(version: int = 1) -> dict: - if version == 2: +FLAT_DESCRIPTION = ( + "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " + "the requested variables. `household` is a JSON object: `state` (the " + "two-letter state code), `people` (a list, each with `name`, `role` -- " + 'head, spouse or dependent -- and `facts`, for example {"age": 40, ' + '"employment income": 30000}), and `household_facts` (facts about the ' + "household as a whole: its tax unit, family, SPM unit or home). The head " + "and spouse are one couple and everyone else is on their own; everyone " + "shares one tax unit, family, SPM unit and household. A fact is a value " + 'for `period`, or {period: value} with a period such as "2026" or ' + '"2026-01". Facts you do not give take the engine\'s defaults. `variables` ' + 'lists the variables to compute for `period` (default "2026"). Each ' + "result gives the variable's entity, label and value for each member of " + "that entity; a monthly variable computed for a year is summed over its " + "months." +) + + +def calculate_tool_schema(version: int = 1, features: frozenset | None = None) -> dict: + if features is None: + features = V2_FEATURES if version == 2 else frozenset() + if features: schema = calculate_tool_schema(1) function = schema["function"] - function["description"] = engine_text(CALCULATE_DESCRIPTION_V2) - function["parameters"]["properties"]["explain"] = { - "type": "array", - "items": {"type": "string"}, - "description": ( - f"Up to {MAX_EXPLAIN} variables to explain: how the engine " - "computed them." - ), - } + properties = function["parameters"]["properties"] + if features == V2_FEATURES: + function["description"] = engine_text(CALCULATE_DESCRIPTION_V2) + else: + description = ( + FLAT_DESCRIPTION if FLAT in features else CALCULATE_DESCRIPTION + ) + notes = [] + if NAMES in features: + notes.append( + "A name may also be written with spaces (as the facts are " + "labeled) or be one of the output keys you are asked for." + ) + if PLACEMENT in features and FLAT not in features: + notes.append( + "An input set on the wrong entity is moved to the one unit it " + "belongs to." + ) + if OVERRIDES in features: + notes.append( + "Setting a variable the engine computes replaces its formula; a " + "contribution set as the amount contributed is read as the " + "amount asked for before limits (`..._desired`)." + ) + if notes: + notes.append("The result's `notes` say what it did.") + function["description"] = engine_text(" ".join([description, *notes])) + if FLAT in features: + del properties["situation"] + properties["household"] = { + "type": "string", + "description": ( + "The household as a JSON object (state, people with role and " + "facts, household_facts), written as a JSON string." + ), + } + function["parameters"]["required"] = ["household", "variables"] + if EXPLAIN in features: + properties["explain"] = { + "type": "array", + "items": {"type": "string"}, + "description": ( + f"Up to {MAX_EXPLAIN} variables to explain: how the engine " + "computed them." + ), + } + if CONCISE in features: + properties["detail"] = { + "type": "string", + "enum": ["brief", "full"], + "description": ( + 'How much an explanation shows: "brief" (the default: the ' + 'value, what it came from, and the main parameters) or "full" ' + "(every step)." + ), + } return schema return { "type": "function", @@ -733,7 +927,12 @@ def _normalize_inputs( for name, value in data.items(): if allow_members and name == "members": continue - if v2 is not None and isinstance(name, str) and name not in registry: + if ( + v2 is not None + and NAMES in v2["features"] + and isinstance(name, str) + and name not in registry + ): resolved, note = resolve_name(name, registry, output_names=False) if resolved is not None: v2["notes"].append(f"{where}: {note}") @@ -745,7 +944,7 @@ def _normalize_inputs( ) continue info = registry[name] - if v2 is not None: + if v2 is not None and OVERRIDES in v2["features"]: desired = f"{name}{DESIRED_SUFFIX}" if ( info.computed @@ -763,7 +962,7 @@ def _normalize_inputs( "replaces its formula" ) if info.entity != entity: - if v2 is not None: + if v2 is not None and PLACEMENT in v2["features"]: v2["misplaced"].append((where, entity, name, value)) continue problems.append( @@ -782,7 +981,10 @@ def _normalize_inputs( def normalize_request( - args, registry: Mapping[str, VariableInfo], version: int = 1 + args, + registry: Mapping[str, VariableInfo], + version: int = 1, + features: frozenset | None = None, ) -> tuple[dict | None, dict | None]: """Check a calculate_household call; return (request, None) or (None, error). @@ -795,7 +997,9 @@ def normalize_request( ``explain``; the request's ``notes`` say what it did. """ problems: list[str] = [] - v2 = {"notes": [], "misplaced": []} if version == 2 else None + if features is None: + features = V2_FEATURES if version == 2 else frozenset() + v2 = {"notes": [], "misplaced": [], "features": features} if features else None def fail() -> tuple[None, dict]: return None, { @@ -818,9 +1022,10 @@ def fail() -> tuple[None, dict]: if not isinstance(args, dict): problems.append("arguments must be an object") return fail() - allowed = {"situation", "variables", "period"} | ( - {"explain"} if version == 2 else set() - ) + allowed = {"variables", "period"} + allowed |= {"household"} if FLAT in features else {"situation"} + allowed |= {"explain"} if EXPLAIN in features else set() + allowed |= {"detail"} if CONCISE in features else set() unknown = sorted(set(args) - allowed) if unknown: problems.append(f"unknown arguments: {', '.join(map(str, unknown))}") @@ -836,6 +1041,13 @@ def fail() -> tuple[None, dict]: ) period = str(YEAR) + if FLAT in features: + args = dict(args) + args["situation"] = flat_situation( + args.get("household"), registry, problems, v2 + ) + if args["situation"] is None: + return fail() situation = args.get("situation") if isinstance(situation, str): if len(situation) > MAX_SITUATION_CHARS: @@ -943,11 +1155,16 @@ def fail() -> tuple[None, dict]: variables = [] elif len(variables) > MAX_VARIABLES: problems.append(f"at most {MAX_VARIABLES} variables per call") - if v2 is not None: + if v2 is not None and PLACEMENT in features: _place_misplaced(v2, normalized, situation, registry, period, problems) requested = [] for name in variables: - if v2 is not None and isinstance(name, str) and name not in registry: + if ( + v2 is not None + and NAMES in features + and isinstance(name, str) + and name not in registry + ): resolved, note = resolve_name(name, registry, output_names=True) if resolved is not None: v2["notes"].append(f"variables: {note}") @@ -959,7 +1176,7 @@ def fail() -> tuple[None, dict]: elif name not in requested: requested.append(name) explain = [] - if v2 is not None and args.get("explain") is not None: + if EXPLAIN in features and args.get("explain") is not None: names = args["explain"] if isinstance(names, str): names = [names] @@ -969,7 +1186,7 @@ def fail() -> tuple[None, dict]: elif len(names) > MAX_EXPLAIN: problems.append(f"explain at most {MAX_EXPLAIN} variables per call") for name in names: - if name not in registry: + if NAMES in features and name not in registry: resolved, note = resolve_name(name, registry, output_names=True) if resolved is not None: v2["notes"].append(f"explain: {note}") @@ -981,15 +1198,154 @@ def fail() -> tuple[None, dict]: elif name not in explain: explain.append(name) + detail = "full" + if CONCISE in features: + detail = args.get("detail") or "brief" + if detail not in ("brief", "full"): + problems.append('detail must be "brief" or "full"') + if problems: return fail() request = {"situation": normalized, "variables": requested, "period": period} if v2 is not None: request["explain"] = explain request["notes"] = v2["notes"] + request["detail"] = detail return request, None +FLAT_ROLES = ("head", "spouse", "dependent") + + +def flat_situation(household, registry, problems: list[str], v2: dict | None): + """The flat feature's household -- state, people with role and facts, and + household facts -- as a situation, with the units the benchmark's own + reference situations have (policybench.scenarios.Scenario.to_pe_household): + the head and spouse one marital unit and everyone else alone, and one tax + unit, SPM unit, family and household holding everyone. A household fact + goes to the unit of its variable's entity. Returns None after recording + problems.""" + if isinstance(household, str): + if len(household) > MAX_SITUATION_CHARS: + problems.append("household is too large") + return None + try: + household = json.loads(household) + except (ValueError, RecursionError) as error: + problems.append(f"household is not valid JSON: {error}"[:MAX_MESSAGE_CHARS]) + return None + if not isinstance(household, dict): + problems.append( + "household must be a JSON object with state, people and household_facts" + ) + return None + extra = sorted(set(household) - {"state", "people", "household_facts"}) + if extra: + problems.append( + "household accepts state, people and household_facts; remove " + + ", ".join(map(str, extra)) + ) + people = household.get("people") + if not isinstance(people, list) or not people: + problems.append("household.people must be a non-empty list of people") + return None + if len(people) > MAX_PEOPLE: + problems.append(f"at most {MAX_PEOPLE} people") + return None + situation: dict[str, Any] = {"people": {}} + roles: dict[str, str] = {} + for index, person in enumerate(people): + if not isinstance(person, dict): + problems.append( + f"people[{index}] must be an object with name, role and facts" + ) + continue + name, role, facts = ( + person.get("name"), + person.get("role"), + person.get("facts", {}), + ) + if not isinstance(name, str) or not name or len(name) > 64: + problems.append(f"people[{index}].name must be a short string") + continue + if name in roles: + problems.append(f"people[{index}]: {name!r} is listed twice") + continue + if role not in FLAT_ROLES: + problems.append( + f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" + ) + continue + if not isinstance(facts, dict): + problems.append(f"people[{index}].facts must be an object") + continue + roles[name] = role + situation["people"][name] = dict(facts) + heads = [n for n, r in roles.items() if r == "head"] + spouses = [n for n, r in roles.items() if r == "spouse"] + if len(heads) != 1: + problems.append("the household needs exactly one head") + if len(spouses) > 1: + problems.append("the household has at most one spouse") + if problems: + return None + for name, role in roles.items(): + facts = situation["people"][name] + facts.setdefault("is_tax_unit_head", role == "head") + facts.setdefault("is_tax_unit_spouse", role == "spouse") + everyone = list(situation["people"]) + couple = heads + spouses + groups = [couple] + [[n] for n in everyone if n not in couple] + situation["marital_units"] = { + f"marital_unit_{i}": {"members": members} + for i, members in enumerate(groups, start=1) + } + situation["tax_units"] = {"tax_unit": {"members": list(everyone)}} + situation["spm_units"] = {"spm_unit": {"members": list(everyone)}} + situation["families"] = {"family": {"members": list(everyone)}} + situation["households"] = {"household": {"members": list(everyone)}} + state = household.get("state") + if state is not None: + situation["households"]["household"]["state_code"] = state + unit_of = { + "tax_unit": situation["tax_units"]["tax_unit"], + "spm_unit": situation["spm_units"]["spm_unit"], + "family": situation["families"]["family"], + "household": situation["households"]["household"], + "marital_unit": situation["marital_units"]["marital_unit_1"], + } + facts = household.get("household_facts") or {} + if not isinstance(facts, dict): + problems.append("household_facts must be an object") + return None + for name, value in facts.items(): + resolved = name + if ( + isinstance(name, str) + and name not in registry + and v2 is not None + and NAMES in v2["features"] + ): + resolved, note = resolve_name(name, registry, output_names=False) + if resolved is not None: + v2["notes"].append(f"household_facts: {note}") + if not isinstance(resolved, str) or resolved not in registry: + problems.append( + f"household_facts: unknown variable {name!r}." + + (_suggest(name, registry) if isinstance(name, str) else "") + ) + continue + entity = registry[resolved].entity + if entity == "person": + problems.append( + f"household_facts: {resolved} is a person's fact; give it under " + "that person" + ) + continue + unit_of[entity][resolved] = value + return None if problems else situation + + DESIRED_SUFFIX = "_desired" _LABEL_SEPARATORS = re.compile(r"[\s\-]+") @@ -1152,12 +1508,40 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic out["notes"] = list(request["notes"]) if request.get("explain"): out["explanations"] = { - name: explain_variable(system, request, name, simulation_cls) + name: ( + brief_explanation( + explain_variable(system, request, name, simulation_cls) + ) + if request.get("detail") == "brief" + else explain_variable(system, request, name, simulation_cls) + ) for name in request["explain"] } return out +BRIEF_INPUTS = 8 +BRIEF_PARAMETERS = 5 + + +def brief_explanation(full: dict) -> dict: + """The concise feature's explanation: the value, the variables it came + from directly (one level down, zero-valued ones left out), and the first + parameters it read with their sources.""" + if "error" in full: + return full + steps = full["steps"] + direct = [ + s.strip() for s in steps[1:] if s.startswith(" ") and not s.startswith(" ") + ] + return { + "value": steps[0].strip() if steps else "", + "from": direct[:BRIEF_INPUTS], + "parameters": full["parameters"][:BRIEF_PARAMETERS], + "more": 'set detail to "full" for every step', + } + + EXPLAIN_MAX_LINES = 60 EXPLAIN_MAX_DEPTH = 6 EXPLAIN_MAX_PARAMETERS = 15 @@ -1356,10 +1740,13 @@ def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> return if args is None: return - version = 1 - if isinstance(args, dict) and args.get("__tool_version__") == 2: - version, args = 2, args.get("arguments") - request, error = normalize_request(args, registry, version) + features = frozenset() + if isinstance(args, dict) and "__tool_features__" in args: + features = frozenset(args.get("__tool_features__") or ()) + args = args.get("arguments") + elif isinstance(args, dict) and args.get("__tool_version__") == 2: + features, args = V2_FEATURES, args.get("arguments") + request, error = normalize_request(args, registry, features=features) if error is not None: conn.send(("ok", error)) continue @@ -1379,9 +1766,15 @@ def __init__( acquire_timeout: float | None = None, board: Path | None = None, tool_version: int = 1, + tool_features: frozenset | None = None, ): self._ctx = mp.get_context("spawn") self._tool_version = tool_version + self._tool_features = ( + tool_features + if tool_features is not None + else (V2_FEATURES if tool_version == 2 else frozenset()) + ) self._fix_path = str(fix_path) # A worker gets the board only when it is not the frozen one, so a # three-argument worker target still starts on the default. @@ -1448,8 +1841,13 @@ def call(self, args) -> tuple[dict, float]: time.perf_counter() - started, ) process, conn = worker - if self._tool_version == 2: + if self._tool_features == V2_FEATURES: args = {"__tool_version__": 2, "arguments": args} + elif self._tool_features: + args = { + "__tool_features__": sorted(self._tool_features), + "arguments": args, + } try: conn.send(args) if conn.poll(self.timeout): @@ -1607,6 +2005,8 @@ def notice(self) -> str: notice = WEB_TOOL_NOTICE if tool else JSON_WEB_TOOL_NOTICE elif self.arm == "pe2": return engine_text(TOOL_NOTICE_V2 if tool else JSON_TOOL_NOTICE_V2) + elif self.arm != "pe" and arm_features(self.arm) is not None: + return engine_text(compose_tool_notice(arm_features(self.arm), tool)) else: return engine_text(TOOL_NOTICE if tool else JSON_TOOL_NOTICE) return engine_text(notice) + ( @@ -1615,11 +2015,15 @@ def notice(self) -> str: @property def uses_calculator(self) -> bool: - return self.arm in TOOL_VERSION_BY_ARM + return arm_features(self.arm) is not None + + @property + def tool_features(self) -> frozenset: + return arm_features(self.arm) or frozenset() @property def tool_version(self) -> int: - return TOOL_VERSION_BY_ARM.get(self.arm, 1) + return 2 if self.tool_features == V2_FEATURES else 1 @property def uses_web(self) -> bool: @@ -1637,7 +2041,7 @@ def arm1_label(self) -> str: def arm2_label(self) -> str: version = "" if self.prompt_version == "v1" else f"@{self.prompt_version}" rule = "+no-benchmark-search" if self.benchmark_search_rule else "" - return f"{self.name}{ARM_SUFFIX[self.arm]}{version}{rule}" + return f"{self.name}{arm_suffix(self.arm)}{version}{rule}" def provider_family(model_id: str) -> str: @@ -1663,8 +2067,11 @@ def model_profile( _uses_responses_api, ) - if arm not in ARMS: - raise SystemExit(f"unknown arm {arm!r}; the arms are {', '.join(ARMS)}") + if arm not in ARMS and arm_features(arm) is None: + raise SystemExit( + f"unknown arm {arm!r}; the arms are {', '.join(ARMS)}, or pe: plus " + f"calculator features joined by + ({', '.join(TOOL_FEATURES)})" + ) if prompt_version not in PROMPT_VERSIONS: raise SystemExit( f"unknown prompt {prompt_version!r}; the prompts are " @@ -2298,7 +2705,7 @@ def run_household( sent_board_prompt_sha256 = base.pop("_sent_board_prompt_sha256") ceiling = base.pop("_budget_ceiling") answer_tool = base.pop("_answer_tool") - calculate_tool = calculate_tool_schema(profile.tool_version) + calculate_tool = calculate_tool_schema(features=profile.tool_features) offered = ([web_search_tool(profile)] if profile.uses_web else []) + ( [calculate_tool] if profile.uses_calculator else [] ) @@ -3133,7 +3540,10 @@ def command_run(args) -> None: "api": profile.api, "answer_contract": profile.contract, "tool_notice": profile.notice, - "calculate_tool": calculate_tool_schema(profile.tool_version) + "calculate_tool": calculate_tool_schema(features=profile.tool_features) + if profile.uses_calculator + else None, + "tool_features": sorted(profile.tool_features) if profile.uses_calculator else None, "tool_version": profile.tool_version if profile.uses_calculator else None, @@ -3164,7 +3574,12 @@ def command_run(args) -> None: return # Arms without the calculator start no engine workers. pool = ( - ToolPool(size=args.workers, board=snapshot, tool_version=profile.tool_version) + ToolPool( + size=args.workers, + board=snapshot, + tool_version=profile.tool_version, + tool_features=profile.tool_features, + ) if profile.uses_calculator else None ) @@ -3879,6 +4294,12 @@ def _successful_calls(record: dict) -> list[dict]: def _person_names(arguments) -> list[str]: try: args = json.loads(arguments) if isinstance(arguments, str) else arguments + if "household" in args: # the flat feature's people list + household = args["household"] + household = ( + json.loads(household) if isinstance(household, str) else household + ) + return [person["name"] for person in household["people"]] situation = args["situation"] situation = json.loads(situation) if isinstance(situation, str) else situation return list(situation["people"]) @@ -4486,7 +4907,7 @@ def main_call(record: dict, scenario, outputs: list[str]) -> dict | None: def situation_fidelity( - records, board, system, registry, defaults, out: Path, version: int = 1 + records, board, system, registry, defaults, out: Path, features=frozenset() ) -> None: """How each household's main situation differs from the reference's. @@ -4508,7 +4929,7 @@ def situation_fidelity( if call is None: rows.append({"scenario_id": sid, "differences": None}) continue - request, _ = normalize_request(call["arguments"], registry, version) + request, _ = normalize_request(call["arguments"], registry, features=features) if request is None: rows.append({"scenario_id": sid, "differences": None}) continue @@ -4561,7 +4982,7 @@ def command_diagnose(args) -> None: out = Path(args.out_dir) records = load_records(out) # A run's calls are read as its own tool version read them. - version = TOOL_VERSION_BY_ARM.get(records_arm(records), 1) + features = arm_features(records_arm(records)) or frozenset() rows_path = out / "scored_rows.csv" if not rows_path.exists(): raise SystemExit("run score first") @@ -4577,7 +4998,7 @@ def command_diagnose(args) -> None: name: system.variables[name].default_value for name in system.variables } if registry is not None: - situation_fidelity(records, board, system, registry, defaults, out, version) + situation_fidelity(records, board, system, registry, defaults, out, features) out_rows = [] if misses.empty: print("no arm-2 misses") @@ -4595,7 +5016,7 @@ def command_diagnose(args) -> None: if registry is not None: def board_value_for(arguments, scenario=scenario, output=row.variable): - request, _ = normalize_request(arguments, registry, version) + request, _ = normalize_request(arguments, registry, features=features) if request is None: return None return output_value(system, request["situation"], scenario, output) @@ -4619,7 +5040,7 @@ def board_value_for(arguments, scenario=scenario, output=row.variable): attribution = {} arguments = verdict.get("call_arguments") if arguments is not None and registry is not None: - request, _ = normalize_request(arguments, registry, version) + request, _ = normalize_request(arguments, registry, features=features) if request is not None: reference_situation = scenario.to_pe_household() diff = situation_diff( @@ -4716,9 +5137,10 @@ def main(argv: list[str] | None = None) -> None: ) p.add_argument( "--arm", - choices=ARMS, default="pe", - help="the tools: none, PolicyEngine (pe), web search (web), or both", + help="the tools: none, PolicyEngine (pe, pe2, pe3, or pe: plus " + "calculator features joined by +), web search (web), or both " + "(web+pe)", ) p.add_argument( "--prompt", diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index 986a89d9..4e586374 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -2294,6 +2294,7 @@ def test_web_notices_name_search_and_keep_the_answer_instruction(): "web": "gpt-6-sol+web", "web+pe": "gpt-6-sol+web+pe-tool", "pe2": "gpt-6-sol+pe-tool-v2", + "pe3": "gpt-6-sol+pe-tool-v3", } with pytest.raises(SystemExit): pilot.model_profile(SOL, "search") diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py new file mode 100644 index 00000000..5c018fd1 --- /dev/null +++ b/tests/test_tool_v3.py @@ -0,0 +1,377 @@ +"""The calculator's features one at a time (tool v3 and the ablation arms). + +Each feature (names, placement, overrides, explain, concise, flat) must work +on its own and change nothing else, so a study can turn one on at a time. The +version 2 feature set must reproduce tool v2's notice byte for byte, and the +flat household format must build the units the benchmark's own reference +situations have. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) + +import tool_pilot_argon as pilot # noqa: E402 + +V = pilot.VariableInfo +REGISTRY = { + "age": V("person", "int", "age", "year"), + "employment_income": V("person", "float", "employment income", "year"), + "is_tax_unit_head": V("person", "bool", "head", "year"), + "is_tax_unit_spouse": V("person", "bool", "spouse", "year"), + "traditional_401k_contributions": V( + "person", "float", "401(k)", "year", computed=True + ), + "traditional_401k_contributions_desired": V("person", "float", "asked", "year"), + "snap": V("spm_unit", "float", "SNAP", "month", computed=True), + "pre_subsidy_rent": V("spm_unit", "float", "rent", "year"), + "qualified_passenger_vehicle_loan_interest": V( + "household", "float", "car loan", "year" + ), + "state_code": V("household", "Enum", "state", "year", ("CA", "LA")), + "tax_unit_childcare_expenses": V("tax_unit", "float", "childcare", "year"), + "marital_unit_flag": V("marital_unit", "bool", "flag", "year"), +} +F = pilot + + +def _situation(): + return { + "people": {"head": {"age": 40}, "kid": {"age": 8}}, + "spm_units": {"spm": {"members": ["head", "kid"]}}, + "households": {"h": {"members": ["head", "kid"], "state_code": "CA"}}, + } + + +def _norm(args, *features): + return pilot.normalize_request(args, REGISTRY, features=frozenset(features)) + + +def test_the_version_2_set_reproduces_tool_v2s_notice(): + assert F.compose_tool_notice(F.V2_FEATURES, True) == F.TOOL_NOTICE_V2 + assert F.compose_tool_notice(F.V2_FEATURES, False) == F.JSON_TOOL_NOTICE_V2 + assert F.compose_tool_notice(frozenset(), True) == F.TOOL_NOTICE + assert F.compose_tool_notice(frozenset(), False) == F.JSON_TOOL_NOTICE + + +def test_each_feature_adds_its_own_sentence_only(): + base = F.compose_tool_notice(frozenset({F.EXPLAIN}), True) + assert "`explain`" in base and "names it read" not in base + assert "written with spaces" not in base + names = F.compose_tool_notice(frozenset({F.NAMES}), True) + assert "written with spaces" in names and "`explain`" not in names + flat = F.compose_tool_notice(frozenset({F.FLAT}), True) + assert "household you describe" in flat and "`household`" in flat + concise = F.compose_tool_notice(frozenset({F.EXPLAIN, F.CONCISE}), True) + assert '`detail` set to "full"' in concise + + +def test_arms_name_their_features(): + assert F.arm_features("pe") == frozenset() + assert F.arm_features("pe2") == F.V2_FEATURES + assert F.arm_features("pe3") == frozenset(F.TOOL_FEATURES) + assert F.arm_features("pe:flat+names") == frozenset({"flat", "names"}) + assert F.arm_features("none") is None and F.arm_features("web") is None + assert F.arm_suffix("pe:flat+names") == "+pe-tool[names+flat]" + with pytest.raises(SystemExit, match="unknown tool features"): + F.arm_features("pe:names+telepathy") + with pytest.raises(SystemExit, match="concise feature needs explain"): + F.arm_features("pe:concise") + profile = F.ModelProfile("m", "gpt-x", "chat", "tool", arm="pe:names") + assert profile.uses_calculator and profile.tool_features == frozenset({"names"}) + assert profile.tool_version == 1 + assert "written with spaces" in profile.notice + assert profile.arm2_label == "m+pe-tool[names]" + + +def test_names_alone_reads_labels_but_places_and_redirects_nothing(): + situation = _situation() + situation["people"]["head"]["Employment income"] = 30000 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.NAMES) + assert error is None + assert request["situation"]["people"]["head"]["employment_income"] == { + "2026": 30000 + } + misplaced = _situation() + misplaced["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 5 + request, error = _norm({"situation": misplaced, "variables": ["snap"]}, F.NAMES) + assert request is None and any( + "set it under households" in p for p in error["problems"] + ) + realized = _situation() + realized["people"]["head"]["traditional_401k_contributions"] = 100 + request, error = _norm({"situation": realized, "variables": ["snap"]}, F.NAMES) + assert "traditional_401k_contributions" in request["situation"]["people"]["head"] + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, F.NAMES + ) + assert request is None and any( + "unknown arguments: explain" in p for p in error["problems"] + ) + + +def test_placement_alone_moves_inputs_but_reads_no_labels(): + situation = _situation() + situation["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 5 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.PLACEMENT) + assert error is None + assert request["situation"]["households"]["h"][ + "qualified_passenger_vehicle_loan_interest" + ] == {"2026": 5} + labeled = _situation() + labeled["people"]["head"]["employment income"] = 1 + request, error = _norm({"situation": labeled, "variables": ["snap"]}, F.PLACEMENT) + assert request is None + + +def test_overrides_alone_redirects_and_flags(): + situation = _situation() + situation["people"]["head"]["traditional_401k_contributions"] = 100 + situation["spm_units"]["spm"]["snap"] = 5 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.OVERRIDES) + assert error is None + head = request["situation"]["people"]["head"] + assert head == { + "age": {"2026": 40}, + "traditional_401k_contributions_desired": {"2026": 100}, + } + assert any("snap is computed" in n for n in request["notes"]) + + +def test_concise_takes_a_detail_and_defaults_to_brief(): + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, + F.EXPLAIN, + F.CONCISE, + ) + assert error is None and request["detail"] == "brief" + request, _ = _norm( + { + "situation": _situation(), + "variables": ["snap"], + "explain": ["snap"], + "detail": "full", + }, + F.EXPLAIN, + F.CONCISE, + ) + assert request["detail"] == "full" + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "detail": "short"}, + F.EXPLAIN, + F.CONCISE, + ) + assert request is None + request, _ = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, + F.EXPLAIN, + ) + assert request["detail"] == "full" + + +def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_parameters(): + full = { + "steps": [ + "snap (2026) = 288", + " snap_normal_allotment (2026) = 288", + " snap_net_income (2026) = 1700", + " is_snap_eligible (2026) = True", + ], + "parameters": [{"name": f"p{i}", "value": i, "sources": []} for i in range(9)], + } + brief = F.brief_explanation(full) + assert brief["value"] == "snap (2026) = 288" + assert brief["from"] == [ + "snap_normal_allotment (2026) = 288", + "is_snap_eligible (2026) = True", + ] + assert [p["name"] for p in brief["parameters"]] == ["p0", "p1", "p2", "p3", "p4"] + assert F.brief_explanation({"error": "x"}) == {"error": "x"} + + +def _flat(**overrides): + household = { + "state": "LA", + "people": [ + { + "name": "Head", + "role": "head", + "facts": {"age": 40, "employment_income": 30000}, + }, + {"name": "Spouse", "role": "spouse", "facts": {"age": 38}}, + {"name": "Kid", "role": "dependent", "facts": {"age": 6}}, + ], + "household_facts": { + "pre_subsidy_rent": 9000, + "tax_unit_childcare_expenses": 1200, + "qualified_passenger_vehicle_loan_interest": 300, + }, + } + household.update(overrides) + return {"household": json.dumps(household), "variables": ["snap"]} + + +def test_flat_builds_the_benchmarks_units(): + request, error = _norm(_flat(), F.FLAT) + assert error is None, error + s = request["situation"] + assert s["marital_units"] == { + "marital_unit_1": {"members": ["Head", "Spouse"]}, + "marital_unit_2": {"members": ["Kid"]}, + } + for group, unit in ( + ("tax_units", "tax_unit"), + ("spm_units", "spm_unit"), + ("families", "family"), + ): + assert s[group][unit]["members"] == ["Head", "Spouse", "Kid"] + assert s["households"]["household"]["state_code"] == {"2026": "LA"} + assert s["spm_units"]["spm_unit"]["pre_subsidy_rent"] == {"2026": 9000} + assert s["tax_units"]["tax_unit"]["tax_unit_childcare_expenses"] == {"2026": 1200} + assert s["households"]["household"][ + "qualified_passenger_vehicle_loan_interest" + ] == {"2026": 300} + assert s["people"]["Head"]["is_tax_unit_head"] == {"2026": True} + assert s["people"]["Spouse"]["is_tax_unit_spouse"] == {"2026": True} + assert s["people"]["Kid"]["is_tax_unit_head"] == {"2026": False} + + +def test_flat_refuses_what_it_cannot_place(): + people = [ + {"name": "A", "role": "head", "facts": {}}, + {"name": "B", "role": "head", "facts": {}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("exactly one head" in p for p in error["problems"]) + request, error = _norm(_flat(household_facts={"age": 3}), F.FLAT) + assert any("person's fact" in p for p in error["problems"]) + people = [{"name": "A", "role": "boss", "facts": {}}] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("role must be one of" in p for p in error["problems"]) + request, error = _norm({"situation": _situation(), "variables": ["snap"]}, F.FLAT) + assert any("unknown arguments: situation" in p for p in error["problems"]) + request, error = _norm(_flat(), F.NAMES) + assert any("unknown arguments: household" in p for p in error["problems"]) + + +def test_flat_with_names_reads_labeled_household_facts(): + request, error = _norm( + _flat(household_facts={"pre subsidy rent": 9000}), F.FLAT, F.NAMES + ) + assert error is None + assert request["situation"]["spm_units"]["spm_unit"]["pre_subsidy_rent"] == { + "2026": 9000 + } + assert any("as pre_subsidy_rent" in n for n in request["notes"]) + + +def test_the_schema_offers_each_features_arguments(): + flat = F.calculate_tool_schema(features=frozenset({F.FLAT}))["function"] + props = flat["parameters"]["properties"] + assert "household" in props and "situation" not in props + assert flat["parameters"]["required"] == ["household", "variables"] + concise = F.calculate_tool_schema(features=frozenset({F.EXPLAIN, F.CONCISE})) + props = concise["function"]["parameters"]["properties"] + assert "explain" in props and props["detail"]["enum"] == ["brief", "full"] + assert F.calculate_tool_schema(features=F.V2_FEATURES) == F.calculate_tool_schema(2) + assert F.calculate_tool_schema(features=frozenset()) == F.calculate_tool_schema(1) + + +def _echo_worker(conn, fix_path, reform): + conn.send(("ready", None)) + while True: + args = conn.recv() + if args is None: + return + conn.send(("ok", {"received": args})) + + +def test_the_pool_tells_its_workers_the_features(): + pool = F.ToolPool( + size=1, + timeout=5.0, + worker_target=_echo_worker, + tool_features=frozenset({F.FLAT, F.NAMES}), + ) + try: + result, _ = pool.call({"x": 1}) + assert result == { + "received": {"__tool_features__": ["flat", "names"], "arguments": {"x": 1}} + } + finally: + pool.close() + + +@pytest.mark.slow +def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): + """Differential: on the installed engine, each frozen-board household + described in the flat format computes every board output exactly as the + benchmark's own situation (Scenario.to_pe_household) does.""" + import household_set as hs + import pandas as pd + from policyengine_us import Simulation + + from policybench.ground_truth import _extract_person_value, _pe_variable_for_output + + system = F.build_tool_system(reform=True) + registry = F.build_registry(system) + reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") + variables = reference.groupby("scenario_id").variable.agg(list) + relations = {"is_tax_unit_head", "is_tax_unit_spouse"} + checked = 0 + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:8]: + situation = hs.situation_for(scenario) + couple = scenario.marital_couple() or () + head = couple[0] if couple else scenario.adults[0].name + people = [ + { + "name": name, + "role": "head" + if name == head + else ("spouse" if name in couple else "dependent"), + "facts": {k: v for k, v in facts.items() if k not in relations}, + } + for name, facts in situation["people"].items() + ] + household_facts = {} + for group in ("tax_units", "spm_units", "families", "households"): + for unit in situation[group].values(): + household_facts.update( + { + k: v + for k, v in unit.items() + if k not in ("members", "state_code") + } + ) + household = { + "state": scenario.state, + "people": people, + "household_facts": household_facts, + } + request, error = F.normalize_request( + {"household": json.dumps(household), "variables": ["snap"]}, + registry, + features=frozenset({F.FLAT}), + ) + assert error is None, (scenario.id, error) + own = Simulation(tax_benefit_system=system, situation=situation) + flat = Simulation(tax_benefit_system=system, situation=request["situation"]) + for variable in variables[scenario.id]: + name = _pe_variable_for_output(variable, "us") + values = [ + float( + _extract_person_value( + simulation.calculate(name, 2026), scenario, variable + ) + ) + for simulation in (own, flat) + ] + assert values[0] == values[1], (scenario.id, variable, values) + checked += 1 + assert checked > 80 From b3ec406e62b19d59f285150ddc8a6207a2cf0bf0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:25:47 -0400 Subject: [PATCH 09/16] Read facts as the prompt labels them (labels feature) The prompts label some facts in words rather than as the variable name with spaces: employment_income is "gross wages and salaries", hours_worked_last_week "usual weekly hours worked", has_esi "has employer-sponsored insurance" (policybench.prompts.INPUT_LABEL_OVERRIDES). In the 20-household test set nearly every household has at least one; hours appear in 17, employer coverage and premiums in 11 each. Tool version 2's names feature cannot read these, and its notice says the facts carry the variable names. A Flash-Lite smoke test of the flat format met three such errors and fell back to setting employer_sponsored_insurance_premiums to 0 and weekly_hours_worked, which the SNAP work rules do not read. The new labels feature reads a fact under its prompt label (case and punctuation aside, renames applied) and, as the reference situations do, also sets stated usual weekly hours as weekly_hours_worked_before_lsr (PE_INPUT_ALIASES). pe3 now offers all seven features. pe2 is unchanged and keeps its notice as it ran; every other feature set's notice says only what the tool reads. Tests: label reading, the alias (and leaving a set value alone), pe2 unchanged, the variable-name claim only in pe2's notice across every feature set, every prompt label reading back to its engine name, and a differential: twelve frozen-board households described flat with each fact as its prompt writes it compute every board output as the benchmark does. Three mutants (no alias, no renames, strict label matching) are caught. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 130 +++++++++++++++++++++------ tests/test_tool_v3.py | 175 ++++++++++++++++++++++++++++++++++++ 2 files changed, 279 insertions(+), 26 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 88d8c4d6..58ec78bb 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -51,6 +51,7 @@ import copy import decimal import difflib +import functools import hashlib import importlib.metadata import importlib.util @@ -166,6 +167,12 @@ # publishable"; "do all these tests with a cheap model"): # - names: a name written with spaces, or a benchmark output key, is read as # its engine variable; +# - labels: a fact written under the label the prompt gives it (a prompt +# labels some facts in words, e.g. employment_income as "gross wages and +# salaries"; policybench.prompts.INPUT_LABEL_OVERRIDES) is read as its +# variable, and a fact the reference also sets under another name +# (policybench.scenarios.PE_INPUT_ALIASES: usual weekly hours) is set there +# too; # - placement: an input set on another entity moves to the one unit holding # its people; # - overrides: a realized contribution is read as its _desired input, and an @@ -174,18 +181,20 @@ # - concise: an explanation is brief unless `detail` is "full"; # - flat: the household is described as people with roles and facts # (`household`), and the tool builds its units as the benchmark does. -# "pe2" is tool version 2 (the first four); "pe3" is all six. Any other set +# "pe2" is tool version 2 (names, placement, overrides and explain); "pe3" is +# all seven. Any other set # runs as an arm named "pe:" plus its features joined by "+", e.g. # "pe:names+flat". -NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( +NAMES, LABELS, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( "names", + "labels", "placement", "overrides", "explain", "concise", "flat", ) -TOOL_FEATURES = (NAMES, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT) +TOOL_FEATURES = (NAMES, LABELS, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT) V2_FEATURES = frozenset({NAMES, PLACEMENT, OVERRIDES, EXPLAIN}) FEATURES_BY_ARM = { "pe": frozenset(), @@ -251,10 +260,14 @@ def _listing(items: list[str]) -> str: def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: """The calculator notice for a set of features: the opening, a sentence per feature, then the turn limit and the answer instruction. No feature - gives TOOL_NOTICE and the version 2 features give TOOL_NOTICE_V2, byte for - byte.""" + gives TOOL_NOTICE and the version 2 features TOOL_NOTICE_V2, as those arms + ran. TOOL_NOTICE_V2 says the facts are labeled with the engine's variable + names, which holds for most facts but not those the prompt labels in words + ("gross wages and salaries"); every other set says only what it reads.""" if not features: return TOOL_NOTICE if tool_contract else JSON_TOOL_NOTICE + if features == V2_FEATURES: + return TOOL_NOTICE_V2 if tool_contract else JSON_TOOL_NOTICE_V2 what = ( "household you describe" if FLAT in features @@ -271,14 +284,24 @@ def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: "or dependent) and facts, and the facts about the household as a " "whole; it builds the household's units as the facts imply." ) - if NAMES in features: + if LABELS in features and NAMES in features: + text += ( + " It accepts the facts above under the labels they have there, its " + "variable names written with spaces, and the output keys you are asked " + "for. Set every listed fact." + ) + elif LABELS in features: + text += ( + " It accepts the facts above under the labels they have there. Set " + "every listed fact." + ) + elif NAMES in features: text += ( - " The facts above are labeled with its variable names, written with " - "spaces; it accepts them as written, and it accepts the output keys you " - "are asked for. Set every listed fact." + " It accepts its variable names written with spaces, and the output " + "keys you are asked for. Set every listed fact." ) notes = [] - if NAMES in features: + if NAMES in features or LABELS in features: notes.append("which names it read") if PLACEMENT in features: notes.append("which inputs it moved to the entity they belong to") @@ -670,10 +693,15 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - FLAT_DESCRIPTION if FLAT in features else CALCULATE_DESCRIPTION ) notes = [] + if LABELS in features: + notes.append( + "A fact may be written under the label it has in the facts " + "you were given." + ) if NAMES in features: notes.append( - "A name may also be written with spaces (as the facts are " - "labeled) or be one of the output keys you are asked for." + "A name may also be written with spaces or be one of the " + "output keys you are asked for." ) if PLACEMENT in features and FLAT not in features: notes.append( @@ -927,13 +955,8 @@ def _normalize_inputs( for name, value in data.items(): if allow_members and name == "members": continue - if ( - v2 is not None - and NAMES in v2["features"] - and isinstance(name, str) - and name not in registry - ): - resolved, note = resolve_name(name, registry, output_names=False) + if v2 is not None and isinstance(name, str) and name not in registry: + resolved, note = read_fact_name(name, registry, v2["features"]) if resolved is not None: v2["notes"].append(f"{where}: {note}") name = resolved @@ -1157,6 +1180,8 @@ def fail() -> tuple[None, dict]: problems.append(f"at most {MAX_VARIABLES} variables per call") if v2 is not None and PLACEMENT in features: _place_misplaced(v2, normalized, situation, registry, period, problems) + if v2 is not None and LABELS in features: + _set_aliases(v2, normalized["people"], registry) requested = [] for name in variables: if ( @@ -1320,13 +1345,8 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None for name, value in facts.items(): resolved = name - if ( - isinstance(name, str) - and name not in registry - and v2 is not None - and NAMES in v2["features"] - ): - resolved, note = resolve_name(name, registry, output_names=False) + if isinstance(name, str) and name not in registry and v2 is not None: + resolved, note = read_fact_name(name, registry, v2["features"]) if resolved is not None: v2["notes"].append(f"household_facts: {note}") if not isinstance(resolved, str) or resolved not in registry: @@ -1375,6 +1395,64 @@ def output_reading(name: str) -> tuple[str, str] | None: return engine, how +def _label_key(text: str) -> str: + """A label as compared: lower case, letters and digits only, so "Roth + 401(k) contributions" and "roth 401k contributions" read alike.""" + return re.sub(r"[^0-9a-z]", "", text.casefold()) + + +@functools.lru_cache(maxsize=1) +def prompt_labels() -> dict[str, str]: + """The labels the prompts give facts in words, keyed by _label_key, as the + engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied). A + label two facts share is left out.""" + from policybench.prompts import INPUT_LABEL_OVERRIDES + from policybench.scenarios import PE_INPUT_RENAMES + + labels: dict[str, str] = {} + shared = set() + for source, label in INPUT_LABEL_OVERRIDES.items(): + key = _label_key(label) + name = PE_INPUT_RENAMES.get(source, source) + if key in labels and labels[key] != name: + shared.add(key) + labels[key] = name + for source, name in PE_INPUT_RENAMES.items(): + labels.setdefault(_label_key(source), name) + return {key: name for key, name in labels.items() if key not in shared} + + +def read_fact_name( + name: str, registry: Mapping[str, VariableInfo], features: frozenset +) -> tuple[str | None, str | None]: + """A fact's engine variable under the labels feature (the label the prompt + gives it) and then the names feature (spaces for underscores). Returns + (engine name, note) or (None, None).""" + if LABELS in features: + labeled = prompt_labels().get(_label_key(name)) + if labeled is not None and labeled in registry: + return labeled, f"read {name!r} as {labeled}" + if NAMES in features: + return resolve_name(name, registry, output_names=False) + return None, None + + +def _set_aliases(v2: dict, people: dict, registry: Mapping[str, VariableInfo]): + """Set each fact the reference also sets under another name there too + (policybench.scenarios.PE_INPUT_ALIASES: the SNAP work rules read stated + usual weekly hours as weekly_hours_worked_before_lsr), unless it is set.""" + from policybench.scenarios import PE_INPUT_ALIASES + + for person, inputs in people.items(): + for source, target in PE_INPUT_ALIASES.items(): + if source in inputs and target in registry and target not in inputs: + inputs[target] = dict(inputs[source]) + v2["notes"].append( + f"people.{person}: also set {target} from {source}, as the " + "benchmark's reference does" + ) + + def resolve_name( name: str, registry: Mapping[str, VariableInfo], output_names: bool ) -> tuple[str | None, str | None]: diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index 5c018fd1..c3a5eb36 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -9,6 +9,7 @@ from __future__ import annotations +import itertools import json import sys from pathlib import Path @@ -37,6 +38,10 @@ "state_code": V("household", "Enum", "state", "year", ("CA", "LA")), "tax_unit_childcare_expenses": V("tax_unit", "float", "childcare", "year"), "marital_unit_flag": V("marital_unit", "bool", "flag", "year"), + "has_esi": V("person", "bool", "ESI", "year"), + "hours_worked_last_week": V("person", "float", "hours", "year"), + "weekly_hours_worked_before_lsr": V("person", "float", "hours", "year"), + "partnership_self_employment_net_earnings": V("person", "float", "p", "year"), } F = pilot @@ -76,6 +81,7 @@ def test_arms_name_their_features(): assert F.arm_features("pe") == frozenset() assert F.arm_features("pe2") == F.V2_FEATURES assert F.arm_features("pe3") == frozenset(F.TOOL_FEATURES) + assert len(F.TOOL_FEATURES) == 7 and F.LABELS in F.arm_features("pe3") assert F.arm_features("pe:flat+names") == frozenset({"flat", "names"}) assert F.arm_features("none") is None and F.arm_features("web") is None assert F.arm_suffix("pe:flat+names") == "+pe-tool[names+flat]" @@ -195,6 +201,82 @@ def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_paramete assert F.brief_explanation({"error": "x"}) == {"error": "x"} +def test_labels_reads_the_facts_as_the_prompt_labels_them(): + situation = _situation() + situation["people"]["head"].update( + { + "gross wages and salaries": 30000, + "Has Employer-Sponsored Insurance": True, + "usual weekly hours worked": 40, + "self-employment partnership income": 5, + } + ) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None, error + head = request["situation"]["people"]["head"] + assert head["employment_income"] == {"2026": 30000} + assert head["has_esi"] == {"2026": True} + assert head["hours_worked_last_week"] == {"2026": 40} + assert head["partnership_self_employment_net_earnings"] == {"2026": 5} + # The reference also sets stated hours where the SNAP work rules read them. + assert head["weekly_hours_worked_before_lsr"] == {"2026": 40} + assert any("also set weekly_hours_worked_before_lsr" in n for n in request["notes"]) + # A name written with spaces is the names feature's, not this one's. + spaced = _situation() + spaced["people"]["head"]["employment income"] = 1 + request, error = _norm({"situation": spaced, "variables": ["snap"]}, F.LABELS) + assert request is None and "unknown variable 'employment income'" in str(error) + + +def test_labels_leaves_hours_already_set_where_the_rules_read_them(): + situation = _situation() + situation["people"]["head"].update( + {"hours_worked_last_week": 40, "weekly_hours_worked_before_lsr": 20} + ) + request, _ = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + head = request["situation"]["people"]["head"] + assert head["weekly_hours_worked_before_lsr"] == {"2026": 20} + + +def test_pe2_reads_no_labels_and_sets_no_alias_as_it_ran(): + situation = _situation() + situation["people"]["head"]["gross wages and salaries"] = 30000 + request, error = _norm( + {"situation": situation, "variables": ["snap"]}, *F.V2_FEATURES + ) + assert request is None and "gross wages and salaries" in str(error) + situation = _situation() + situation["people"]["head"]["hours_worked_last_week"] = 40 + request, _ = _norm({"situation": situation, "variables": ["snap"]}, *F.V2_FEATURES) + assert ( + "weekly_hours_worked_before_lsr" not in request["situation"]["people"]["head"] + ) + + +def test_only_pe2_says_the_facts_carry_variable_names(): + claim = "labeled with its variable names" + assert claim in F.TOOL_NOTICE_V2 + for size in range(1, len(F.TOOL_FEATURES) + 1): + for features in itertools.combinations(F.TOOL_FEATURES, size): + features = frozenset(features) + if F.CONCISE in features and F.EXPLAIN not in features: + continue + notice = F.compose_tool_notice(features, True) + assert (claim in notice) == (features == F.V2_FEATURES), features + assert notice.endswith("as instructed above.") + assert ("names it read" in notice) == bool(features & {F.NAMES, F.LABELS}) + + +def test_every_prompt_label_reads_back_to_its_engine_name(): + from policybench.prompts import INPUT_LABEL_OVERRIDES + from policybench.scenarios import PE_INPUT_RENAMES + + labels = F.prompt_labels() + for field, label in INPUT_LABEL_OVERRIDES.items(): + assert labels[F._label_key(label)] == PE_INPUT_RENAMES.get(field, field) + assert labels[F._label_key("roth 401k contributions")] == "roth_401k_contributions" + + def _flat(**overrides): household = { "state": "LA", @@ -375,3 +457,96 @@ def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): assert values[0] == values[1], (scenario.id, variable, values) checked += 1 assert checked > 80 + + +@pytest.mark.slow +def test_a_household_described_under_the_prompts_labels_computes_the_same(): + """Differential: each frozen-board household described flat with every + fact as its prompt writes it (a label in words, or the name with spaces), + and stated hours only where the prompt states them, computes every board + output as the benchmark does. Take-up conventions, which the prompt does + not list, keep their engine names.""" + import household_set as hs + import pandas as pd + from policyengine_us import Simulation + + from policybench.ground_truth import _extract_person_value, _pe_variable_for_output + from policybench.prompts import _humanize_input_label + from policybench.scenarios import PE_INPUT_ALIASES, PE_INPUT_RENAMES + + system = F.build_tool_system(reform=True) + registry = F.build_registry(system) + reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") + variables = reference.groupby("scenario_id").variable.agg(list) + field_of = {name: field for field, name in PE_INPUT_RENAMES.items()} + dropped = {"is_tax_unit_head", "is_tax_unit_spouse", *PE_INPUT_ALIASES.values()} + + worded = [] + + def labeled(facts): + out = {} + for k, v in facts.items(): + if k in dropped: + continue + if k.startswith(F.CONVENTION_PREFIXES): + out[k] = v + continue + label = _humanize_input_label(field_of.get(k, k)) + if label != k.replace("_", " "): + worded.append(label) + out[label] = v + return out + + checked = 0 + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:12]: + situation = hs.situation_for(scenario) + couple = scenario.marital_couple() or () + head = couple[0] if couple else scenario.adults[0].name + people = [ + { + "name": name, + "role": "head" + if name == head + else ("spouse" if name in couple else "dependent"), + "facts": labeled(facts), + } + for name, facts in situation["people"].items() + ] + household_facts = {} + for group in ("tax_units", "spm_units", "families", "households"): + for unit in situation[group].values(): + household_facts.update( + labeled( + { + k: v + for k, v in unit.items() + if k not in ("members", "state_code") + } + ) + ) + household = { + "state": scenario.state, + "people": people, + "household_facts": household_facts, + } + request, error = F.normalize_request( + {"household": json.dumps(household), "variables": ["snap"]}, + registry, + features=frozenset({F.FLAT, F.LABELS, F.NAMES}), + ) + assert error is None, (scenario.id, error) + own = Simulation(tax_benefit_system=system, situation=situation) + flat = Simulation(tax_benefit_system=system, situation=request["situation"]) + for variable in variables[scenario.id]: + name = _pe_variable_for_output(variable, "us") + values = [ + float( + _extract_person_value( + simulation.calculate(name, 2026), scenario, variable + ) + ) + for simulation in (own, flat) + ] + assert values[0] == values[1], (scenario.id, variable, values) + checked += 1 + assert checked > 120 and len(worded) > 10, (checked, worded) From fd3f5d096562754da767819bf6334a84449acf73 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:31:44 -0400 Subject: [PATCH 10/16] Labels: read the v2-pilot relabelings; flat: names optional The v2-pilot prompt relabels two facts itself ("home energy assistance received" for spm_unit_energy_subsidy, "employer-paid health insurance premiums ..." for employer_sponsored_insurance_premiums); the labels feature now reads them, derived from V2_PILOT_CONVENTIONS so the two stay in step. In the flat format a person given no name is named by role (head, spouse, dependent_1, ...): Flash-Lite's first flat call in two smoke households left the names out. The label differential now writes each fact as the v2-pilot prompt does. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 46 +++++++++++++++++++++++++++---------- tests/test_tool_v3.py | 44 +++++++++++++++++++++++++++++++++-- 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 58ec78bb..775725c9 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -664,9 +664,10 @@ def refuse_policyengine(label: str, value) -> None: FLAT_DESCRIPTION = ( "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " "the requested variables. `household` is a JSON object: `state` (the " - "two-letter state code), `people` (a list, each with `name`, `role` -- " - 'head, spouse or dependent -- and `facts`, for example {"age": 40, ' - '"employment income": 30000}), and `household_facts` (facts about the ' + "two-letter state code), `people` (a list, each with `role` -- head, " + 'spouse or dependent --, `facts`, for example {"age": 40, "employment ' + 'income": 30000}, and optionally `name`, by default the role, numbered for ' + "dependents), and `household_facts` (facts about the " "household as a whole: its tax unit, family, SPM unit or home). The head " "and spouse are one couple and everyone else is on their own; everyone " "shares one tax unit, family, SPM unit and household. A fact is a value " @@ -1248,7 +1249,8 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): reference situations have (policybench.scenarios.Scenario.to_pe_household): the head and spouse one marital unit and everyone else alone, and one tax unit, SPM unit, family and household holding everyone. A household fact - goes to the unit of its variable's entity. Returns None after recording + goes to the unit of its variable's entity. A person given no name is named + by role (head, spouse, dependent_1, ...). Returns None after recording problems.""" if isinstance(household, str): if len(household) > MAX_SITUATION_CHARS: @@ -1279,6 +1281,7 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None situation: dict[str, Any] = {"people": {}} roles: dict[str, str] = {} + dependents = 0 for index, person in enumerate(people): if not isinstance(person, dict): problems.append( @@ -1290,17 +1293,23 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): person.get("role"), person.get("facts", {}), ) + if role not in FLAT_ROLES: + problems.append( + f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" + ) + continue + if role == "dependent": + dependents += 1 + if name is None: + name = role if role != "dependent" else f"dependent_{dependents}" + if v2 is not None: + v2["notes"].append(f"people[{index}]: named {name}") if not isinstance(name, str) or not name or len(name) > 64: problems.append(f"people[{index}].name must be a short string") continue if name in roles: problems.append(f"people[{index}]: {name!r} is listed twice") continue - if role not in FLAT_ROLES: - problems.append( - f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" - ) - continue if not isinstance(facts, dict): problems.append(f"people[{index}].facts must be an object") continue @@ -1404,14 +1413,15 @@ def _label_key(text: str) -> str: @functools.lru_cache(maxsize=1) def prompt_labels() -> dict[str, str]: """The labels the prompts give facts in words, keyed by _label_key, as the - engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied). A - label two facts share is left out.""" + engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied): + the board's labels and the v2-pilot prompt's relabelings. A label two facts + share is left out.""" from policybench.prompts import INPUT_LABEL_OVERRIDES from policybench.scenarios import PE_INPUT_RENAMES labels: dict[str, str] = {} shared = set() - for source, label in INPUT_LABEL_OVERRIDES.items(): + for source, label in {**INPUT_LABEL_OVERRIDES, **v2_pilot_labels()}.items(): key = _label_key(label) name = PE_INPUT_RENAMES.get(source, source) if key in labels and labels[key] != name: @@ -1422,6 +1432,18 @@ def prompt_labels() -> dict[str, str]: return {key: name for key, name in labels.items() if key not in shared} +def v2_pilot_labels() -> dict[str, str]: + """The facts the v2-pilot prompt relabels (a V2_PILOT_CONVENTIONS rewrite + of ": " to ": "), by input name.""" + relabeled = {} + for pattern, replacement in V2_PILOT_CONVENTIONS: + old = pattern.pattern + if pattern.groups or not old.endswith(": ") or not replacement.endswith(": "): + continue + relabeled[old[:-2].replace(" ", "_")] = replacement[:-2] + return relabeled + + def read_fact_name( name: str, registry: Mapping[str, VariableInfo], features: frozenset ) -> tuple[str | None, str | None]: diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index c3a5eb36..d8903c93 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -274,6 +274,16 @@ def test_every_prompt_label_reads_back_to_its_engine_name(): labels = F.prompt_labels() for field, label in INPUT_LABEL_OVERRIDES.items(): assert labels[F._label_key(label)] == PE_INPUT_RENAMES.get(field, field) + # The v2-pilot prompt's own relabelings, as that prompt writes them. + relabeled = F.v2_pilot_labels() + assert set(relabeled) == { + "employer_sponsored_insurance_premiums", + "spm_unit_energy_subsidy", + } + for field, label in relabeled.items(): + written = F.v2_pilot_text(f"- {field.replace('_', ' ')}: $5") + assert written == f"- {label}: $5" + assert labels[F._label_key(label)] == field assert labels[F._label_key("roth 401k contributions")] == "roth_401k_contributions" @@ -324,6 +334,33 @@ def test_flat_builds_the_benchmarks_units(): assert s["people"]["Kid"]["is_tax_unit_head"] == {"2026": False} +def test_flat_names_a_person_by_role_when_no_name_is_given(): + people = [ + {"role": "head", "facts": {"age": 40}}, + {"role": "dependent", "facts": {"age": 6}}, + {"name": "Kid", "role": "dependent", "facts": {"age": 4}}, + {"role": "dependent", "facts": {"age": 2}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert error is None, error + assert list(request["situation"]["people"]) == [ + "head", + "dependent_1", + "Kid", + "dependent_3", + ] + assert "people[0]: named head" in request["notes"] + people = [ + {"role": "head", "facts": {}}, + {"name": "head", "role": "dependent", "facts": {}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("listed twice" in p for p in error["problems"]) + people = [{"name": "", "role": "head", "facts": {}}] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("name must be a short string" in p for p in error["problems"]) + + def test_flat_refuses_what_it_cannot_place(): people = [ {"name": "A", "role": "head", "facts": {}}, @@ -462,7 +499,8 @@ def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): @pytest.mark.slow def test_a_household_described_under_the_prompts_labels_computes_the_same(): """Differential: each frozen-board household described flat with every - fact as its prompt writes it (a label in words, or the name with spaces), + fact as its v2-pilot prompt writes it (a label in words, or the name with + spaces), and stated hours only where the prompt states them, computes every board output as the benchmark does. Take-up conventions, which the prompt does not list, keep their engine names.""" @@ -491,7 +529,9 @@ def labeled(facts): if k.startswith(F.CONVENTION_PREFIXES): out[k] = v continue - label = _humanize_input_label(field_of.get(k, k)) + label = F.v2_pilot_text(f"- {_humanize_input_label(field_of.get(k, k))}: ")[ + 2:-2 + ] if label != k.replace("_", " "): worded.append(label) out[label] = v From 5c0f3ec5843b4a839b8255f9836a359e4d6f3c80 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:39:43 -0400 Subject: [PATCH 11/16] Flat: a household of one's person fact goes to that person Two of three Flash-Lite smoke households gave the head's rent as a household fact. pre_subsidy_rent is a person variable in policyengine-us, so the flat format refused it. Under the placement feature a person fact given for a household of one now goes to that person, as placement moves a person variable off a one-person unit; with two or more people it is still refused, now saying to give it under the person whose facts list it. The test registry's entities now match the engine's (pre_subsidy_rent is a person variable). Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 18 ++++++++++++++++- tests/test_tool_v3.py | 39 +++++++++++++++++++++++++++++-------- 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 775725c9..e3e3f1b4 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1366,9 +1366,25 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): continue entity = registry[resolved].entity if entity == "person": + # As placement moves a person variable off a one-person unit: only + # a household of one says whose fact it is. + if v2 is not None and PLACEMENT in v2["features"] and len(everyone) == 1: + facts_of = situation["people"][everyone[0]] + if resolved in facts_of and facts_of[resolved] != value: + problems.append( + f"household_facts: {resolved} is also set for " + f"{everyone[0]} with a different value" + ) + continue + facts_of[resolved] = value + v2["notes"].append( + f"household_facts: moved {resolved}, a person's fact, to " + f"{everyone[0]}" + ) + continue problems.append( f"household_facts: {resolved} is a person's fact; give it under " - "that person" + "the person whose facts list it" ) continue unit_of[entity][resolved] = value diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index d8903c93..be60efdf 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -31,7 +31,8 @@ ), "traditional_401k_contributions_desired": V("person", "float", "asked", "year"), "snap": V("spm_unit", "float", "SNAP", "month", computed=True), - "pre_subsidy_rent": V("spm_unit", "float", "rent", "year"), + "pre_subsidy_rent": V("person", "float", "rent", "year"), + "spm_unit_energy_subsidy": V("spm_unit", "float", "energy", "year"), "qualified_passenger_vehicle_loan_interest": V( "household", "float", "car loan", "year" ), @@ -300,7 +301,7 @@ def _flat(**overrides): {"name": "Kid", "role": "dependent", "facts": {"age": 6}}, ], "household_facts": { - "pre_subsidy_rent": 9000, + "spm_unit_energy_subsidy": 900, "tax_unit_childcare_expenses": 1200, "qualified_passenger_vehicle_loan_interest": 300, }, @@ -324,7 +325,7 @@ def test_flat_builds_the_benchmarks_units(): ): assert s[group][unit]["members"] == ["Head", "Spouse", "Kid"] assert s["households"]["household"]["state_code"] == {"2026": "LA"} - assert s["spm_units"]["spm_unit"]["pre_subsidy_rent"] == {"2026": 9000} + assert s["spm_units"]["spm_unit"]["spm_unit_energy_subsidy"] == {"2026": 900} assert s["tax_units"]["tax_unit"]["tax_unit_childcare_expenses"] == {"2026": 1200} assert s["households"]["household"][ "qualified_passenger_vehicle_loan_interest" @@ -361,6 +362,29 @@ def test_flat_names_a_person_by_role_when_no_name_is_given(): assert any("name must be a short string" in p for p in error["problems"]) +def test_flat_placement_gives_a_household_of_ones_person_fact_to_that_person(): + one = [{"name": "Solo", "role": "head", "facts": {"age": 70}}] + args = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, error + solo = request["situation"]["people"]["Solo"] + assert solo["pre_subsidy_rent"] == {"2026": 9000} + assert "pre_subsidy_rent" not in request["situation"]["spm_units"]["spm_unit"] + assert any("moved pre_subsidy_rent" in n for n in request["notes"]) + one[0]["facts"]["pre_subsidy_rent"] = 100 + conflict = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(conflict, F.FLAT, F.PLACEMENT) + assert any("different value" in p for p in error["problems"]) + # Two people: whose rent it is is not said, so it is refused. + request, error = _norm( + _flat(household_facts={"pre_subsidy_rent": 9000}), F.FLAT, F.PLACEMENT + ) + assert any("whose facts list it" in p for p in error["problems"]) + # Without placement, even a household of one is refused. + request, error = _norm(args, F.FLAT) + assert any("whose facts list it" in p for p in error["problems"]) + + def test_flat_refuses_what_it_cannot_place(): people = [ {"name": "A", "role": "head", "facts": {}}, @@ -381,13 +405,12 @@ def test_flat_refuses_what_it_cannot_place(): def test_flat_with_names_reads_labeled_household_facts(): request, error = _norm( - _flat(household_facts={"pre subsidy rent": 9000}), F.FLAT, F.NAMES + _flat(household_facts={"spm unit energy subsidy": 900}), F.FLAT, F.NAMES ) assert error is None - assert request["situation"]["spm_units"]["spm_unit"]["pre_subsidy_rent"] == { - "2026": 9000 - } - assert any("as pre_subsidy_rent" in n for n in request["notes"]) + spm_unit = request["situation"]["spm_units"]["spm_unit"] + assert spm_unit["spm_unit_energy_subsidy"] == {"2026": 900} + assert any("as spm_unit_energy_subsidy" in n for n in request["notes"]) def test_the_schema_offers_each_features_arguments(): From 9c43f4f7ea0e576e8686af00b962a1896fe458a3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 17:43:24 -0400 Subject: [PATCH 12/16] Keep tool version 2's request as it was: detail only under concise A differential against 5bbd571a (the commit the test batches ran) found version 2 requests carrying detail "full"; only execute_request reads it, through .get, so results were unchanged, but the request now matches exactly. 4,000 generated calls through versions 1 and 2, the notices, descriptions, schemas and 15 model profiles are identical to 5bbd571a. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 3 ++- tests/test_tool_v3.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index e3e3f1b4..2712b046 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1224,7 +1224,7 @@ def fail() -> tuple[None, dict]: elif name not in explain: explain.append(name) - detail = "full" + detail = None if CONCISE in features: detail = args.get("detail") or "brief" if detail not in ("brief", "full"): @@ -1236,6 +1236,7 @@ def fail() -> tuple[None, dict]: if v2 is not None: request["explain"] = explain request["notes"] = v2["notes"] + if detail is not None: request["detail"] = detail return request, None diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index be60efdf..e0064022 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -179,7 +179,7 @@ def test_concise_takes_a_detail_and_defaults_to_brief(): {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, F.EXPLAIN, ) - assert request["detail"] == "full" + assert "detail" not in request # every step, as tool version 2 explains def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_parameters(): From 6478e8f40f01b2253245f23a8e872d4a11784642 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 18:35:53 -0400 Subject: [PATCH 13/16] Tool v3 review r1: frame, resume, labels, flat and text fixes GPT-6.1 Sol's review (REQUEST_CHANGES, 7 findings) at 9c43f4f7: 1. A model's arguments could turn on features its arm does not offer: the worker read "__tool_features__" (and, as at 5bbd571a, "__tool_version__") from the arguments dict. A pool now sends features in a tuple frame no JSON can be, the worker validates them (names, concise beside explain), and keys like "__tool_version__" in a model's arguments are refused as unknown arguments. No run record ever carried them. 2. A resume could mix treatments. Each household record now keeps its tool features and calculator schema hash; a resume compares features, notice and schema, and refuses a non-legacy record without them. 3. Labels missed what the prompt renders. The label index is now built from the prompt's own renderer (_format_input_line, then v2_pilot_text) for every promptable, labeled and reference-stated input: fallback labels, the Marketplace plan line, "state", partnership or S-corp income. Facts the reference drops (PE_DROPPED_INPUTS) are ignored with a note. A label that names two of the engine's inputs is not read. 4. Flat refuses person keys other than name, role and facts. 5. Flat stores household facts under the name written, so normalization compares them once normalized; roles set the relationship inputs and a stated one that contradicts its role is refused. 6. The flat example uses employment_income; placement replaces the wrong-entity error sentence; "full" detail is the trace, capped. 7. Diagnostics read flat people by their default names. Evidence: 210 passed, 4 frozen-engine skips. A labels-only differential builds 30 frozen households from their prompts' facts and matches the reference situation on every output, and so do all 100 frozen and all 20 test-set households in a one-off run, with or without take-up inputs. pe and pe2 are identical to 5bbd571a over 4,000 calls and 2,000 pool messages. 11 mutants of the fixes are caught. An end-to-end Flash-Lite pe3 smoke run had 0 tool errors on 3 households. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 454 +++++++++++++++++++++++++++--------- tests/test_tool_v2.py | 7 +- tests/test_tool_v3.py | 354 ++++++++++++++++++++++------ 3 files changed, 629 insertions(+), 186 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 2712b046..5f586522 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -311,8 +311,8 @@ def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: if EXPLAIN in features: explain = ( "ask it to `explain` up to three variables to see in brief how it " - "computed them, with the parameters and laws it used, or every step " - 'with `detail` set to "full"' + "computed them, with the parameters and laws it used, or the fuller " + 'trace with `detail` set to "full"' if CONCISE in features else "ask it to `explain` up to three variables to see how it computed " "them, with the parameters and laws it used" @@ -665,9 +665,9 @@ def refuse_policyengine(label: str, value) -> None: "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " "the requested variables. `household` is a JSON object: `state` (the " "two-letter state code), `people` (a list, each with `role` -- head, " - 'spouse or dependent --, `facts`, for example {"age": 40, "employment ' - 'income": 30000}, and optionally `name`, by default the role, numbered for ' - "dependents), and `household_facts` (facts about the " + 'spouse or dependent --, `facts`, for example {"age": 40, ' + '"employment_income": 30000}, and optionally `name`, by default the role, ' + "numbered for dependents), and `household_facts` (facts about the " "household as a whole: its tax unit, family, SPM unit or home). The head " "and spouse are one couple and everyone else is on their own; everyone " "shares one tax unit, family, SPM unit and household. A fact is a value " @@ -676,7 +676,13 @@ def refuse_policyengine(label: str, value) -> None: 'lists the variables to compute for `period` (default "2026"). Each ' "result gives the variable's entity, label and value for each member of " "that entity; a monthly variable computed for a year is summed over its " - "months." + "months. An unknown name or a malformed household returns an error with " + "suggestions instead of results." +) +# The version 1 description's last sentence, which placement amends. +_ENTITY_ERROR_SENTENCE = ( + "An unknown name, a variable set on the wrong entity, or a malformed " + "situation returns an error with suggestions instead of results." ) @@ -693,6 +699,12 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - description = ( FLAT_DESCRIPTION if FLAT in features else CALCULATE_DESCRIPTION ) + if PLACEMENT in features and FLAT not in features: + description = description.replace( + _ENTITY_ERROR_SENTENCE, + "An unknown name or a malformed situation returns an error " + "with suggestions instead of results.", + ) notes = [] if LABELS in features: notes.append( @@ -706,8 +718,9 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - ) if PLACEMENT in features and FLAT not in features: notes.append( - "An input set on the wrong entity is moved to the one unit it " - "belongs to." + "An input set on the wrong entity is moved to the one unit " + "holding its people; where no single unit holds them, it " + "returns an error." ) if OVERRIDES in features: notes.append( @@ -744,7 +757,8 @@ def calculate_tool_schema(version: int = 1, features: frozenset | None = None) - "description": ( 'How much an explanation shows: "brief" (the default: the ' 'value, what it came from, and the main parameters) or "full" ' - "(every step)." + f"(the trace, up to {EXPLAIN_MAX_LINES} steps and " + f"{EXPLAIN_MAX_PARAMETERS} parameters)." ), } return schema @@ -821,13 +835,9 @@ class VariableInfo: def stated_inputs() -> frozenset[str]: """Every variable a household fact the benchmark states can set.""" - from policybench.scenarios import ( - PE_INPUT_ALIASES, - PE_INPUT_RENAMES, - get_promptable_input_specs, - ) + from policybench.scenarios import PE_INPUT_ALIASES, PE_INPUT_RENAMES - names = {spec.source_name for spec in get_promptable_input_specs()} + names = {spec.source_name for spec in _promptable_specs()} names |= BASE_STATED_INPUTS names |= set(PE_INPUT_ALIASES.values()) | set(PE_INPUT_RENAMES.values()) return frozenset(names) @@ -961,6 +971,9 @@ def _normalize_inputs( if resolved is not None: v2["notes"].append(f"{where}: {note}") name = resolved + elif (field := dropped_fact(name, registry, v2["features"])) is not None: + v2["notes"].append(_dropped_note(where, name, field)) + continue if not isinstance(name, str) or name not in registry: problems.append( f"{where}: unknown variable {name!r}." @@ -1068,7 +1081,7 @@ def fail() -> tuple[None, dict]: if FLAT in features: args = dict(args) args["situation"] = flat_situation( - args.get("household"), registry, problems, v2 + args.get("household"), registry, problems, v2, period ) if args["situation"] is None: return fail() @@ -1242,17 +1255,67 @@ def fail() -> tuple[None, dict]: FLAT_ROLES = ("head", "spouse", "dependent") +FLAT_PERSON_KEYS = ("name", "role", "facts") +# The relationship inputs a role sets (as Scenario.to_pe_household's people +# carry them). +FLAT_RELATIONS = {"is_tax_unit_head": "head", "is_tax_unit_spouse": "spouse"} + + +def flat_names(people) -> list: + """The name each listed person goes by: the one given, or, for a person + given none, the role (head, spouse, dependent_1, ...: dependents are + numbered by their place among the dependents listed).""" + names, dependents = [], 0 + for person in people if isinstance(people, list) else []: + if not isinstance(person, dict): + names.append(None) + continue + role = person.get("role") + if role == "dependent": + dependents += 1 + name = person.get("name") + if name is None and role in FLAT_ROLES: + name = role if role != "dependent" else f"dependent_{dependents}" + names.append(name) + return names -def flat_situation(household, registry, problems: list[str], v2: dict | None): +def _same_value(registry, name: str, a, b, period: str) -> bool | None: + """Whether two values of an input are the same once normalized (None when + either is malformed; normalization reports it).""" + scratch: list[str] = [] + info = registry[name] + first = _normalize_value("", name, a, info, period, scratch) + second = _normalize_value("", name, b, info, period, scratch) + return None if scratch else first == second + + +def flat_situation( + household, + registry, + problems: list[str], + v2: dict | None, + period: str = str(YEAR), +): """The flat feature's household -- state, people with role and facts, and household facts -- as a situation, with the units the benchmark's own reference situations have (policybench.scenarios.Scenario.to_pe_household): the head and spouse one marital unit and everyone else alone, and one tax - unit, SPM unit, family and household holding everyone. A household fact - goes to the unit of its variable's entity. A person given no name is named - by role (head, spouse, dependent_1, ...). Returns None after recording - problems.""" + unit, SPM unit, family and household holding everyone. Roles set the + relationship inputs; a stated one that contradicts its role is refused. + A household fact goes to the unit of its variable's entity, under the + name it was written with, so normalization reads and compares it as any + other input. A person given no name is named by role (flat_names). + Returns None after recording problems.""" + features = v2["features"] if v2 is not None else frozenset() + + def engine_name(name): + if isinstance(name, str) and name in registry: + return name + if isinstance(name, str): + return read_fact_name(name, registry, features)[0] + return None + if isinstance(household, str): if len(household) > MAX_SITUATION_CHARS: problems.append("household is too large") @@ -1282,29 +1345,27 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None situation: dict[str, Any] = {"people": {}} roles: dict[str, str] = {} - dependents = 0 - for index, person in enumerate(people): + for index, (person, name) in enumerate(zip(people, flat_names(people))): if not isinstance(person, dict): problems.append( f"people[{index}] must be an object with name, role and facts" ) continue - name, role, facts = ( - person.get("name"), - person.get("role"), - person.get("facts", {}), - ) + stray = sorted(map(str, set(person) - set(FLAT_PERSON_KEYS))) + if stray: + problems.append( + f"people[{index}] takes only name, role and facts; give " + f"{', '.join(stray)} under facts" + ) + continue + role, facts = person.get("role"), person.get("facts", {}) if role not in FLAT_ROLES: problems.append( f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" ) continue - if role == "dependent": - dependents += 1 - if name is None: - name = role if role != "dependent" else f"dependent_{dependents}" - if v2 is not None: - v2["notes"].append(f"people[{index}]: named {name}") + if person.get("name") is None and v2 is not None: + v2["notes"].append(f"people[{index}]: named {name}") if not isinstance(name, str) or not name or len(name) > 64: problems.append(f"people[{index}].name must be a short string") continue @@ -1326,8 +1387,21 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): return None for name, role in roles.items(): facts = situation["people"][name] - facts.setdefault("is_tax_unit_head", role == "head") - facts.setdefault("is_tax_unit_spouse", role == "spouse") + for relation, relation_role in FLAT_RELATIONS.items(): + if relation not in registry: + continue + implied = role == relation_role + stated = [k for k in facts if engine_name(k) == relation] + if not stated: + facts[relation] = implied + elif any( + _same_value(registry, relation, facts[k], implied, period) is False + for k in stated + ): + problems.append( + f"people.{name}: {relation} contradicts the role {role}; the " + "role sets it" + ) everyone = list(situation["people"]) couple = heads + spouses groups = [couple] + [[n] for n in everyone if n not in couple] @@ -1354,12 +1428,16 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): problems.append("household_facts must be an object") return None for name, value in facts.items(): - resolved = name - if isinstance(name, str) and name not in registry and v2 is not None: - resolved, note = read_fact_name(name, registry, v2["features"]) - if resolved is not None: - v2["notes"].append(f"household_facts: {note}") - if not isinstance(resolved, str) or resolved not in registry: + resolved = engine_name(name) + if resolved is None: + field = ( + dropped_fact(name, registry, features) + if isinstance(name, str) + else None + ) + if field is not None: + v2["notes"].append(_dropped_note("household_facts", name, field)) + continue problems.append( f"household_facts: unknown variable {name!r}." + (_suggest(name, registry) if isinstance(name, str) else "") @@ -1369,26 +1447,25 @@ def flat_situation(household, registry, problems: list[str], v2: dict | None): if entity == "person": # As placement moves a person variable off a one-person unit: only # a household of one says whose fact it is. - if v2 is not None and PLACEMENT in v2["features"] and len(everyone) == 1: - facts_of = situation["people"][everyone[0]] - if resolved in facts_of and facts_of[resolved] != value: - problems.append( - f"household_facts: {resolved} is also set for " - f"{everyone[0]} with a different value" - ) - continue - facts_of[resolved] = value - v2["notes"].append( - f"household_facts: moved {resolved}, a person's fact, to " - f"{everyone[0]}" + if v2 is None or PLACEMENT not in features or len(everyone) != 1: + problems.append( + f"household_facts: {resolved} is a person's fact; give it " + "under the person whose facts list it" ) continue - problems.append( - f"household_facts: {resolved} is a person's fact; give it under " - "the person whose facts list it" + target = situation["people"][everyone[0]] + v2["notes"].append( + f"household_facts: moved {resolved}, a person's fact, to {everyone[0]}" ) + else: + target = unit_of[entity] + if name in target: + if _same_value(registry, resolved, target[name], value, period) is False: + problems.append( + f"household_facts: {resolved} is set twice with different values" + ) continue - unit_of[entity][resolved] = value + target[name] = value return None if problems else situation @@ -1427,55 +1504,123 @@ def _label_key(text: str) -> str: return re.sub(r"[^0-9a-z]", "", text.casefold()) +# The prompt's household block gives the state as "state: CT". +PROMPT_EXTRA_LABELS = {"state": "state_code"} + + +@functools.lru_cache(maxsize=1) +def _promptable_specs(): + from policybench.scenarios import get_promptable_input_specs + + return get_promptable_input_specs() + + +def _sample_value(value_type: str): + if value_type == "bool": + return True + if value_type in ("str", "Enum"): + return "x" + return 1.0 + + +def rendered_labels(field: str, value_type: str = "float") -> tuple[str, ...]: + """The labels a fact is shown under: in the board's prompt + (policybench.prompts._format_input_line, special renderers included) and + in the v2-pilot prompt (v2_pilot_text applied to that line).""" + from policybench.prompts import _format_input_line, _humanize_input_label + + line = _format_input_line(field, _sample_value(value_type)) + label = _humanize_input_label(field) + if not (line == f"- {label}" or line.startswith(f"- {label}: ")): + label = line[2:].split(": ", 1)[0] + labels = [label] + suffix = line[2 + len(label) :] + rewritten = v2_pilot_text(line) + if rewritten != line and rewritten.startswith("- ") and rewritten.endswith(suffix): + labels.append(rewritten[2 : len(rewritten) - len(suffix)]) + return tuple(labels) + + @functools.lru_cache(maxsize=1) -def prompt_labels() -> dict[str, str]: - """The labels the prompts give facts in words, keyed by _label_key, as the - engine names the facts (policybench.scenarios.PE_INPUT_RENAMES applied): - the board's labels and the v2-pilot prompt's relabelings. A label two facts - share is left out.""" +def prompt_labels() -> dict[str, frozenset[tuple[str, bool]]]: + """Every label the prompts show a fact under, keyed by _label_key, each + with the facts it may name as (engine name, dropped): the promptable + inputs, every labeled input (policybench.prompts.INPUT_LABEL_OVERRIDES), + the inputs the reference sets besides them (BASE_STATED_INPUTS), inputs + the engine has renamed and inputs the reference leaves out + (policybench.scenarios.PE_INPUT_RENAMES, PE_DROPPED_INPUTS), each in the + board's and the v2-pilot prompt's wording. Dropped facts are shown but are + not engine inputs. A label may name facts of another country; a lookup + reads only the one its engine has.""" from policybench.prompts import INPUT_LABEL_OVERRIDES - from policybench.scenarios import PE_INPUT_RENAMES - - labels: dict[str, str] = {} - shared = set() - for source, label in {**INPUT_LABEL_OVERRIDES, **v2_pilot_labels()}.items(): - key = _label_key(label) - name = PE_INPUT_RENAMES.get(source, source) - if key in labels and labels[key] != name: - shared.add(key) - labels[key] = name - for source, name in PE_INPUT_RENAMES.items(): - labels.setdefault(_label_key(source), name) - return {key: name for key, name in labels.items() if key not in shared} - - -def v2_pilot_labels() -> dict[str, str]: - """The facts the v2-pilot prompt relabels (a V2_PILOT_CONVENTIONS rewrite - of ": " to ": "), by input name.""" - relabeled = {} - for pattern, replacement in V2_PILOT_CONVENTIONS: - old = pattern.pattern - if pattern.groups or not old.endswith(": ") or not replacement.endswith(": "): - continue - relabeled[old[:-2].replace(" ", "_")] = replacement[:-2] - return relabeled + from policybench.scenarios import PE_DROPPED_INPUTS, PE_INPUT_RENAMES + + fields = {spec.output_name: spec.value_type for spec in _promptable_specs()} + dropped = {f for names in PE_DROPPED_INPUTS.values() for f in names} + for source in ( + "age", + *BASE_STATED_INPUTS, + *INPUT_LABEL_OVERRIDES, + *PE_INPUT_RENAMES, + *dropped, + ): + fields.setdefault(source, "int" if source == "age" else "float") + labels: dict[str, set[tuple[str, bool]]] = {} + for source, value_type in fields.items(): + target = (PE_INPUT_RENAMES.get(source, source), source in dropped) + keys = {_label_key(label) for label in rendered_labels(source, value_type)} + if source in PE_INPUT_RENAMES or source in dropped: + keys.add(_label_key(source)) + for key in keys: + labels.setdefault(key, set()).add(target) + for label, name in PROMPT_EXTRA_LABELS.items(): + labels.setdefault(_label_key(label), set()).add((name, False)) + return {key: frozenset(targets) for key, targets in labels.items()} def read_fact_name( name: str, registry: Mapping[str, VariableInfo], features: frozenset ) -> tuple[str | None, str | None]: - """A fact's engine variable under the labels feature (the label the prompt - gives it) and then the names feature (spaces for underscores). Returns - (engine name, note) or (None, None).""" + """A fact's engine variable under the labels feature (any label the + prompt shows it under, when exactly one fact it may name is the engine's) + and then the names feature (spaces for underscores). Returns (engine + name, note) or (None, None).""" if LABELS in features: - labeled = prompt_labels().get(_label_key(name)) - if labeled is not None and labeled in registry: - return labeled, f"read {name!r} as {labeled}" + found = [ + engine + for engine, dropped in prompt_labels().get(_label_key(name), ()) + if not dropped and engine in registry + ] + if len(found) == 1: + return found[0], f"read {name!r} as {found[0]}" if NAMES in features: return resolve_name(name, registry, output_names=False) return None, None +def dropped_fact( + name: str, registry: Mapping[str, VariableInfo], features: frozenset +) -> str | None: + """Under the labels feature, the input a fact names when the prompt shows + it but the reference leaves it out (policybench.scenarios.PE_DROPPED_INPUTS: + no longer an engine input); otherwise None.""" + if LABELS not in features: + return None + found = [ + engine + for engine, dropped in prompt_labels().get(_label_key(name), ()) + if dropped + ] + return found[0] if len(found) == 1 else None + + +def _dropped_note(where: str, name: str, field: str) -> str: + return ( + f"{where}: {name!r} ({field}) is not an input of this engine; the " + "benchmark's reference leaves it out too, so it was ignored" + ) + + def _set_aliases(v2: dict, people: dict, registry: Mapping[str, VariableInfo]): """Set each fact the reference also sets under another name there too (policybench.scenarios.PE_INPUT_ALIASES: the SNAP work rules read stated @@ -1643,7 +1788,7 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic def brief_explanation(full: dict) -> dict: """The concise feature's explanation: the value, the variables it came - from directly (one level down, zero-valued ones left out), and the first + from directly (one level down, the first BRIEF_INPUTS), and the first parameters it read with their sources.""" if "error" in full: return full @@ -1655,7 +1800,7 @@ def brief_explanation(full: dict) -> dict: "value": steps[0].strip() if steps else "", "from": direct[:BRIEF_INPUTS], "parameters": full["parameters"][:BRIEF_PARAMETERS], - "more": 'set detail to "full" for every step', + "more": 'set detail to "full" for the trace', } @@ -1843,6 +1988,46 @@ def build_tool_system( return CountryTaxBenefitSystem(reform=load_fix_reform(path, snapshot)) +# How a pool tells its worker the calculator's features: a tuple, which a +# model's JSON arguments can never be, so a model cannot turn on a feature its +# arm does not offer. A pool without features sends the arguments as they are. +CALL_FRAME = "calculate" + + +def tool_call_frame(features: frozenset, arguments) -> tuple: + return (CALL_FRAME, tuple(sorted(features)), arguments) + + +def frame_features(names) -> frozenset: + """A frame's features, refused (ValueError) unless they are a tuple of + known feature names with concise only beside explain.""" + if not isinstance(names, tuple) or not all(isinstance(n, str) for n in names): + raise ValueError("the features must be a tuple of names") + features = frozenset(names) + if features - set(TOOL_FEATURES): + raise ValueError( + f"unknown tool features {sorted(features - set(TOOL_FEATURES))}" + ) + if CONCISE in features and EXPLAIN not in features: + raise ValueError("the concise feature needs explain") + return features + + +def worker_request(message, registry) -> tuple[dict | None, dict | None]: + """A worker's reading of one pool message: a call frame's arguments under + its features, or anything else as a model's arguments with no features + (keys such as "__tool_features__" or "__tool_version__" in them are the + model's, and refused as unknown arguments).""" + features = frozenset() + if isinstance(message, tuple) and len(message) == 3 and message[0] == CALL_FRAME: + try: + features = frame_features(message[1]) + except ValueError as error: + return None, {"error": "configuration", "message": str(error)} + message = message[2] + return normalize_request(message, registry, features=features) + + def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> None: for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"): os.environ[name] = "1" @@ -1857,13 +2042,7 @@ def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> return if args is None: return - features = frozenset() - if isinstance(args, dict) and "__tool_features__" in args: - features = frozenset(args.get("__tool_features__") or ()) - args = args.get("arguments") - elif isinstance(args, dict) and args.get("__tool_version__") == 2: - features, args = V2_FEATURES, args.get("arguments") - request, error = normalize_request(args, registry, features=features) + request, error = worker_request(args, registry) if error is not None: conn.send(("ok", error)) continue @@ -1958,13 +2137,8 @@ def call(self, args) -> tuple[dict, float]: time.perf_counter() - started, ) process, conn = worker - if self._tool_features == V2_FEATURES: - args = {"__tool_version__": 2, "arguments": args} - elif self._tool_features: - args = { - "__tool_features__": sorted(self._tool_features), - "arguments": args, - } + if self._tool_features: + args = tool_call_frame(self._tool_features, args) try: conn.send(args) if conn.poll(self.timeout): @@ -2863,6 +3037,7 @@ def run_household( "benchmark_search_rule": profile.benchmark_search_rule, "sent_board_prompt_sha256": sent_board_prompt_sha256, "tool_notice_sha256": hashlib.sha256(profile.notice.encode()).hexdigest(), + **profile_treatment(profile), "turns": [], "status": "running", } @@ -3571,6 +3746,52 @@ def records_wording(records: Mapping[str, dict]) -> str | None: return hashes.pop() if hashes else None +# Arms whose calculator features never changed, so a record written before +# households recorded their features has these. +LEGACY_TOOL_ARMS = {"pe": frozenset(), "web+pe": frozenset(), "pe2": V2_FEATURES} + + +def profile_treatment(profile: "ModelProfile") -> dict: + """What a household record keeps of the treatment it ran under, beyond + the notice hash: the calculator's features and its schema's hash.""" + if not profile.uses_calculator: + return {"tool_features": None, "calculate_tool_sha256": None} + schema = calculate_tool_schema(features=profile.tool_features) + return { + "tool_features": sorted(profile.tool_features), + "calculate_tool_sha256": hashlib.sha256( + json.dumps(schema, sort_keys=True).encode() + ).hexdigest(), + } + + +def record_treatment_problem(record: dict, profile: "ModelProfile") -> str | None: + """Why a household record ran under another treatment than the profile's + (features, notice or calculator schema), or None. A record written before + a field existed counts as matching only for an arm whose treatment never + changed (none, web, and LEGACY_TOOL_ARMS); otherwise it is unrecorded.""" + sid = record.get("scenario_id") + want = profile_treatment(profile) + legacy = profile.arm in ("none", "web") or profile.arm in LEGACY_TOOL_ARMS + if "tool_features" in record: + if record["tool_features"] != want["tool_features"]: + return f"{sid} ran with tool features {record['tool_features']}" + elif not legacy: + return f"{sid} has no recorded tool features" + notice = record.get("tool_notice_sha256") + if notice is None and not legacy: + return f"{sid} has no recorded tool notice" + if ( + notice is not None + and notice != hashlib.sha256(profile.notice.encode()).hexdigest() + ): + return f"{sid} ran under another tool notice" + schema = record.get("calculate_tool_sha256") + if schema is not None and schema != want["calculate_tool_sha256"]: + return f"{sid} ran with another calculator schema" + return None + + def records_arm(records: Mapping[str, dict]) -> str: """The one arm a run directory's records belong to; records written before the web arms are the PE tool's.""" @@ -3629,6 +3850,17 @@ def command_run(args) -> None: f"{households_dir} holds records sent under other v2-pilot wording " f"({records_wording(previous)}); use a fresh --out-dir" ) + mixed = [ + problem + for record in previous.values() + if (problem := record_treatment_problem(record, profile)) is not None + ] + if mixed: + raise SystemExit( + f"{households_dir} holds records from another treatment (" + + "; ".join(mixed[:3]) + + "); use a fresh --out-dir" + ) spend = Spend(args.budget_usd, spent_before) todo = [sid for sid in ids if sid not in done] manifest = { @@ -4416,7 +4648,7 @@ def _person_names(arguments) -> list[str]: household = ( json.loads(household) if isinstance(household, str) else household ) - return [person["name"] for person in household["people"]] + return [name or "" for name in flat_names(household["people"])] situation = args["situation"] situation = json.loads(situation) if isinstance(situation, str) else situation return list(situation["people"]) diff --git a/tests/test_tool_v2.py b/tests/test_tool_v2.py index d6e381b5..68248090 100644 --- a/tests/test_tool_v2.py +++ b/tests/test_tool_v2.py @@ -261,8 +261,13 @@ def test_the_pool_tells_its_workers_the_tool_version(): ) try: result, _ = pool.call({"situation": "{}"}) + # A frame no model's JSON arguments can be (tool v3 review r1). assert result == { - "received": {"__tool_version__": 2, "arguments": {"situation": "{}"}} + "received": ( + "calculate", + ("explain", "names", "overrides", "placement"), + {"situation": "{}"}, + ) } finally: pool.close() diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index e0064022..4b1e9c99 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -179,7 +179,7 @@ def test_concise_takes_a_detail_and_defaults_to_brief(): {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, F.EXPLAIN, ) - assert "detail" not in request # every step, as tool version 2 explains + assert "detail" not in request # the trace, as tool version 2 explains def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_parameters(): @@ -268,24 +268,75 @@ def test_only_pe2_says_the_facts_carry_variable_names(): assert ("names it read" in notice) == bool(features & {F.NAMES, F.LABELS}) -def test_every_prompt_label_reads_back_to_its_engine_name(): - from policybench.prompts import INPUT_LABEL_OVERRIDES - from policybench.scenarios import PE_INPUT_RENAMES +def test_every_label_a_prompt_shows_reads_back_to_its_engine_name(): + from policybench.prompts import _format_input_line + from policybench.scenarios import PE_DROPPED_INPUTS, PE_INPUT_RENAMES labels = F.prompt_labels() - for field, label in INPUT_LABEL_OVERRIDES.items(): - assert labels[F._label_key(label)] == PE_INPUT_RENAMES.get(field, field) - # The v2-pilot prompt's own relabelings, as that prompt writes them. - relabeled = F.v2_pilot_labels() - assert set(relabeled) == { - "employer_sponsored_insurance_premiums", - "spm_unit_energy_subsidy", + fields = {spec.output_name: spec.value_type for spec in F._promptable_specs()} + fields |= {"age": "int", "employment_income": "float"} + dropped = {f for names in PE_DROPPED_INPUTS.values() for f in names} + for field, value_type in {**fields, **dict.fromkeys(dropped, "float")}.items(): + for label in F.rendered_labels(field, value_type): + want = (PE_INPUT_RENAMES.get(field, field), field in dropped) + assert want in labels[F._label_key(label)], (field, label) + # Fallback labels (the name with spaces) are the prompt's too. + assert ("first_home_mortgage_balance", False) in labels[ + F._label_key("first home mortgage balance") + ] + # Special renderers: the Marketplace plan line is labeled on its own. + line = _format_input_line("selected_marketplace_plan_benchmark_ratio", 0.8) + assert line.startswith("- selected Marketplace plan: ") + assert ("selected_marketplace_plan_benchmark_ratio", False) in labels[ + F._label_key("selected Marketplace plan") + ] + # The v2-pilot prompt's relabelings, as that prompt writes them. + for field, label in ( + ("spm_unit_energy_subsidy", "home energy assistance received"), + ( + "employer_sponsored_insurance_premiums", + "employer-paid health insurance premiums (paid by the employer, " + "not by the household)", + ), + ): + assert label in F.rendered_labels(field) + assert (field, False) in labels[F._label_key(label)] + assert ("state_code", False) in labels[F._label_key("state")] + assert ("first_home_mortgage_interest", True) in labels[ + F._label_key("first home mortgage interest") + ] + assert ("roth_401k_contributions_desired", False) in labels[ + F._label_key("Roth 401k contributions desired") + ] + assert ("partnership_s_corp_income", False) in labels[ + F._label_key("partnership or S-corp income") + ] + + +def test_labels_reads_no_label_that_names_two_of_the_engines_inputs(monkeypatch): + two = frozenset({("has_esi", False), ("age", False)}) + monkeypatch.setattr(F, "prompt_labels", lambda: {F._label_key("covered"): two}) + situation = _situation() + situation["people"]["head"]["covered"] = True + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert request is None and "unknown variable 'covered'" in str(error) + one = frozenset({("has_esi", False), ("uk_only_input", False)}) + monkeypatch.setattr(F, "prompt_labels", lambda: {F._label_key("covered"): one}) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None and request["situation"]["people"]["head"]["has_esi"] + + +def test_labels_ignores_a_fact_the_reference_leaves_out(): + situation = _situation() + situation["tax_units"] = { + "t": {"members": ["head", "kid"], "first home mortgage interest": 900} } - for field, label in relabeled.items(): - written = F.v2_pilot_text(f"- {field.replace('_', ' ')}: $5") - assert written == f"- {label}: $5" - assert labels[F._label_key(label)] == field - assert labels[F._label_key("roth 401k contributions")] == "roth_401k_contributions" + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None, error + assert request["situation"]["tax_units"]["t"] == {"members": ["head", "kid"]} + assert any("leaves it out too" in n for n in request["notes"]) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.NAMES) + assert request is None and "first home mortgage interest" in str(error) def _flat(**overrides): @@ -385,6 +436,74 @@ def test_flat_placement_gives_a_household_of_ones_person_fact_to_that_person(): assert any("whose facts list it" in p for p in error["problems"]) +def test_flat_refuses_facts_given_beside_the_person(): + people = [{"role": "head", "age": 40, "employment_income": 30000}] + request, error = _norm(_flat(people=people), F.FLAT) + assert request is None + assert any( + "takes only name, role and facts; give age, employment_income under facts" in p + for p in error["problems"] + ) + + +def test_flat_reads_a_fact_written_two_ways_once_whatever_the_order(): + for facts in ( + {"spm_unit_energy_subsidy": 900, "spm unit energy subsidy": 800}, + {"spm unit energy subsidy": 800, "spm_unit_energy_subsidy": 900}, + ): + request, error = _norm(_flat(household_facts=facts), F.FLAT, F.NAMES) + assert request is None + assert any("set twice with different values" in p for p in error["problems"]) + same = {"spm_unit_energy_subsidy": 900, "spm unit energy subsidy": {"2026": 900}} + request, error = _norm(_flat(household_facts=same), F.FLAT, F.NAMES) + assert error is None, error + + +def test_flat_roles_set_the_relationship_inputs_and_refuse_a_contradiction(): + for written in ("is_tax_unit_head", "is tax unit head"): + people = [{"name": "A", "role": "head", "facts": {written: False}}] + request, error = _norm(_flat(people=people), F.FLAT, F.NAMES) + assert request is None, written + assert any("contradicts the role head" in p for p in error["problems"]) + people = [{"name": "A", "role": "head", "facts": {written: True}}] + request, error = _norm(_flat(people=people), F.FLAT, F.NAMES) + assert error is None, (written, error) + head = request["situation"]["people"]["A"] + assert head["is_tax_unit_head"] == {"2026": True} + assert head["is_tax_unit_spouse"] == {"2026": False} + + +def test_flat_placement_compares_values_once_normalized(): + one = [ + {"name": "Solo", "role": "head", "facts": {"pre_subsidy_rent": {"2026": 9000}}} + ] + args = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, error + assert request["situation"]["people"]["Solo"]["pre_subsidy_rent"] == {"2026": 9000} + + +def test_the_descriptions_promise_only_what_the_features_do(): + flat = F.calculate_tool_schema(features=frozenset({F.FLAT}))["function"] + assert '"employment_income": 30000' in flat["description"] + example = [ + {"role": "head", "facts": {"age": 40, "employment_income": 30000}}, + ] + request, error = _norm(_flat(people=example), F.FLAT) + assert error is None, error + entity_error = "a variable set on the wrong entity" + placed = F.calculate_tool_schema(features=frozenset({F.PLACEMENT}))["function"] + assert entity_error not in placed["description"] + assert "moved to the one unit holding its people" in placed["description"] + named = F.calculate_tool_schema(features=frozenset({F.NAMES}))["function"] + assert entity_error in named["description"] + concise = F.calculate_tool_schema(features=frozenset({F.EXPLAIN, F.CONCISE})) + detail = concise["function"]["parameters"]["properties"]["detail"]["description"] + assert f"up to {F.EXPLAIN_MAX_LINES} steps" in detail + for features in (frozenset({F.EXPLAIN, F.CONCISE}), frozenset(F.TOOL_FEATURES)): + assert "every step" not in F.compose_tool_notice(features, True) + + def test_flat_refuses_what_it_cannot_place(): people = [ {"name": "A", "role": "head", "facts": {}}, @@ -434,7 +553,7 @@ def _echo_worker(conn, fix_path, reform): conn.send(("ok", {"received": args})) -def test_the_pool_tells_its_workers_the_features(): +def test_the_pool_tells_its_workers_the_features_in_a_frame(): pool = F.ToolPool( size=1, timeout=5.0, @@ -443,13 +562,91 @@ def test_the_pool_tells_its_workers_the_features(): ) try: result, _ = pool.call({"x": 1}) - assert result == { - "received": {"__tool_features__": ["flat", "names"], "arguments": {"x": 1}} - } + assert result == {"received": ("calculate", ("flat", "names"), {"x": 1})} finally: pool.close() +def test_a_model_cannot_turn_on_a_feature_its_arm_does_not_offer(): + forged = { + "__tool_features__": ["names"], + "arguments": { + "situation": {"people": {"head": {"Employment income": 12345}}}, + "variables": ["age"], + }, + } + for message in (forged, {"__tool_version__": 2, "arguments": forged["arguments"]}): + request, error = F.worker_request(message, REGISTRY) + assert request is None + assert "unknown arguments: " in error["problems"][0] + request, error = F.worker_request( + F.tool_call_frame(frozenset({F.NAMES}), forged["arguments"]), REGISTRY + ) + assert error is None + assert request["situation"]["people"]["head"]["employment_income"] == { + "2026": 12345 + } + for names in (("telepathy",), ("concise",), ({},), ["names"]): + request, error = F.worker_request(("calculate", names, {}), REGISTRY) + assert request is None and error["error"] == "configuration", names + + +def _profile(arm): + return F.ModelProfile("m", "gpt-x", "chat", "tool", arm=arm) + + +def _record(profile, **changes): + record = { + "scenario_id": "mt_000", + "tool_notice_sha256": __import__("hashlib") + .sha256(profile.notice.encode()) + .hexdigest(), + **F.profile_treatment(profile), + } + record.update(changes) + return record + + +def test_a_resume_refuses_households_run_under_another_treatment(): + pe3 = _profile("pe3") + assert F.record_treatment_problem(_record(pe3), pe3) is None + six = sorted(set(F.TOOL_FEATURES) - {F.LABELS}) + assert "tool features" in F.record_treatment_problem( + _record(pe3, tool_features=six), pe3 + ) + assert "another tool notice" in F.record_treatment_problem( + _record(pe3, tool_notice_sha256="0" * 64), pe3 + ) + assert "calculator schema" in F.record_treatment_problem( + _record(pe3, calculate_tool_sha256="0" * 64), pe3 + ) + old = _record(pe3) + del old["tool_features"] + assert "no recorded tool features" in F.record_treatment_problem(old, pe3) + # Arms whose treatment never changed accept records from before the fields. + for arm in ("pe", "pe2", "web+pe", "none", "web"): + profile = _profile(arm) + legacy = _record(profile) + for key in ("tool_features", "calculate_tool_sha256", "tool_notice_sha256"): + legacy.pop(key) + assert F.record_treatment_problem(legacy, profile) is None, arm + pe2 = _profile("pe2") + assert "tool features" in F.record_treatment_problem(_record(_profile("pe")), pe2) + + +def test_diagnostics_read_flat_people_by_their_default_names(): + household = { + "people": [ + {"role": "head", "facts": {}}, + {"role": "spouse", "facts": {}}, + {"name": "Kid", "role": "dependent", "facts": {}}, + {"role": "dependent", "facts": {}}, + ] + } + names = F._person_names({"household": json.dumps(household)}) + assert names == ["head", "spouse", "Kid", "dependent_2"] + + @pytest.mark.slow def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): """Differential: on the installed engine, each frozen-board household @@ -520,73 +717,72 @@ def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): @pytest.mark.slow -def test_a_household_described_under_the_prompts_labels_computes_the_same(): - """Differential: each frozen-board household described flat with every - fact as its v2-pilot prompt writes it (a label in words, or the name with - spaces), - and stated hours only where the prompt states them, computes every board - output as the benchmark does. Take-up conventions, which the prompt does - not list, keep their engine names.""" +def test_a_household_described_from_its_prompt_computes_the_same(): + """Differential, with the labels feature alone: each of 30 frozen-board + households described flat from the facts its v2-pilot prompt shows, each + under the label the prompt shows it under (dropped facts included), with + the state and roles the prompt gives, computes every board output as the + benchmark's own situation does. The take-up conventions, which the prompt + states in words, are set under their engine names.""" import household_set as hs import pandas as pd from policyengine_us import Simulation from policybench.ground_truth import _extract_person_value, _pe_variable_for_output - from policybench.prompts import _humanize_input_label - from policybench.scenarios import PE_INPUT_ALIASES, PE_INPUT_RENAMES + from policybench.scenarios import ( + DEFAULT_TAKEUP_INPUTS, + is_excluded_prompt_input_name, + ) system = F.build_tool_system(reform=True) registry = F.build_registry(system) reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") variables = reference.groupby("scenario_id").variable.agg(list) - field_of = {name: field for field, name in PE_INPUT_RENAMES.items()} - dropped = {"is_tax_unit_head", "is_tax_unit_spouse", *PE_INPUT_ALIASES.values()} + shown = set() - worded = [] - - def labeled(facts): - out = {} - for k, v in facts.items(): - if k in dropped: - continue - if k.startswith(F.CONVENTION_PREFIXES): - out[k] = v - continue - label = F.v2_pilot_text(f"- {_humanize_input_label(field_of.get(k, k))}: ")[ - 2:-2 - ] - if label != k.replace("_", " "): - worded.append(label) - out[label] = v - return out + def label(field, value): + value_type = ( + "bool" + if isinstance(value, bool) + else ("str" if isinstance(value, str) else "float") + ) + written = F.rendered_labels(field, value_type)[-1] + shown.add((field, written)) + return written + + def prompt_facts(inputs): + return { + label(k, v): v + for k, v in inputs.items() + if not is_excluded_prompt_input_name(k) + } checked = 0 - for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:12]: - situation = hs.situation_for(scenario) + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:30]: couple = scenario.marital_couple() or () head = couple[0] if couple else scenario.adults[0].name - people = [ - { - "name": name, - "role": "head" - if name == head - else ("spouse" if name in couple else "dependent"), - "facts": labeled(facts), - } - for name, facts in situation["people"].items() - ] + people = [] + for person in scenario.all_people: + facts = {"age": person.age} + if person.name.startswith("adult") or abs(person.employment_income) > 1e-6: + facts[label("employment_income", 1.0)] = person.employment_income + facts |= prompt_facts(person.inputs) + facts |= DEFAULT_TAKEUP_INPUTS["person"] + role = ( + "head" + if person.name == head + else ("spouse" if person.name in couple else "dependent") + ) + people.append({"name": person.name, "role": role, "facts": facts}) household_facts = {} - for group in ("tax_units", "spm_units", "families", "households"): - for unit in situation[group].values(): - household_facts.update( - labeled( - { - k: v - for k, v in unit.items() - if k not in ("members", "state_code") - } - ) - ) + for inputs in ( + scenario.tax_unit_inputs, + scenario.spm_unit_inputs, + scenario.household_inputs, + ): + household_facts |= prompt_facts(inputs) + for entity in ("tax_unit", "spm_unit", "household"): + household_facts |= DEFAULT_TAKEUP_INPUTS[entity] household = { "state": scenario.state, "people": people, @@ -595,10 +791,12 @@ def labeled(facts): request, error = F.normalize_request( {"household": json.dumps(household), "variables": ["snap"]}, registry, - features=frozenset({F.FLAT, F.LABELS, F.NAMES}), + features=frozenset({F.FLAT, F.LABELS}), ) assert error is None, (scenario.id, error) - own = Simulation(tax_benefit_system=system, situation=situation) + own = Simulation( + tax_benefit_system=system, situation=hs.situation_for(scenario) + ) flat = Simulation(tax_benefit_system=system, situation=request["situation"]) for variable in variables[scenario.id]: name = _pe_variable_for_output(variable, "us") @@ -612,4 +810,12 @@ def labeled(facts): ] assert values[0] == values[1], (scenario.id, variable, values) checked += 1 - assert checked > 120 and len(worded) > 10, (checked, worded) + worded = {w for f, w in shown if w != f.replace("_", " ")} + assert checked > 300 and len(worded) > 10, (checked, sorted(worded)) + # No label a US prompt can show names two of this engine's inputs. + ambiguous = { + key: targets + for key, targets in F.prompt_labels().items() + if len({e for e, dropped in targets if not dropped and e in registry}) > 1 + } + assert not ambiguous, ambiguous From d570b8e7af46a99f0dd4110be2ec8ea9147de835 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 20:09:47 -0400 Subject: [PATCH 14/16] Tool v3 review r2: report malformed repeats, resumable harness errors - flat: a household fact repeated with a malformed value was dropped without a word. _same_value now reports the repeat's validation problems, and only a repeat validated as equal is dropped quietly. A malformed first value is reported once, by normal validation, and never as a conflict. - The fallback record for a household the harness failed on carried no treatment, so one failure blocked resuming any custom arm. The record is built by harness_error_record, which now includes the notice hash and profile_treatment. resume_treatment_problems checks only the records a resume keeps, since a rerunnable record is replaced under the current profile. pe, web+pe and pe2 still behave exactly as at 5bbd571a: 4,000 calls and 2,000 pool messages match. 5 mutants of these fixes were all caught. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 102 +++++++++++++++++++++++++----------- tests/test_tool_v3.py | 54 +++++++++++++++++++ 2 files changed, 126 insertions(+), 30 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 5f586522..b7b66b72 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1280,14 +1280,29 @@ def flat_names(people) -> list: return names -def _same_value(registry, name: str, a, b, period: str) -> bool | None: - """Whether two values of an input are the same once normalized (None when - either is malformed; normalization reports it).""" - scratch: list[str] = [] +def _same_value( + registry, + name: str, + a, + b, + period: str, + problems: list[str] | None = None, + where: str = "", +) -> bool | None: + """Whether two values of an input are the same once normalized, or None + when either is malformed. ``a`` stays in the situation, so normalization + reports it later; a caller that drops ``b`` passes ``problems`` so a + malformed ``b`` is reported, not lost.""" info = registry[name] - first = _normalize_value("", name, a, info, period, scratch) - second = _normalize_value("", name, b, info, period, scratch) - return None if scratch else first == second + first_problems: list[str] = [] + second_problems: list[str] = [] + first = _normalize_value(where, name, a, info, period, first_problems) + second = _normalize_value(where, name, b, info, period, second_problems) + if problems is not None: + problems.extend(second_problems) + if first_problems or second_problems: + return None + return first == second def flat_situation( @@ -1460,7 +1475,12 @@ def engine_name(name): else: target = unit_of[entity] if name in target: - if _same_value(registry, resolved, target[name], value, period) is False: + # The repeat is dropped, so only one validated as equal is kept quiet. + same = _same_value( + registry, resolved, target[name], value, period, problems, + "household_facts", + ) + if same is False: problems.append( f"household_facts: {resolved} is set twice with different values" ) @@ -3792,6 +3812,46 @@ def record_treatment_problem(record: dict, profile: "ModelProfile") -> str | Non return None +def resume_treatment_problems( + previous: Mapping[str, dict], done: set[str], profile: "ModelProfile" +) -> list[str]: + """record_treatment_problem for each record a resume keeps. A rerunnable + record (RERUNNABLE_STATUSES) is replaced under this profile, so it cannot + mix treatments and does not block the resume.""" + return [ + problem + for sid, record in previous.items() + if sid in done + and (problem := record_treatment_problem(record, profile)) is not None + ] + + +def harness_error_record( + sid: str, variables: list[str], profile: "ModelProfile", error: Exception +) -> dict: + """The record of a household the harness itself failed on: rerunnable, + with every setup and treatment key a resume and score check, so a folder + holding it still resumes and scores.""" + record = { + "scenario_id": sid, + "model": profile.name, + "arm": profile.arm, + "prompt_version": profile.prompt_version, + "prompt_wording_sha256": prompt_wording_sha256(profile), + "benchmark_search_rule": profile.benchmark_search_rule, + "tool_notice_sha256": hashlib.sha256(profile.notice.encode()).hexdigest(), + **profile_treatment(profile), + "variables": variables, + "status": "harness_error", + "error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS], + "turns": [], + "predictions": {v: None for v in variables}, + "explanations": {v: None for v in variables}, + } + record["totals"] = household_totals(record) + return record + + def records_arm(records: Mapping[str, dict]) -> str: """The one arm a run directory's records belong to; records written before the web arms are the PE tool's.""" @@ -3850,11 +3910,7 @@ def command_run(args) -> None: f"{households_dir} holds records sent under other v2-pilot wording " f"({records_wording(previous)}); use a fresh --out-dir" ) - mixed = [ - problem - for record in previous.values() - if (problem := record_treatment_problem(record, profile)) is not None - ] + mixed = resume_treatment_problems(previous, done, profile) if mixed: raise SystemExit( f"{households_dir} holds records from another treatment (" @@ -3959,23 +4015,9 @@ def one(sid: str) -> dict: profile=profile, ) except Exception as error: # one household must not end the run - record = { - "scenario_id": sid, - "model": profile.name, - # The setup keys a resume and score check, so a folder with - # this record still resumes and scores. - "arm": profile.arm, - "prompt_version": profile.prompt_version, - "prompt_wording_sha256": prompt_wording_sha256(profile), - "benchmark_search_rule": profile.benchmark_search_rule, - "variables": board["variables"][sid], - "status": "harness_error", - "error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS], - "turns": [], - "predictions": {v: None for v in board["variables"][sid]}, - "explanations": {v: None for v in board["variables"][sid]}, - } - record["totals"] = household_totals(record) + record = harness_error_record( + sid, board["variables"][sid], profile, error + ) if sid in previous: prior = previous[sid] record["prior_attempts"] = prior.get("prior_attempts", []) + [ diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index 4b1e9c99..dcf6317a 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -819,3 +819,57 @@ def prompt_facts(inputs): if len({e for e, dropped in targets if not dropped and e in registry}) > 1 } assert not ambiguous, ambiguous + + +def test_flat_reports_a_malformed_repeat_of_a_fact_it_drops(): + # Review r2: a repeat that fails validation was dropped without a word. + one = [{"name": "Solo", "role": "head", "facts": {"employment_income": 30000}}] + for bad in ("40000", [40000], {"26": 40000}, {}): + args = _flat(people=one, household_facts={"employment_income": bad}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None, bad + assert any("employment_income" in m for m in error["problems"]), (bad, error) + # An equal repeat, scalar or by period, still merges quietly; a different + # one is still refused. + for same in (30000, {"2026": 30000}): + args = _flat(people=one, household_facts={"employment_income": same}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, (same, error) + args = _flat(people=one, household_facts={"employment_income": 40000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None + assert any("set twice with different values" in m for m in error["problems"]) + + +def test_a_harness_error_record_resumes_under_its_own_arm(): + # Review r2: the fallback record carried no treatment, so one harness + # failure blocked every custom-arm resume. + for arm in ("pe3", "pe:labels+flat", "pe", "pe2", "web+pe", "none", "web"): + profile = _profile(arm) + record = F.harness_error_record("mt_000", ["snap"], profile, RuntimeError("x")) + assert record["status"] in F.RERUNNABLE_STATUSES + assert F.record_treatment_problem(record, profile) is None, arm + assert record["predictions"] == {"snap": None} + assert record["totals"]["cost_usd"] == 0 + + +def test_a_resume_checks_the_treatment_only_of_records_it_keeps(): + pe3 = _profile("pe3") + stale = _record(pe3, tool_features=["names"]) + previous = { + "mt_000": dict(stale, status="complete"), + "mt_001": dict(stale, scenario_id="mt_001", status="harness_error"), + } + # Kept, so refused. + assert len(F.resume_treatment_problems(previous, {"mt_000"}, pe3)) == 1 + # Rerun under this profile, so no mixing and no block. + assert F.resume_treatment_problems(previous, set(), pe3) == [] + + +def test_flat_reports_a_malformed_first_value_once_and_not_as_a_conflict(): + one = [{"name": "Solo", "role": "head", "facts": {"employment_income": "30000"}}] + args = _flat(people=one, household_facts={"employment_income": 40000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None + assert not any("set twice" in m for m in error["problems"]), error + assert sum("employment_income" in m for m in error["problems"]) == 1, error From fb35d8273b0d1f6fd5df1aea76b6e22d15af0e8f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 20:45:58 -0400 Subject: [PATCH 15/16] Format the r2 fix Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index b7b66b72..abab108f 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1477,7 +1477,12 @@ def engine_name(name): if name in target: # The repeat is dropped, so only one validated as equal is kept quiet. same = _same_value( - registry, resolved, target[name], value, period, problems, + registry, + resolved, + target[name], + value, + period, + problems, "household_facts", ) if same is False: From 6b00a6f0c4997e73b9968b412b57646957cad631 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 01:29:28 -0400 Subject: [PATCH 16/16] Tool v3 review r3: check every record's treatment, at resume and at score A rerun is not sure to replace a rerunnable record: a budget stop, lost workers or a narrower --scenarios leave it in place, and scoring read whatever was there. treatment_problems now checks every record in the folder, rerunnable ones included, and command_score refuses a folder that mixes treatments, rebuilding the notices on the board's engine. Runs written before 6e1578fc (test batch 0's frontier calculator arms, b6d835a8) were told policyengine-us 2.41.0 while computing on 2.60.4. A folder written wholly so is one treatment and scores with a note; a mix of the two texts is refused, and a resume into it stays refused, since its new records would state the board's engine. All 39 run folders score (628 records); the six frontier calculator folders take the note. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 63 ++++++++++++--- tests/test_tool_v3.py | 154 ++++++++++++++++++++++++++++++++++-- 2 files changed, 201 insertions(+), 16 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index abab108f..f09e0b6f 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -48,6 +48,7 @@ import argparse import concurrent.futures +import contextlib import copy import decimal import difflib @@ -343,6 +344,18 @@ def set_stated_engine(version: str) -> None: _STATED_ENGINE["version"] = version +@contextlib.contextmanager +def stating_engine(version: str): + """The texts as a run on `version` sent them, for a check that rebuilds a + notice; the stated engine is restored afterwards.""" + before = _STATED_ENGINE["version"] + set_stated_engine(version) + try: + yield + finally: + set_stated_engine(before) + + def engine_text(text: str) -> str: return text.replace( f"policyengine-us {REFERENCE_ENGINE}", @@ -3817,20 +3830,43 @@ def record_treatment_problem(record: dict, profile: "ModelProfile") -> str | Non return None -def resume_treatment_problems( - previous: Mapping[str, dict], done: set[str], profile: "ModelProfile" +def treatment_problems( + records: Mapping[str, dict], profile: "ModelProfile" ) -> list[str]: - """record_treatment_problem for each record a resume keeps. A rerunnable - record (RERUNNABLE_STATUSES) is replaced under this profile, so it cannot - mix treatments and does not block the resume.""" + """record_treatment_problem for every record in a run directory, a + rerunnable one (RERUNNABLE_STATUSES) too. A rerun is not sure to replace + it: a budget stop, lost workers or a narrower --scenarios leave it in + place, and scoring reads whatever is there (review r3). So a resume and + a score both refuse a directory holding any record from another treatment.""" return [ problem - for sid, record in previous.items() - if sid in done - and (problem := record_treatment_problem(record, profile)) is not None + for record in records.values() + if (problem := record_treatment_problem(record, profile)) is not None ] +def scored_treatment_problems( + records: Mapping[str, dict], profile: "ModelProfile", engine: str +) -> tuple[list[str], str | None]: + """treatment_problems as score applies it, with the note score prints. + The records' notices state the board's engine, as command_run set it. + Runs before the texts stated the board's engine (6e1578fc) named the + reference engine whatever the board's: a directory written wholly so is + one treatment and scores as it ran. A mix of the two texts does not. A + resume stays strict, since its new records would state the board's.""" + with stating_engine(engine): + problems = treatment_problems(records, profile) + if not problems or engine == REFERENCE_ENGINE: + return problems, None + with stating_engine(REFERENCE_ENGINE): + if treatment_problems(records, profile): + return problems, None + return [], ( + f"NOTE: these records' notices name policyengine-us {REFERENCE_ENGINE}, " + f"as runs before 6e1578fc stated it; the board's engine is {engine}" + ) + + def harness_error_record( sid: str, variables: list[str], profile: "ModelProfile", error: Exception ) -> dict: @@ -3915,7 +3951,7 @@ def command_run(args) -> None: f"{households_dir} holds records sent under other v2-pilot wording " f"({records_wording(previous)}); use a fresh --out-dir" ) - mixed = resume_treatment_problems(previous, done, profile) + mixed = treatment_problems(previous, profile) if mixed: raise SystemExit( f"{households_dir} holds records from another treatment (" @@ -4148,6 +4184,15 @@ def command_score(args) -> None: profile = model_profile( records_model(records), records_arm(records), *records_prompt(records) ) + mixed, note = scored_treatment_problems(records, profile, board_engine(snapshot)) + if note: + print(note) + if mixed: + raise SystemExit( + f"{out / 'households'} holds records from another treatment (" + + "; ".join(mixed[:3]) + + "); scoring would mix them" + ) arm1_label, arm2_label = profile.arm1_label, profile.arm2_label complete = sorted( sid for sid, r in records.items() if r.get("status") == "complete" diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py index dcf6317a..71bc73c3 100644 --- a/tests/test_tool_v3.py +++ b/tests/test_tool_v3.py @@ -853,17 +853,157 @@ def test_a_harness_error_record_resumes_under_its_own_arm(): assert record["totals"]["cost_usd"] == 0 -def test_a_resume_checks_the_treatment_only_of_records_it_keeps(): +def test_a_rerunnable_record_from_another_treatment_is_refused_too(): + # Review r3: a rerun is not sure to replace a rerunnable record (a budget + # stop, lost workers, a narrower --scenarios), so every record is checked. pe3 = _profile("pe3") stale = _record(pe3, tool_features=["names"]) - previous = { + records = { "mt_000": dict(stale, status="complete"), - "mt_001": dict(stale, scenario_id="mt_001", status="harness_error"), + "mt_001": dict(stale, scenario_id="mt_001", status="provider_error"), + "mt_002": F.harness_error_record("mt_002", ["snap"], pe3, RuntimeError("x")), } - # Kept, so refused. - assert len(F.resume_treatment_problems(previous, {"mt_000"}, pe3)) == 1 - # Rerun under this profile, so no mixing and no block. - assert F.resume_treatment_problems(previous, set(), pe3) == [] + problems = F.treatment_problems(records, pe3) + assert [p.split()[0] for p in problems] == ["mt_000", "mt_001"] + + +SIX = ("concise", "explain", "flat", "names", "overrides", "placement") + + +def _scored_record(sid, profile, treatment_of=None, status="complete", **changes): + """A household record as command_run writes it, with the treatment of + `treatment_of` (another arm's profile) when given.""" + import hashlib + + board = pilot.load_board() + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + variables = board["variables"][sid] + source = treatment_of or profile + record = { + "scenario_id": sid, + "model": profile.name, + "arm": profile.arm, + "prompt_version": profile.prompt_version, + "benchmark_search_rule": profile.benchmark_search_rule, + "tool_notice_sha256": hashlib.sha256(source.notice.encode()).hexdigest(), + **F.profile_treatment(source), + "variables": variables, + "predictions": {v: float(reference[(sid, v)]) for v in variables}, + "explanations": {v: "x" for v in variables}, + "status": status, + "turns": [{"total_cost_usd": 1.25, "calls": []}], + "cap_hit": False, + "wall_seconds": 1.0, + } + record.update(changes) + record["totals"] = F.household_totals(record) + return record + + +def _write(directory, *records): + (directory / "households").mkdir(parents=True, exist_ok=True) + for record in records: + path = directory / "households" / f"{record['scenario_id']}.json" + path.write_text(json.dumps(record)) + + +def test_a_resume_refuses_an_old_treatments_provider_error_before_running( + tmp_path, monkeypatch +): + # The review's case: a pe3 folder holding a provider_error written under + # six features, with the budget already spent. Before the fix the resume + # passed, `one` stopped on the budget without replacing it, and scoring + # read its predictions as pe3's. + import argparse + + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + six = F.model_profile(F.MODEL_NAME, "pe:" + "+".join(SIX)) + _write(tmp_path, _scored_record("scenario_000", pe3, six, status="provider_error")) + monkeypatch.setattr(F, "_require_engine", lambda board: None) + monkeypatch.setenv(F.provider_key(pe3.model_id), "test-key") + + def no_pool(**kwargs): + raise AssertionError("the run got as far as starting the calculator") + + monkeypatch.setattr(F, "ToolPool", no_pool) + args = argparse.Namespace( + out_dir=str(tmp_path), + board=None, + model=F.MODEL_NAME, + arm="pe3", + prompt="v1", + forbid_benchmark_search=False, + scenarios=None, + limit=None, + workers=1, + budget_usd=1.0, + reserve_usd=0.0, + ) + with pytest.raises(SystemExit, match="another treatment"): + F.command_run(args) + assert not list(tmp_path.glob("manifest-*.json")) + + +def test_score_refuses_a_folder_mixing_treatments(tmp_path): + import argparse + + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + six = F.model_profile(F.MODEL_NAME, "pe:" + "+".join(SIX)) + _write( + tmp_path, + _scored_record("scenario_000", pe3), + _scored_record("scenario_001", pe3, six, status="provider_error"), + ) + args = argparse.Namespace(out_dir=str(tmp_path), complete_only=False) + with pytest.raises(SystemExit, match="another treatment"): + F.command_score(args) + assert not (tmp_path / "score.json").exists() + # The same folder without the stray record scores. + (tmp_path / "households" / "scenario_001.json").unlink() + F.command_score(args) + assert json.loads((tmp_path / "score.json").read_text())["households"] == 1 + + +def test_score_rebuilds_notices_on_the_boards_engine_and_restores_it( + tmp_path, monkeypatch +): + # A household set's records state its engine (2.60.4) in their notices and + # schema, so a check on the reference engine's text would refuse them all. + import argparse + + monkeypatch.setattr(F, "board_engine", lambda snapshot=None: "2.60.4") + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + with F.stating_engine("2.60.4"): + current = _scored_record("scenario_000", pe3) + _write(tmp_path, current) + args = argparse.Namespace(out_dir=str(tmp_path), complete_only=False) + F.command_score(args) + assert F._STATED_ENGINE["version"] == F.REFERENCE_ENGINE + # Beside it, a record sent under the reference engine's text mixes two. + _write(tmp_path, _scored_record("scenario_001", pe3)) + with pytest.raises(SystemExit, match="another tool notice"): + F.command_score(args) + assert F._STATED_ENGINE["version"] == F.REFERENCE_ENGINE + + +def test_score_takes_a_folder_written_wholly_before_the_texts_stated_the_engine( + tmp_path, monkeypatch, capsys +): + # Test batch 0's frontier calculator arms (b6d835a8) were told 2.41.0 while + # the board and calculator ran 2.60.4: one treatment, scored as it ran. A + # resume into it is still refused, since new records would say 2.60.4. + import argparse + + monkeypatch.setattr(F, "board_engine", lambda snapshot=None: "2.60.4") + pe = F.model_profile(F.MODEL_NAME, "pe") + _write( + tmp_path, _scored_record("scenario_000", pe), _scored_record("scenario_001", pe) + ) + F.command_score(argparse.Namespace(out_dir=str(tmp_path), complete_only=False)) + assert "name policyengine-us 2.41.0" in capsys.readouterr().out + assert json.loads((tmp_path / "score.json").read_text())["households"] == 2 + with F.stating_engine("2.60.4"): + assert len(F.treatment_problems(F.load_records(tmp_path), pe)) == 2 def test_flat_reports_a_malformed_first_value_once_and_not_as_a_conflict():