From 854eb4e9b35e129379e369d56124e68fef0128b2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 07:39:26 -0400 Subject: [PATCH 1/6] Add an internal tool-use pilot for Gemini 4 Argon (not a board condition) scripts/tool_pilot_argon.py gives Argon a calculate_household function that runs the reference engine (policyengine-us 2.41.0 with the reference builder's latest_final convention reform) on a situation the model writes as JSON, and compares it with Argon's frozen no-tool board row on the same 100 households. It writes only under --out-dir and labels its rows gemini-4-argon+pe-tool; the board's single no_tools condition is unchanged. Subcommands: oracle (reference situations through the tool engine), run, score (the board's analyze_no_tools on the scored outputs) and diagnose (one category per arm-2 miss). Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 1834 ++++++++++++++++++++++++++++++++ tests/test_tool_pilot_argon.py | 694 ++++++++++++ 2 files changed, 2528 insertions(+) create mode 100644 scripts/tool_pilot_argon.py create mode 100644 tests/test_tool_pilot_argon.py diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py new file mode 100644 index 00000000..f3ec8b32 --- /dev/null +++ b/scripts/tool_pilot_argon.py @@ -0,0 +1,1834 @@ +"""Internal pilot: Gemini 4 Argon with a PolicyEngine tool on the board's 100 US +households. Not a board condition, and nothing here is published. + +The benchmark has one condition, ``no_tools`` (docs/benchmark_card.md, +"Evaluation conditions"). This script stays outside the board pipeline: it +writes only under --out-dir, never touches the run store, the snapshot, the +cache or the dashboard, and labels its rows ``gemini-4-argon+pe-tool`` so no +board export can take them for a board row. + +Arms, on the same 100 households: + 1. No tools: Argon's frozen board predictions (release dashboard-data-20261014, + PR #221), read from the snapshot and never re-run. + 2. PolicyEngine tool: the board prompt plus TOOL_NOTICE, the board's + ``submit_outputs`` answer tool, and ``calculate_household``, which runs the + reference engine on a situation the model writes as JSON. The reference + engine is the installed policyengine-us (2.41.0) with the reference + builder's convention reform, reference_audit/2026-09-28/fixes/latest_final.py + (reference_audit/2026-10-12-engine-upgrade/scripts/build_references_upgrade.py + computes every reference on that system). The ``oracle`` subcommand checks + that the reference situations reproduce every scored reference through it. + +The model writes JSON only. The tool parses it, checks every name against the +engine's variable registry, and passes plain data to policyengine_us.Simulation +in a separate worker process with a per-call timeout. No model-written text is +executed. + +Subcommands: + oracle reference situations through the tool engine + run call the model (needs GEMINI_API_KEY); one JSON record per household + score both arms through the board's analyze_no_tools on the scored outputs + diagnose classify every arm-2 miss +""" + +from __future__ import annotations + +import argparse +import concurrent.futures +import copy +import difflib +import hashlib +import importlib.metadata +import importlib.util +import json +import math +import multiprocessing as mp +import os +import queue +import re +import subprocess +import sys +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Mapping + +import pandas as pd + +ROOT = Path(__file__).resolve().parents[1] +SNAPSHOT = ( + 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" +# The bytes the 2.41.0 references were built with: the sidecar's 2026-10-13 +# engine_upgrade revision pins latest_final.py to this sha256. The modules it +# composes are checked end to end by the oracle subcommand. +FIX_MODULE_SHA256 = "dbbdd228b99933c6f336af97970821c25c863a18606285e8d6d9ffb3bc9af78e" + +MODEL_NAME = "gemini-4-argon" +ARM1_LABEL = MODEL_NAME +ARM2_LABEL = "gemini-4-argon+pe-tool" +YEAR = 2026 +REFERENCE_ENGINE = "2.41.0" + +TOOL_NAME = "calculate_household" +TOOL_TURN_CAP = 8 +MAX_CALCULATIONS_PER_TURN = 4 +# Tool turns, then one forced answer turn, two repair turns and a margin for +# replies that call no function. +MAX_TOTAL_TURNS = TOOL_TURN_CAP + 5 +MAX_REPAIR_ROUNDS = 2 +MAX_PROVIDER_ATTEMPTS = 3 +TOOL_TIMEOUT_SECONDS = 60 +WORKER_START_TIMEOUT_SECONDS = 900 +DEFAULT_BUDGET_USD = 50.0 + +MAX_PEOPLE = 16 +MAX_UNITS_PER_GROUP = 16 +MAX_VARIABLES = 40 +MAX_SITUATION_CHARS = 200_000 +MAX_PROBLEMS = 20 +MAX_MESSAGE_CHARS = 2_000 +PERIOD_RE = re.compile(r"^(?:20[0-4]\d)(?:-(?:0[1-9]|1[0-2]))?$") + +GROUPS = { + "tax_units": "tax_unit", + "spm_units": "spm_unit", + "families": "family", + "marital_units": "marital_unit", + "households": "household", +} +ENTITY_GROUP = {"person": "people", **{v: k for k, v in GROUPS.items()}} + +TOOL_NOTICE = ( + "\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. You may call it in up to 8 turns. Then call `submit_outputs` exactly " + "once with your final answers, as instructed above." +) + +CALCULATE_DESCRIPTION = ( + "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " + "the requested variables. `situation` is a PolicyEngine household JSON " + "object with the entity groups people, tax_units, spm_units, families, " + "marital_units and households. Name each person and each unit; every unit " + "lists its `members` by person name, and a group left out puts everyone in " + "one unit. Set each input under its policyengine-us variable name as " + '{period: value}, for example "age": {"2026": 40}; a period is a year ' + '("2026") or a month ("2026-01"), and a bare value is read as the call\'s ' + "`period`. Inputs you do not set take the engine's defaults. `variables` " + 'lists policyengine-us variables to compute for `period` (default "2026"). ' + "Each result gives the variable's entity, label and value for each member of " + "that entity; a monthly variable computed for a year is summed over its " + "months. An unknown name, a variable set on the wrong entity, or a malformed " + "situation returns an error with suggestions instead of results." +) + + +def calculate_tool_schema() -> dict: + return { + "type": "function", + "function": { + "name": TOOL_NAME, + "description": CALCULATE_DESCRIPTION, + "parameters": { + "type": "object", + "properties": { + "situation": { + "type": "string", + "description": ( + "The household situation as a JSON object, written " + "as a JSON string." + ), + }, + "variables": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "policyengine-us variable names to compute (at most " + f"{MAX_VARIABLES})." + ), + }, + "period": { + "type": "string", + "description": 'Year or month to compute, default "2026".', + }, + }, + "required": ["situation", "variables"], + }, + }, + } + + +# --------------------------------------------------------------------------- +# The tool: validation (pure, testable without the engine) and execution. +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class VariableInfo: + entity: str + value_type: str + label: str | None = None + definition_period: str | None = None + possible_values: tuple[str, ...] | None = None + + +def build_registry(system) -> dict[str, VariableInfo]: + registry = {} + for name, variable in system.variables.items(): + value_type = getattr(variable.value_type, "__name__", str(variable.value_type)) + possible = None + if value_type == "Enum" and variable.possible_values is not None: + possible = tuple(member.name for member in variable.possible_values) + registry[name] = VariableInfo( + entity=variable.entity.key, + value_type=value_type, + label=variable.label, + definition_period=str(variable.definition_period), + possible_values=possible, + ) + return registry + + +def _suggest(name: str, registry: Mapping[str, VariableInfo], entity=None) -> str: + names = [ + n for n, info in registry.items() if entity is None or info.entity == entity + ] + close = difflib.get_close_matches(name, names, n=5, cutoff=0.6) + return f" Close matches: {', '.join(close)}." if close else "" + + +def _scalar_problem(value) -> str | None: + if isinstance(value, bool) or isinstance(value, str): + return None + if isinstance(value, (int, float)): + if not math.isfinite(float(value)): + return "is not a finite number" + return None + return "must be a number, true/false or a string" + + +def _type_problem(value, info: VariableInfo) -> str | None: + if info.value_type in {"float", "int"}: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return f"takes a number ({info.value_type})" + elif info.value_type == "bool": + if not (isinstance(value, bool) or value in (0, 1)): + return "takes true or false" + elif info.value_type == "Enum": + if not isinstance(value, str): + return "takes one of its named values as a string" + if info.possible_values and value not in info.possible_values: + shown = ", ".join(info.possible_values[:30]) + return f"takes one of: {shown}" + elif info.value_type == "str": + if not isinstance(value, str): + return "takes a string" + return None + + +def _normalize_inputs( + where: str, + entity: str, + data: dict, + registry: Mapping[str, VariableInfo], + period: str, + problems: list[str], + allow_members: bool, +) -> dict: + out: dict[str, Any] = {} + for name, value in data.items(): + if allow_members and name == "members": + continue + if not isinstance(name, str) or name not in registry: + problems.append( + f"{where}: unknown variable {name!r}." + + (_suggest(name, registry) if isinstance(name, str) else "") + ) + continue + info = registry[name] + if info.entity != entity: + 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): + 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 + out[name] = normalized + return out + + +def normalize_request( + args, registry: Mapping[str, VariableInfo] +) -> 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. + """ + problems: list[str] = [] + + def fail() -> tuple[None, dict]: + return None, { + "error": "validation", + "problems": problems[:MAX_PROBLEMS], + "problem_count": len(problems), + } + + if isinstance(args, str): + try: + args = json.loads(args) + except json.JSONDecodeError as error: + problems.append(f"arguments are not valid JSON: {error}") + return fail() + if not isinstance(args, dict): + problems.append("arguments must be an object") + return fail() + unknown = sorted(set(args) - {"situation", "variables", "period"}) + if unknown: + problems.append(f"unknown arguments: {', '.join(map(str, unknown))}") + + period = args.get("period", str(YEAR)) + if period in (None, ""): + period = str(YEAR) + if isinstance(period, int) and not isinstance(period, bool): + period = str(period) + if not isinstance(period, str) or not PERIOD_RE.match(period): + problems.append( + f'period {period!r} must be a year ("2026") or a month ("2026-01")' + ) + period = str(YEAR) + + situation = args.get("situation") + if isinstance(situation, str): + if len(situation) > MAX_SITUATION_CHARS: + problems.append("situation is too large") + return fail() + try: + situation = json.loads(situation) + except json.JSONDecodeError as error: + problems.append(f"situation is not valid JSON: {error}") + return fail() + if not isinstance(situation, dict): + problems.append("situation must be a JSON object") + return fail() + if len(json.dumps(situation, default=str)) > MAX_SITUATION_CHARS: + problems.append("situation is too large") + return fail() + extra = sorted(set(situation) - {"people", *GROUPS}) + if extra: + problems.append( + "situation accepts only people, " + + ", ".join(GROUPS) + + f"; remove {', '.join(map(str, extra))}" + ) + + people = situation.get("people") + normalized: dict[str, Any] = {} + if not isinstance(people, dict) or not people: + problems.append("situation.people must be a non-empty object of named people") + people = {} + elif len(people) > MAX_PEOPLE: + problems.append(f"at most {MAX_PEOPLE} people") + normalized_people = {} + for name, data in people.items(): + if not isinstance(name, str) or not name or len(name) > 64: + problems.append(f"person name {name!r} must be a short string") + continue + if not isinstance(data, 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 + ) + normalized["people"] = normalized_people + + for group, entity in GROUPS.items(): + if group not in situation: + continue + units = situation[group] + if not isinstance(units, dict) or not units: + problems.append(f"{group} must be a non-empty object of named units") + continue + if len(units) > MAX_UNITS_PER_GROUP: + problems.append(f"at most {MAX_UNITS_PER_GROUP} {group}") + seen: dict[str, str] = {} + normalized_units = {} + for unit_name, unit in units.items(): + where = f"{group}.{unit_name}" + if not isinstance(unit, dict): + problems.append(f"{where} must be an object with members") + continue + members = unit.get("members") + if ( + not isinstance(members, list) + or not members + or not all(isinstance(m, str) for m in members) + ): + problems.append(f"{where}.members must list person names") + members = [] + for member in members: + if member not in people: + problems.append(f"{where}: member {member!r} is not in people") + elif member in seen: + problems.append( + f"{where}: {member} is already in {group}.{seen[member]}" + ) + else: + seen[member] = unit_name + normalized_unit = {"members": list(members)} + normalized_unit.update( + _normalize_inputs(where, entity, unit, registry, period, problems, True) + ) + normalized_units[unit_name] = normalized_unit + normalized[group] = normalized_units + + variables = args.get("variables") + if isinstance(variables, str): + variables = [variables] + if ( + not isinstance(variables, list) + or not variables + or not all(isinstance(v, str) for v in variables) + ): + problems.append("variables must be a non-empty list of variable names") + variables = [] + elif len(variables) > MAX_VARIABLES: + problems.append(f"at most {MAX_VARIABLES} variables per call") + requested = [] + for name in variables: + if name not in registry: + problems.append( + f"variables: unknown variable {name!r}." + _suggest(name, registry) + ) + elif name not in requested: + requested.append(name) + + if problems: + return fail() + return {"situation": normalized, "variables": requested, "period": period}, None + + +def _jsonable(value): + if hasattr(value, "item"): + value = value.item() + if isinstance(value, bytes): + value = value.decode() + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + value = float(value) + if not math.isfinite(value): + return None + return round(value, 2) + return str(value) + + +def execute_request(system, registry, request: dict, simulation_cls=None) -> dict: + 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"]) + ) + except Exception as error: # the engine's own message is the useful one + return { + "error": "engine", + "message": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS], + } + results = {} + for name in request["variables"]: + info = registry[name] + try: + array = simulation.calculate(name, request["period"]) + except Exception as error: + results[name] = { + "error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS] + } + continue + if hasattr(array, "decode_to_str"): + array = array.decode_to_str() + ids = [str(i) for i in simulation.populations[info.entity].ids] + results[name] = { + "entity": info.entity, + "label": info.label, + "period": request["period"], + "values": {i: _jsonable(v) for i, v in zip(ids, list(array))}, + } + return {"results": results} + + +def load_fix_reform(path: Path = FIX_MODULE): + if path == FIX_MODULE and _sha256(path) != FIX_MODULE_SHA256: + raise SystemExit(f"{path} does not match the pinned sha256 {FIX_MODULE_SHA256}") + spec = importlib.util.spec_from_file_location("tool_pilot_latest_final", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module.reform + + +def build_tool_system(reform: bool = True, path: Path = FIX_MODULE): + from policyengine_us import CountryTaxBenefitSystem + + if not reform: + return CountryTaxBenefitSystem() + return CountryTaxBenefitSystem(reform=load_fix_reform(path)) + + +def _worker_main(conn, fix_path: str, reform: bool) -> 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)) + registry = build_registry(system) + conn.send(("ready", None)) + while True: + try: + args = conn.recv() + except EOFError: + return + if args is None: + return + request, error = normalize_request(args, registry) + if error is not None: + conn.send(("ok", error)) + continue + conn.send(("ok", execute_request(system, registry, request))) + + +class ToolPool: + """Worker processes that hold the tool system, each call under a timeout.""" + + def __init__( + self, + size: int, + fix_path: Path = FIX_MODULE, + timeout: float = TOOL_TIMEOUT_SECONDS, + reform: bool = True, + ): + self._ctx = mp.get_context("spawn") + self._fix_path = str(fix_path) + self._reform = reform + self.timeout = timeout + self._idle: queue.Queue = queue.Queue() + self.timeouts = 0 + self.crashes = 0 + starting = [self._start() for _ in range(size)] + for worker in starting: + self._idle.put(self._await_ready(worker)) + + def _start(self): + parent, child = self._ctx.Pipe() + process = self._ctx.Process( + target=_worker_main, + args=(child, self._fix_path, self._reform), + daemon=True, + ) + process.start() + return process, parent + + def _await_ready(self, worker): + process, conn = worker + if not conn.poll(WORKER_START_TIMEOUT_SECONDS): + process.kill() + raise RuntimeError("tool worker did not start") + status, _ = conn.recv() + if status != "ready": + raise RuntimeError(f"tool worker failed to start: {status}") + return worker + + def call(self, args) -> tuple[dict, float]: + worker = self._idle.get() + process, conn = worker + started = time.perf_counter() + try: + conn.send(args) + if conn.poll(self.timeout): + _, result = conn.recv() + self._idle.put(worker) + return result, time.perf_counter() - started + self.timeouts += 1 + result = { + "error": "timeout", + "message": f"the calculation exceeded {self.timeout:.0f} seconds", + } + except (EOFError, BrokenPipeError, OSError) as error: + self.crashes += 1 + result = {"error": "engine", "message": f"tool worker failed: {error}"} + elapsed = time.perf_counter() - started + process.kill() + process.join(timeout=10) + self._idle.put(self._await_ready(self._start())) + return result, elapsed + + def close(self) -> None: + while not self._idle.empty(): + process, conn = self._idle.get() + try: + conn.send(None) + except (BrokenPipeError, OSError): + pass + process.join(timeout=5) + if process.is_alive(): + process.kill() + + +# --------------------------------------------------------------------------- +# Households, prompts and the conversation loop. +# --------------------------------------------------------------------------- + + +def load_board(snapshot: Path = SNAPSHOT): + from policybench.reference_exclusions import ( + load_reference_exclusions, + split_reference, + ) + from policybench.scenarios import scenario_from_dict + from policybench.spec import expand_programs_for_scenario + + meta = json.loads((snapshot / "reference_outputs.csv.meta.json").read_text()) + reference = pd.read_csv(snapshot / "reference_outputs.csv") + scored, excluded = split_reference(reference, load_reference_exclusions(snapshot)) + scenarios_frame = pd.read_csv(snapshot / "scenarios.csv") + scenarios = {} + variables = {} + for _, row in scenarios_frame.iterrows(): + scenario = scenario_from_dict(json.loads(row["scenario_json"])) + scenarios[scenario.id] = scenario + variables[scenario.id] = expand_programs_for_scenario( + meta["programs"], scenario + ) + return { + "meta": meta, + "reference": reference, + "scored": scored, + "excluded": excluded, + "scenarios_frame": scenarios_frame, + "scenarios": scenarios, + "variables": variables, + } + + +def build_initial_request(scenario, variables: list[str]) -> tuple[list[dict], dict]: + """The board's request for this household, plus the tool notice and tool.""" + from policybench.config import MODELS + from policybench.eval_no_tools import ( + _build_answer_tool, + _chat_completion_request_kwargs, + _completion_budget_ceiling, + ) + + model_id = MODELS[MODEL_NAME] + messages, kwargs = _chat_completion_request_kwargs(scenario, variables, model_id) + board_prompt = messages[0]["content"] + messages = [{"role": "user", "content": board_prompt + TOOL_NOTICE}] + answer_tool = _build_answer_tool(variables, country=scenario.country) + kwargs = {k: v for k, v in kwargs.items() if k not in {"messages", "caching"}} + kwargs["tools"] = [calculate_tool_schema(), answer_tool] + kwargs["tool_choice"] = "required" + kwargs["_board_prompt_sha256"] = hashlib.sha256(board_prompt.encode()).hexdigest() + kwargs["_budget_ceiling"] = _completion_budget_ceiling(model_id) + return messages, kwargs + + +def forced_answer_choice() -> dict: + from policybench.eval_no_tools import ANSWER_FUNCTION_NAME, _tool_choice_for + + return _tool_choice_for(ANSWER_FUNCTION_NAME) + + +def _tool_calls(message) -> list: + return list(getattr(message, "tool_calls", None) or []) + + +def _call_name(call) -> str | None: + function = getattr(call, "function", None) + if function is None and isinstance(call, dict): + function = call.get("function") + if function is None: + return None + return getattr(function, "name", None) or ( + function.get("name") if isinstance(function, dict) else None + ) + + +def _call_arguments(call): + function = getattr(call, "function", None) + if function is None and isinstance(call, dict): + function = call.get("function") + arguments = getattr(function, "arguments", None) + if arguments is None and isinstance(function, dict): + arguments = function.get("arguments") + return arguments + + +def _call_id(call) -> str: + return getattr(call, "id", None) or ( + call.get("id") if isinstance(call, dict) else "" + ) + + +def _message_dict(message) -> dict: + if hasattr(message, "model_dump"): + data = message.model_dump() + elif isinstance(message, dict): + data = dict(message) + else: + data = dict(vars(message)) + return {k: v for k, v in data.items() if v is not None} + + +def _has_thought_signature(message) -> bool: + for call in _tool_calls(message): + if "__thought__" in str(_call_id(call)): + return True + fields = getattr(call, "provider_specific_fields", None) or {} + if isinstance(fields, dict) and fields.get("thought_signature"): + return True + fields = getattr(message, "provider_specific_fields", None) or {} + return bool(isinstance(fields, dict) and fields.get("thought_signatures")) + + +class Spend: + def __init__(self, budget_usd: float, spent_usd: float = 0.0): + self.budget_usd = budget_usd + self.spent_usd = spent_usd + self._lock = threading.Lock() + + def add(self, usd: float | None) -> None: + with self._lock: + self.spent_usd += float(usd or 0.0) + + def exhausted(self, reserve_usd: float = 0.0) -> bool: + with self._lock: + return self.spent_usd + reserve_usd > self.budget_usd + + +class BudgetStop(RuntimeError): + pass + + +def _request_with_retries( + completion_fn, kwargs: dict +) -> tuple[Any, float, float, float]: + from policybench.eval_no_tools import _run_request_with_wall_timeout, _should_retry + + for attempt in range(MAX_PROVIDER_ATTEMPTS): + started_at = time.time() + started = time.perf_counter() + try: + response = _run_request_with_wall_timeout(completion_fn, kwargs) + return response, started_at, time.time(), time.perf_counter() - started + except Exception as error: + if attempt == MAX_PROVIDER_ATTEMPTS - 1 or not _should_retry(error): + raise + time.sleep(2 * (2**attempt)) + raise RuntimeError("unreachable") + + +def run_household( + scenario, + variables: list[str], + tool_call: Callable[[Any], tuple[dict, float]], + completion_fn: Callable[..., Any], + spend: Spend, + turn_cap: int = TOOL_TURN_CAP, +) -> dict: + """One household's conversation: tool turns, then the board's answer tool.""" + from policybench.config import MODELS + from policybench.eval_no_tools import ( + _apply_thinking_configuration, + _build_answer_tool, + _enforce_explanation_value_contract, + _extract_usage_metadata, + _missing_explanations, + _response_finish_reason, + _serialize_response_payload, + extract_explanations, + extract_predictions, + ) + + model_id = MODELS[MODEL_NAME] + messages, base = build_initial_request(scenario, variables) + board_prompt_sha256 = base.pop("_board_prompt_sha256") + ceiling = base.pop("_budget_ceiling") + calculate_tool = base["tools"][0] + record: dict[str, Any] = { + "scenario_id": scenario.id, + "variables": variables, + "board_prompt_sha256": board_prompt_sha256, + "tool_notice_sha256": hashlib.sha256(TOOL_NOTICE.encode()).hexdigest(), + "turns": [], + "status": "running", + } + predictions: dict[str, float | None] = {v: None for v in variables} + explanations: dict[str, str | None] = {v: None for v in variables} + pending = list(variables) + tool_turns = 0 + repair_rounds = 0 + forced = False + answer_variables = list(variables) + household_started = time.perf_counter() + + for turn_index in range(MAX_TOTAL_TURNS): + if spend.exhausted(): + record["status"] = "budget_stop" + break + if repair_rounds: + kind = "repair" + tool_choice = forced_answer_choice() + tools = [ + calculate_tool, + _build_answer_tool(answer_variables, country=scenario.country), + ] + elif forced or tool_turns >= turn_cap: + kind = "forced_answer" + forced = True + tool_choice = forced_answer_choice() + tools = base["tools"] + else: + kind = "open" + tool_choice = "required" + tools = base["tools"] + kwargs = {k: v for k, v in base.items() if k not in {"tools", "tool_choice"}} + kwargs.update( + { + "model": model_id, + "messages": messages, + "tools": tools, + "tool_choice": tool_choice, + } + ) + _apply_thinking_configuration(kwargs, model_id) + budget_key = ( + "max_completion_tokens" + if "max_completion_tokens" in kwargs + else "max_tokens" + ) + escalated_from = None + try: + response, started_at, completed_at, latency = _request_with_retries( + completion_fn, kwargs + ) + finish_reason = _response_finish_reason(response, responses_api=False) + if finish_reason == "length" and kwargs.get(budget_key, 0) < ceiling: + usage = _extract_usage_metadata(response, model_id, messages, "") + spend.add(usage.get("total_cost_usd")) + record["turns"].append( + { + "index": turn_index, + "kind": kind, + "truncated_and_retried": True, + "completion_budget": kwargs.get(budget_key), + "latency_seconds": latency, + **usage, + } + ) + escalated_from = kwargs.get(budget_key) + kwargs[budget_key] = ceiling + response, started_at, completed_at, latency = _request_with_retries( + completion_fn, kwargs + ) + finish_reason = _response_finish_reason(response, responses_api=False) + except Exception as error: + record["status"] = "provider_error" + record["error"] = f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS] + break + + message = response.choices[0].message + content = getattr(message, "content", None) + calls = _tool_calls(message) + raw = _serialize_response_payload( + content=content, tool_calls=calls, function_call=None + ) + usage = _extract_usage_metadata( + response, model_id, messages, raw or content or "" + ) + spend.add(usage.get("total_cost_usd")) + turn = { + "index": turn_index, + "kind": kind, + "tool_choice": tool_choice, + "completion_budget": kwargs.get(budget_key), + "escalated_from": escalated_from, + "request_started_at": started_at, + "request_completed_at": completed_at, + "latency_seconds": latency, + "finish_reason": finish_reason, + "content": content, + "reasoning_content": getattr(message, "reasoning_content", None), + "has_thought_signature": _has_thought_signature(message), + "calls": [], + **usage, + } + record["turns"].append(turn) + messages.append(_message_dict(message)) + + answer_calls = [c for c in calls if _call_name(c) == "submit_outputs"] + calculate_calls = [c for c in calls if _call_name(c) == TOOL_NAME] + other_calls = [ + c for c in calls if c not in answer_calls and c not in calculate_calls + ] + + if answer_calls: + got = extract_predictions( + content, answer_variables, tool_calls=answer_calls + ) + got_explanations = extract_explanations( + content, answer_variables, tool_calls=answer_calls + ) + got, got_explanations = _enforce_explanation_value_contract( + got, got_explanations, answer_variables + ) + for variable in answer_variables: + if got.get(variable) is not None: + predictions[variable] = got[variable] + if got_explanations.get(variable): + explanations[variable] = got_explanations[variable] + turn["calls"].append( + { + "name": "submit_outputs", + "arguments": _call_arguments(answer_calls[0]), + "ignored_calls": len(calculate_calls) + + len(other_calls) + + len(answer_calls) + - 1, + } + ) + pending = [v for v in variables if predictions[v] is None] + missing_explanation = _missing_explanations(explanations, variables) + still = [ + v for v in variables if v in set(pending) | set(missing_explanation) + ] + for call in calls: + messages.append( + { + "role": "tool", + "tool_call_id": _call_id(call), + "name": _call_name(call), + "content": json.dumps( + {"accepted": not still, "missing": still} + if call is answer_calls[0] + else {"error": "not run: submit_outputs ended the turn"} + ), + } + ) + if not still: + record["status"] = "complete" + break + if repair_rounds >= MAX_REPAIR_ROUNDS: + record["status"] = "incomplete" + break + repair_rounds += 1 + answer_variables = still + messages.append( + { + "role": "user", + "content": ( + "Your submit_outputs call omitted or left unusable these " + f"outputs: {', '.join(still)}. Call submit_outputs again " + "with only these outputs, each with a numeric value and an " + "explanation ending with `value = X`." + ), + } + ) + continue + + if calculate_calls: + tool_turns += 1 + for position, call in enumerate(calculate_calls): + if position >= MAX_CALCULATIONS_PER_TURN: + result, tool_latency = ( + { + "error": "not_run", + "message": ( + f"at most {MAX_CALCULATIONS_PER_TURN} " + "calculations per turn" + ), + }, + 0.0, + ) + else: + result, tool_latency = tool_call(_call_arguments(call)) + turn["calls"].append( + { + "name": TOOL_NAME, + "arguments": _call_arguments(call), + "result": result, + "tool_latency_seconds": tool_latency, + } + ) + messages.append( + { + "role": "tool", + "tool_call_id": _call_id(call), + "name": TOOL_NAME, + "content": json.dumps(result), + } + ) + for call in other_calls: + messages.append( + { + "role": "tool", + "tool_call_id": _call_id(call), + "name": _call_name(call) or "unknown", + "content": json.dumps({"error": "unknown function"}), + } + ) + if tool_turns >= turn_cap: + messages.append( + { + "role": "user", + "content": ( + f"You have used all {turn_cap} calculation turns. Call " + "submit_outputs now with your final answers." + ), + } + ) + continue + + for call in other_calls: + messages.append( + { + "role": "tool", + "tool_call_id": _call_id(call), + "name": _call_name(call) or "unknown", + "content": json.dumps({"error": "unknown function"}), + } + ) + messages.append( + { + "role": "user", + "content": "Call calculate_household or submit_outputs.", + } + ) + else: + record["status"] = "turn_limit" + + if record["status"] == "running": + record["status"] = "incomplete" + record["predictions"] = predictions + record["explanations"] = explanations + record["tool_turns"] = tool_turns + record["cap_hit"] = tool_turns >= turn_cap + record["repair_rounds"] = repair_rounds + record["wall_seconds"] = time.perf_counter() - household_started + record["messages"] = messages + record["totals"] = household_totals(record) + return record + + +def _num(value) -> float: + try: + number = float(value) + except (TypeError, ValueError): + return 0.0 + return number if math.isfinite(number) else 0.0 + + +def household_totals(record: dict) -> dict: + turns = record.get("turns", []) + calls = [c for t in turns for c in t.get("calls", []) if c.get("name") == TOOL_NAME] + errors = [ + c for c in calls if isinstance(c.get("result"), dict) and "error" in c["result"] + ] + return { + "model_turns": len(turns), + "tool_calls": len(calls), + "tool_errors": len(errors), + "tool_error_kinds": sorted({c["result"]["error"] for c in errors}), + "prompt_tokens": sum(_num(t.get("prompt_tokens")) for t in turns), + "completion_tokens": sum(_num(t.get("completion_tokens")) for t in turns), + "reasoning_tokens": sum(_num(t.get("reasoning_tokens")) for t in turns), + "cached_prompt_tokens": sum(_num(t.get("cached_prompt_tokens")) for t in turns), + "cost_usd": sum(_num(t.get("total_cost_usd")) for t in turns), + "model_seconds": sum(_num(t.get("latency_seconds")) for t in turns), + "tool_seconds": sum(_num(c.get("tool_latency_seconds")) for c in calls), + } + + +# --------------------------------------------------------------------------- +# Subcommands. +# --------------------------------------------------------------------------- + + +def _git(*args: str) -> str: + try: + return subprocess.run( + ["git", "-C", str(ROOT), *args], capture_output=True, text=True, check=True + ).stdout.strip() + except (subprocess.CalledProcessError, FileNotFoundError): + return "" + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _versions() -> dict: + out = {} + for name in ("policyengine-us", "policyengine-core", "litellm"): + try: + out[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + out[name] = None + return out + + +def _require_engine() -> None: + version = _versions()["policyengine-us"] + if version != REFERENCE_ENGINE: + raise SystemExit( + f"policyengine-us {version} is installed; the pilot pins " + f"{REFERENCE_ENGINE}, the reference engine of release " + "dashboard-data-20261014." + ) + + +def command_oracle(args) -> None: + """Reference situations through the tool engine: every scored reference must + come back within the board's tolerance.""" + _require_engine() + from policybench.ground_truth import _extract_person_value, _pe_variable_for_output + + board = load_board() + 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) + registry = build_registry(system) + for sid, scenario in board["scenarios"].items(): + situation = scenario.to_pe_household() + request, error = normalize_request( + {"situation": situation, "variables": ["age"]}, registry + ) + if error is not None: + raise SystemExit( + f"{sid}: the validator rejects the reference situation: {error}" + ) + from policyengine_us import Simulation + + simulation = Simulation( + tax_benefit_system=system, situation=request["situation"] + ) + # The same household through the tool's execute path and the + # diagnosis's output mapping, as a model's call would go. + outputs = [v for v in board["variables"][sid] if (sid, v) in values.index] + call_args = { + "situation": json.dumps(situation), + "variables": sorted( + {_pe_variable_for_output(v, "us") for v in outputs} + ), + } + tool_request, tool_error = normalize_request(call_args, registry) + if tool_error is not None: + raise SystemExit(f"{sid}: {tool_error}") + call = { + "arguments": call_args, + "result": execute_request(system, registry, tool_request), + } + for variable in outputs: + value = float( + _extract_person_value( + simulation.calculate( + _pe_variable_for_output(variable, "us"), YEAR + ), + scenario, + variable, + ) + ) + rows.append( + { + "engine": label, + "scenario_id": sid, + "variable": variable, + "scored": (sid, variable) in scored_keys, + "reference": float(values[(sid, variable)]), + "tool": value, + "tool_path": tool_value_for_output(call, scenario, variable), + } + ) + frame = pd.DataFrame(rows) + frame["hit"] = [ + hit(variable, tool, reference) + for variable, tool, reference in zip( + frame.variable, frame.tool, frame.reference + ) + ] + frame["tool_path_hit"] = [ + hit(variable, tool, reference) + for variable, tool, reference in zip( + frame.variable, frame.tool_path, frame.reference + ) + ] + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + frame.to_csv(out / "oracle.csv", index=False) + summary = ( + frame[frame.scored] + .groupby("engine") + .agg( + scored_outputs=("hit", "size"), + reproduced=("hit", "sum"), + reproduced_through_tool_path=("tool_path_hit", "sum"), + ) + ) + print(summary.to_string()) + (out / "oracle_summary.json").write_text( + json.dumps( + { + "versions": _versions(), + "fix_module_sha256": _sha256(FIX_MODULE), + "summary": summary.reset_index().to_dict(orient="records"), + }, + indent=2, + ) + ) + + +def hit(variable: str, prediction, reference) -> bool: + """The board's exact-match test (policybench.analysis.row_hit_scores).""" + from policybench.analysis import row_hit_scores + + if prediction is None or reference is None: + return False + try: + if pd.isna(prediction) or pd.isna(reference): + return False + except (TypeError, ValueError): + return False + return row_hit_scores(variable, float(reference), float(prediction))["exact"] == 1.0 + + +def command_run(args) -> None: + _require_engine() + if not os.environ.get("GEMINI_API_KEY"): + raise SystemExit("GEMINI_API_KEY is not set") + from litellm import completion + + board = load_board() + out = Path(args.out_dir) + households_dir = out / "households" + households_dir.mkdir(parents=True, exist_ok=True) + ids = sorted(board["scenarios"]) + if args.scenarios: + ids = [sid for sid in ids if sid in set(args.scenarios)] + if args.limit: + ids = ids[: args.limit] + + spent_before = 0.0 + done = set() + for path in households_dir.glob("*.json"): + record = json.loads(path.read_text()) + spent_before += _num(record.get("totals", {}).get("cost_usd")) + if record.get("status") == "complete": + done.add(record["scenario_id"]) + spend = Spend(args.budget_usd, spent_before) + todo = [sid for sid in ids if sid not in done] + manifest = { + "started_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "git_head": _git("rev-parse", "HEAD"), + "git_dirty": bool(_git("status", "--porcelain")), + "script_sha256": _sha256(Path(__file__)), + "versions": _versions(), + "fix_module": str(FIX_MODULE.relative_to(ROOT)), + "fix_module_sha256": _sha256(FIX_MODULE), + "model": MODEL_NAME, + "tool_notice": TOOL_NOTICE, + "calculate_tool": calculate_tool_schema(), + "turn_cap": TOOL_TURN_CAP, + "max_calculations_per_turn": MAX_CALCULATIONS_PER_TURN, + "tool_timeout_seconds": TOOL_TIMEOUT_SECONDS, + "budget_usd": args.budget_usd, + "spent_before_usd": spent_before, + "households": todo, + "workers": args.workers, + } + (out / f"manifest-{manifest['started_at_utc'].replace(':', '')}.json").write_text( + json.dumps(manifest, indent=2) + ) + print(f"{len(todo)} households to run; ${spent_before:.4f} already spent") + if not todo: + return + pool = ToolPool(size=args.workers) + try: + + def one(sid: str) -> dict: + if spend.exhausted(args.reserve_usd): + return {"scenario_id": sid, "status": "budget_stop"} + record = run_household( + board["scenarios"][sid], + board["variables"][sid], + pool.call, + completion, + spend, + ) + (households_dir / f"{sid}.json").write_text( + json.dumps(record, indent=1, default=str) + ) + totals = record["totals"] + print( + f"{sid} {record['status']} turns={totals['model_turns']} " + f"calls={totals['tool_calls']} errors={totals['tool_errors']} " + f"cost=${totals['cost_usd']:.4f} spent=${spend.spent_usd:.4f}", + flush=True, + ) + return record + + with concurrent.futures.ThreadPoolExecutor( + max_workers=args.workers + ) as executor: + results = list(executor.map(one, todo)) + finally: + pool.close() + stopped = [r["scenario_id"] for r in results if r.get("status") == "budget_stop"] + print( + f"spent ${spend.spent_usd:.4f} of ${args.budget_usd:.2f}; " + f"budget stops: {stopped}" + ) + + +def load_records(out_dir: Path) -> dict[str, dict]: + records = {} + for path in sorted((out_dir / "households").glob("*.json")): + record = json.loads(path.read_text()) + records[record["scenario_id"]] = record + return records + + +def arm2_predictions(records: Mapping[str, dict]) -> pd.DataFrame: + rows = [] + for sid, record in records.items(): + for variable in record["variables"]: + value = record["predictions"].get(variable) + rows.append( + { + "model": ARM2_LABEL, + "scenario_id": sid, + "variable": variable, + "prediction": value, + "explanation": record["explanations"].get(variable), + "error": None if value is not None else record["status"], + } + ) + return pd.DataFrame(rows) + + +def arm1_predictions(snapshot: Path = SNAPSHOT) -> pd.DataFrame: + frame = pd.read_csv(snapshot / "predictions.csv.gz", low_memory=False) + return frame[frame["model"] == ARM1_LABEL].copy() + + +def command_score(args) -> None: + from policybench.analysis import analyze_no_tools + from policybench.spec import output_group_id + + board = load_board() + out = Path(args.out_dir) + records = load_records(out) + complete = sorted( + sid for sid, r in records.items() if r.get("status") == "complete" + ) + households = sorted(records) if args.all_records else complete + if not households: + raise SystemExit("no household records to score") + scored = board["scored"][board["scored"].scenario_id.isin(households)] + arm1 = arm1_predictions() + arm1 = arm1[arm1.scenario_id.isin(households)] + arm2 = arm2_predictions({sid: records[sid] for sid in households}) + scenarios = board["scenarios_frame"][ + board["scenarios_frame"].scenario_id.isin(households) + ] + analysis = analyze_no_tools( + scored, pd.concat([arm1, arm2], ignore_index=True), scenarios=scenarios + ) + summary = analysis["model_summary"] + columns = [ + c + for c in ( + "model", + "weighted_exact", + "mean_exact", + "total_n", + "parsed_n", + "weighted_within_1pct", + "bounded_score", + ) + if c in summary.columns + ] + print(summary[columns].to_string(index=False)) + + rows = scored[["scenario_id", "variable", "value"]].rename( + columns={"value": "reference"} + ) + a1 = arm1.set_index(["scenario_id", "variable"])["prediction"] + a2 = arm2.set_index(["scenario_id", "variable"])["prediction"] + rows["arm1"] = [a1.get((s, v)) for s, v in zip(rows.scenario_id, rows.variable)] + rows["arm2"] = [a2.get((s, v)) for s, v in zip(rows.scenario_id, rows.variable)] + rows["arm1_hit"] = [ + hit(v, p, r) for v, p, r in zip(rows.variable, rows.arm1, rows.reference) + ] + rows["arm2_hit"] = [ + hit(v, p, r) for v, p, r in zip(rows.variable, rows.arm2, rows.reference) + ] + rows["program"] = rows.variable.map(output_group_id) + rows.to_csv(out / "scored_rows.csv", index=False) + per_program = ( + rows.groupby("program") + .agg( + scored_outputs=("variable", "size"), + arm1_hits=("arm1_hit", "sum"), + arm2_hits=("arm2_hit", "sum"), + ) + .reset_index() + ) + per_program["arm1_rate"] = per_program.arm1_hits / per_program.scored_outputs + per_program["arm2_rate"] = per_program.arm2_hits / per_program.scored_outputs + per_program["change_pp"] = 100 * (per_program.arm2_rate - per_program.arm1_rate) + per_program.to_csv(out / "per_program.csv", index=False) + print(per_program.to_string(index=False)) + + usage = usage_by_household(records, arm1, households) + usage.to_csv(out / "usage_by_household.csv", index=False) + usage_summary = summarize_usage(usage) + print(json.dumps(usage_summary, indent=2)) + (out / "score.json").write_text( + json.dumps( + { + "households": len(households), + "incomplete_households": sorted(set(records) - set(complete)), + "scored_outputs": int(len(scored)), + "summary": summary[columns].to_dict(orient="records"), + "unweighted_hits": { + "arm1": int(rows.arm1_hit.sum()), + "arm2": int(rows.arm2_hit.sum()), + }, + "usage": usage_summary, + }, + indent=2, + default=float, + ) + ) + + +def usage_by_household(records, arm1: pd.DataFrame, households) -> pd.DataFrame: + a1 = arm1.copy() + for column in ( + "prompt_tokens", + "completion_tokens", + "reasoning_tokens", + "total_cost_usd", + "elapsed_seconds", + "cached_prompt_tokens", + ): + a1[column] = pd.to_numeric(a1[column], errors="coerce").fillna(0.0) + grouped = a1.groupby("scenario_id").agg( + arm1_prompt_tokens=("prompt_tokens", "sum"), + arm1_completion_tokens=("completion_tokens", "sum"), + arm1_reasoning_tokens=("reasoning_tokens", "sum"), + arm1_cost_usd=("total_cost_usd", "sum"), + arm1_seconds=("elapsed_seconds", "sum"), + arm1_calls=("call_id", "nunique"), + ) + rows = [] + for sid in households: + totals = records[sid]["totals"] + rows.append( + { + "scenario_id": sid, + "status": records[sid]["status"], + "arm2_model_turns": totals["model_turns"], + "arm2_tool_calls": totals["tool_calls"], + "arm2_tool_errors": totals["tool_errors"], + "arm2_cap_hit": records[sid].get("cap_hit"), + "arm2_prompt_tokens": totals["prompt_tokens"], + "arm2_cached_prompt_tokens": totals["cached_prompt_tokens"], + "arm2_completion_tokens": totals["completion_tokens"], + "arm2_reasoning_tokens": totals["reasoning_tokens"], + "arm2_cost_usd": totals["cost_usd"], + "arm2_model_seconds": totals["model_seconds"], + "arm2_tool_seconds": totals["tool_seconds"], + "arm2_wall_seconds": records[sid].get("wall_seconds"), + } + ) + frame = pd.DataFrame(rows).merge( + grouped.reset_index(), on="scenario_id", how="left" + ) + return frame + + +def summarize_usage(usage: pd.DataFrame) -> dict: + def stats(column): + series = pd.to_numeric(usage[column], errors="coerce") + return { + "total": float(series.sum()), + "mean": float(series.mean()), + "median": float(series.median()), + "max": float(series.max()), + } + + keys = [ + c + for c in usage.columns + if c.startswith(("arm1_", "arm2_")) and c != "arm2_cap_hit" + ] + out = {key: stats(key) for key in keys} + out["arm2_cap_hit_households"] = int( + usage["arm2_cap_hit"].fillna(False).astype(bool).sum() + ) + return out + + +# --------------------------------------------------------------------------- +# Miss diagnosis. +# --------------------------------------------------------------------------- + +MISS_CATEGORIES = ( + "wrong_input_mapping", + "wrong_variable_or_period", + "tool_error_or_cap", + "answer_differs_from_tool", + "no_answer", +) + + +def _successful_calls(record: dict) -> list[dict]: + calls = [] + for turn in record.get("turns", []): + for call in turn.get("calls", []): + result = call.get("result") + if ( + call.get("name") == TOOL_NAME + and isinstance(result, dict) + and "results" in result + ): + calls.append(call) + return calls + + +def _person_names(arguments) -> list[str]: + try: + args = json.loads(arguments) if isinstance(arguments, str) else arguments + situation = args["situation"] + situation = json.loads(situation) if isinstance(situation, str) else situation + return list(situation["people"]) + except (KeyError, TypeError, ValueError): + return [] + + +def tool_value_for_output(call: dict, scenario, output: str): + """The value one successful call returned for a benchmark output, or None. + + The output's engine variable must be in the call's results at the 2026 + annual period. Person outputs read the model's person in the scenario's + position (the prompt lists people in scenario order); group outputs take + the board's aggregation over the returned units. + """ + from policybench.ground_truth import ( + _aggregation_for_output, + _pe_variable_for_output, + ) + + pe_variable = _pe_variable_for_output(output, "us") + entry = call["result"]["results"].get(pe_variable) + if not isinstance(entry, dict) or "values" not in entry: + return None + if str(entry.get("period")) != str(YEAR): + return None + values = entry["values"] + aggregation = _aggregation_for_output(output, "us") + if aggregation == "person": + from policybench.ground_truth import _person_index_for_output + + index = _person_index_for_output(scenario, output) + names = _person_names(call.get("arguments")) + if index is None or index >= len(names) or names[index] not in values: + return None + value = values[names[index]] + else: + numbers = [ + float(v) if not isinstance(v, str) else math.nan for v in values.values() + ] + value = sum(numbers) + if aggregation == "any_positive": + value = float(value > 0) + if isinstance(value, bool): + return float(value) + try: + return float(value) + except (TypeError, ValueError): + return None + + +def classify_miss(record: dict, scenario, output: str, submitted, reference) -> dict: + """One of MISS_CATEGORIES, with the evidence behind it. + + Assumes the output missed. The last successful call that computed the + output's engine variable for 2026 decides: if the submitted value differs + from what the tool returned there, the answer departs from the tool; if it + matches, the tool computed a different value from the model's situation + than the reference situation gives, so the inputs were mapped differently + (a reference issue is a reading of that mapping, decided on review). + """ + from policybench.ground_truth import _pe_variable_for_output + + pe_variable = _pe_variable_for_output(output, "us") + calls = _successful_calls(record) + evidence = {"pe_variable": pe_variable, "successful_calls": len(calls)} + if submitted is None or (isinstance(submitted, float) and math.isnan(submitted)): + return {"category": "no_answer", **evidence} + computed = [ + c for c in calls if tool_value_for_output(c, scenario, output) is not None + ] + if not computed: + any_period = [ + c + for c in calls + if pe_variable in c["result"]["results"] + and isinstance(c["result"]["results"][pe_variable], dict) + and "values" in c["result"]["results"][pe_variable] + ] + if not calls: + return { + "category": "tool_error_or_cap", + **evidence, + "cap_hit": record.get("cap_hit"), + } + if any_period: + return { + "category": "wrong_variable_or_period", + "detail": "variable computed only for another period", + **evidence, + } + return { + "category": "wrong_variable_or_period", + "detail": "output's engine variable never computed", + **evidence, + } + last = computed[-1] + tool_value = tool_value_for_output(last, scenario, output) + evidence["tool_value"] = tool_value + evidence["call_arguments"] = last.get("arguments") + if not hit(output, submitted, tool_value): + return {"category": "answer_differs_from_tool", **evidence} + return {"category": "wrong_input_mapping", **evidence} + + +def situation_diff( + model_situation: dict, reference_situation: dict, scenario +) -> list[dict]: + """Inputs set differently in the model's situation and the reference's. + + People are matched by position (the prompt lists them in scenario order); + each group's single unit is matched to the reference's. An input absent on + one side compares as the other side's value against the engine default, + which the caller supplies through ``defaults``. + """ + diffs = [] + model_people = list(model_situation.get("people", {}).items()) + reference_people = list(reference_situation["people"].items()) + for index, (ref_name, ref_inputs) in enumerate(reference_people): + model_inputs = model_people[index][1] if index < len(model_people) else {} + model_name = model_people[index][0] if index < len(model_people) else None + for name in sorted(set(ref_inputs) | set(model_inputs)): + ref = ref_inputs.get(name, {}) + mod = model_inputs.get(name, {}) + if _period_value(ref) != _period_value(mod): + diffs.append( + { + "entity": "person", + "reference_unit": ref_name, + "model_unit": model_name, + "variable": name, + "reference": ref, + "model": mod, + } + ) + if len(model_people) != len(reference_people): + diffs.append( + { + "entity": "person", + "variable": "", + "reference": len(reference_people), + "model": len(model_people), + } + ) + for group in GROUPS: + ref_units = list(reference_situation.get(group, {}).values()) + model_units = list(model_situation.get(group, {}).values()) + ref_unit = ref_units[0] if ref_units else {} + model_unit = model_units[0] if model_units else {} + # A group the model leaves out is one unit holding everyone. + if (len(model_units) or 1) != len(ref_units): + diffs.append( + { + "entity": group, + "variable": "", + "reference": len(ref_units), + "model": len(model_units), + } + ) + for name in sorted((set(ref_unit) | set(model_unit)) - {"members"}): + ref = ref_unit.get(name, {}) + mod = model_unit.get(name, {}) + if _period_value(ref) != _period_value(mod): + diffs.append( + {"entity": group, "variable": name, "reference": ref, "model": mod} + ) + return diffs + + +def _period_value(value) -> Any: + if isinstance(value, dict): + return tuple(sorted((str(k), _canon(v)) for k, v in value.items())) + return ((str(YEAR), _canon(value)),) if value != {} else () + + +def _canon(value): + if isinstance(value, bool): + return float(value) + if isinstance(value, (int, float)): + return round(float(value), 6) + return value + + +def command_diagnose(args) -> None: + board = load_board() + out = Path(args.out_dir) + records = load_records(out) + rows_path = out / "scored_rows.csv" + if not rows_path.exists(): + raise SystemExit("run score first") + rows = pd.read_csv(rows_path) + misses = rows[~rows.arm2_hit.astype(bool)] + registry = None + defaults = {} + if not args.no_engine_defaults: + system = build_tool_system(reform=True) + registry = build_registry(system) + defaults = { + name: system.variables[name].default_value for name in system.variables + } + out_rows = [] + for row in misses.itertuples(): + record = records[row.scenario_id] + scenario = board["scenarios"][row.scenario_id] + submitted = None if pd.isna(row.arm2) else float(row.arm2) + verdict = classify_miss( + record, scenario, row.variable, submitted, row.reference + ) + diff = [] + arguments = verdict.get("call_arguments") + if arguments is not None and registry is not None: + request, _ = normalize_request(arguments, registry) + if request is not None: + diff = situation_diff( + request["situation"], scenario.to_pe_household(), scenario + ) + diff = [d for d in diff if not _is_default_only(d, defaults)] + out_rows.append( + { + "scenario_id": row.scenario_id, + "state": scenario.state, + "variable": row.variable, + "program": row.program, + "reference": row.reference, + "arm1": row.arm1, + "arm1_hit": row.arm1_hit, + "arm2": submitted, + "category": verdict["category"], + "detail": verdict.get("detail"), + "tool_value": verdict.get("tool_value"), + "status": record.get("status"), + "cap_hit": record.get("cap_hit"), + "input_diff": json.dumps(diff, default=str), + } + ) + frame = pd.DataFrame(out_rows) + frame.to_csv(out / "misses.csv", index=False) + print(frame.groupby("category").size().to_string()) + print(frame.groupby(["program", "category"]).size().to_string()) + + +def _is_default_only(diff: dict, defaults: Mapping[str, Any]) -> bool: + """A difference that is only an input left at the engine default on one side.""" + name = diff.get("variable") + if name not in defaults: + return False + default = defaults[name] + if hasattr(default, "name"): + default = default.name + sides = [] + for side in (diff.get("reference"), diff.get("model")): + if side in ({}, None): + sides.append(None) + elif isinstance(side, dict): + sides.append(tuple(sorted(_canon(v) for v in side.values()))) + else: + sides.append((_canon(side),)) + canon_default = (_canon(default),) + if sides[0] is None: + return sides[1] == canon_default + if sides[1] is None: + return sides[0] == canon_default + return False + + +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) + for name in ("oracle", "run", "score", "diagnose"): + p = sub.add_parser(name) + p.add_argument("--out-dir", required=True) + if name == "run": + p.add_argument("--scenarios", nargs="*") + p.add_argument("--limit", type=int) + p.add_argument("--workers", type=int, default=4) + p.add_argument("--budget-usd", type=float, default=DEFAULT_BUDGET_USD) + p.add_argument( + "--reserve-usd", + type=float, + default=1.0, + help="stop starting households when this much budget remains", + ) + if name == "score": + p.add_argument("--all-records", action="store_true") + if name == "diagnose": + p.add_argument("--no-engine-defaults", action="store_true") + args = parser.parse_args(argv) + { + "oracle": command_oracle, + "run": command_run, + "score": command_score, + "diagnose": command_diagnose, + }[args.command](args) + + +if __name__ == "__main__": + main() diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py new file mode 100644 index 00000000..21e33fbd --- /dev/null +++ b/tests/test_tool_pilot_argon.py @@ -0,0 +1,694 @@ +"""The Argon tool pilot's validator, conversation loop, scoring and diagnosis. + +Every model call is faked. The registry is a small stand-in for the engine's, +so the validator tests need no PolicyEngine system; the tests that build the +real tool system are marked slow. +""" + +from __future__ import annotations + +import importlib.metadata +import json +import math +import sys +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest +from hypothesis import given, settings +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 + +from policybench.scenarios import Person, Scenario # noqa: E402 + +V = pilot.VariableInfo +REGISTRY = { + "age": V("person", "int", "age", "year"), + "employment_income": V("person", "float", "employment income", "year"), + "is_disabled": V("person", "bool", "is disabled", "year"), + "is_medicaid_eligible": V("person", "bool", "Medicaid eligible", "year"), + "snap": V("spm_unit", "float", "SNAP", "month"), + "takes_up_snap_if_eligible": V("spm_unit", "bool", "takes up SNAP", "year"), + "income_tax_before_refundable_credits": V( + "tax_unit", "float", "income tax before refundable credits", "year" + ), + "free_school_meals": V("spm_unit", "float", "free school meals", "year"), + "state_code": V("household", "Enum", "state", "year", ("CA", "CO", "TX")), + "filing_status": V( + "tax_unit", "Enum", "filing status", "year", ("SINGLE", "JOINT") + ), +} + + +def situation(**overrides): + base = { + "people": { + "head": {"age": {"2026": 40}, "employment_income": 20000}, + "kid": {"age": {"2026": 4}}, + }, + "spm_units": {"spm": {"members": ["head", "kid"]}}, + "households": {"h": {"members": ["head", "kid"], "state_code": "CO"}}, + } + base.update(overrides) + return base + + +# --------------------------------------------------------------------------- +# The validator. +# --------------------------------------------------------------------------- + + +def test_accepts_a_household_and_yearizes_bare_values(): + request, error = pilot.normalize_request( + {"situation": json.dumps(situation()), "variables": ["snap"]}, REGISTRY + ) + assert error is None + people = request["situation"]["people"] + assert people["head"]["employment_income"] == {"2026": 20000} + assert request["situation"]["households"]["h"]["state_code"] == {"2026": "CO"} + assert request["variables"] == ["snap"] + assert request["period"] == "2026" + + +def test_a_bare_value_takes_the_calls_period(): + request, error = pilot.normalize_request( + {"situation": situation(), "variables": ["snap"], "period": "2026-03"}, + REGISTRY, + ) + assert error is None + assert request["situation"]["people"]["head"]["employment_income"] == { + "2026-03": 20000 + } + + +@pytest.mark.parametrize( + "args, fragment", + [ + ({"situation": "{not json", "variables": ["snap"]}, "not valid JSON"), + ({"situation": situation(), "variables": ["snapp"]}, "Close matches: snap"), + ( + {"situation": situation(axes=[[{"name": "age"}]]), "variables": ["snap"]}, + "remove axes", + ), + ( + { + "situation": situation( + people={"head": {"snap": {"2026": 1}}}, + spm_units={"s": {"members": ["head"]}}, + households={"h": {"members": ["head"]}}, + ), + "variables": ["snap"], + }, + "snap is a spm_unit variable; set it under spm_units", + ), + ( + { + "situation": situation(spm_units={"s": {"members": ["ghost"]}}), + "variables": ["snap"], + }, + "member 'ghost' is not in people", + ), + ( + { + "situation": situation( + spm_units={ + "a": {"members": ["head"]}, + "b": {"members": ["head", "kid"]}, + } + ), + "variables": ["snap"], + }, + "head is already in spm_units.a", + ), + ( + { + "situation": situation(people={"head": {"age": {"26": 40}}}), + "variables": ["age"], + }, + "period '26'", + ), + ( + { + "situation": situation(people={"head": {"age": {"2026": [40]}}}), + "variables": ["age"], + }, + "must be a number", + ), + ( + { + "situation": situation( + people={"head": {"employment_income": {"2026": float("nan")}}} + ), + "variables": ["age"], + }, + "not a finite number", + ), + ( + { + "situation": situation( + people={"head": {"employment_income": {"2026": True}}} + ), + "variables": ["age"], + }, + "takes a number", + ), + ( + { + "situation": situation( + households={"h": {"members": ["head", "kid"], "state_code": "ZZ"}} + ), + "variables": ["snap"], + }, + "takes one of: CA, CO, TX", + ), + ({"situation": situation(), "variables": []}, "non-empty list"), + ( + {"situation": situation(), "variables": ["snap"], "period": "next year"}, + "must be a year", + ), + ({"situation": situation(), "variables": ["snap"], "code": "x"}, "unknown"), + ({"situation": {"people": {}}, "variables": ["snap"]}, "non-empty object"), + ], +) +def test_rejects_with_a_specific_problem(args, fragment): + request, error = pilot.normalize_request(args, REGISTRY) + assert request is None + assert error["error"] == "validation" + assert any(fragment in problem for problem in error["problems"]), error + + +def test_too_many_variables_are_refused(): + names = [f"v{i}" for i in range(pilot.MAX_VARIABLES + 1)] + registry = {**REGISTRY, **{n: V("person", "float") for n in names}} + _, error = pilot.normalize_request( + {"situation": situation(), "variables": names}, registry + ) + assert any("at most" in p for p in error["problems"]) + + +json_scalars = st.none() | st.booleans() | st.integers() | st.floats() | st.text() +json_values = st.recursive( + json_scalars, + lambda children: ( + st.lists(children, max_size=4) + | st.dictionaries( + st.sampled_from([*REGISTRY, "members", "people", "2026", "x", "axes"]), + children, + max_size=4, + ) + ), + max_leaves=25, +) + + +@settings(max_examples=300, deadline=None) +@given(args=json_values) +def test_validator_never_raises_and_accepts_only_known_names(args): + """Invariant: any input gets exactly one of a request or an error, and an + accepted request names only registry variables on their own entities.""" + request, error = pilot.normalize_request(args, REGISTRY) + assert (request is None) != (error is None) + if request is None: + return + for name in request["variables"]: + assert name in REGISTRY + for person in request["situation"]["people"].values(): + for name, periods in person.items(): + assert REGISTRY[name].entity == "person" + assert all(pilot.PERIOD_RE.match(k) for k in periods) + for group, entity in pilot.GROUPS.items(): + for unit in request["situation"].get(group, {}).values(): + for name in unit: + assert name == "members" or REGISTRY[name].entity == entity + + +@settings(max_examples=200, deadline=None) +@given( + income=st.floats(min_value=0, max_value=1e7, allow_nan=False), + age=st.integers(min_value=0, max_value=110), + period=st.sampled_from(["2026", "2025", "2026-01", "2026-12"]), +) +def test_valid_households_round_trip(income, age, period): + sit = situation(people={"head": {"age": age, "employment_income": income}}) + sit["spm_units"] = {"spm": {"members": ["head"]}} + sit["households"] = {"h": {"members": ["head"], "state_code": "TX"}} + request, error = pilot.normalize_request( + {"situation": sit, "variables": ["snap"], "period": period}, REGISTRY + ) + assert error is None + assert request["situation"]["people"]["head"] == { + "age": {period: age}, + "employment_income": {period: income}, + } + + +def test_execute_request_reads_entity_names_and_reports_engine_errors(): + class FakeArray(list): + pass + + class FakeSimulation: + def __init__(self, tax_benefit_system, situation): + if "bad" in situation["people"]: + raise ValueError("bad person") + self.populations = { + "person": SimpleNamespace(ids=list(situation["people"])), + "spm_unit": SimpleNamespace(ids=["spm"]), + } + + def calculate(self, name, period): + if name == "snap": + return FakeArray([1234.567]) + return FakeArray([True, False]) + + request = { + "situation": situation(), + "variables": ["snap", "is_medicaid_eligible"], + "period": "2026", + } + result = pilot.execute_request(None, REGISTRY, request, FakeSimulation) + assert result["results"]["snap"]["values"] == {"spm": 1234.57} + assert result["results"]["is_medicaid_eligible"]["values"] == { + "head": True, + "kid": False, + } + bad = { + "situation": {"people": {"bad": {}}}, + "variables": ["snap"], + "period": "2026", + } + error = pilot.execute_request(None, REGISTRY, bad, FakeSimulation) + assert error["error"] == "engine" and "bad person" in error["message"] + + +# --------------------------------------------------------------------------- +# The conversation loop, with a faked model. +# --------------------------------------------------------------------------- + + +def make_scenario() -> Scenario: + return Scenario( + id="scenario_t", + state="CO", + filing_status="single", + adults=[Person(name="head", age=40, employment_income=20000)], + children=[Person(name="child1", age=4, employment_income=0)], + ) + + +VARIABLES = ["snap", "federal_income_tax_before_refundable_credits"] + + +def tool_call(name, arguments, call_id): + return SimpleNamespace( + id=call_id, + type="function", + function=SimpleNamespace(name=name, arguments=json.dumps(arguments)), + ) + + +def response(calls, prompt=1000, completion=200, reasoning=50, content=None): + message = SimpleNamespace( + role="assistant", content=content, tool_calls=calls, function_call=None + ) + usage = SimpleNamespace( + prompt_tokens=prompt, + completion_tokens=completion, + total_tokens=prompt + completion, + reasoning_tokens=reasoning, + prompt_tokens_details=SimpleNamespace(cached_tokens=0), + completion_tokens_details=None, + ) + return SimpleNamespace( + choices=[SimpleNamespace(message=message, finish_reason="tool_calls")], + usage=usage, + model="gemini-4-argon", + ) + + +def submit(values: dict, call_id="s1"): + outputs = { + k: {"value": v, "explanation": f"From the tool. value = {v}"} + for k, v in values.items() + } + return tool_call("submit_outputs", {"outputs": outputs}, call_id) + + +def calculate(call_id="c1"): + return tool_call( + pilot.TOOL_NAME, + {"situation": json.dumps(situation()), "variables": ["snap"]}, + call_id, + ) + + +class FakeModel: + def __init__(self, responses): + self.responses = list(responses) + self.requests = [] + + def __call__(self, **kwargs): + self.requests.append( + { + **kwargs, + "messages": json.loads(json.dumps(kwargs["messages"], default=str)), + } + ) + return self.responses.pop(0) + + +def fake_tool(args): + return { + "results": { + "snap": {"entity": "spm_unit", "period": "2026", "values": {"spm": 2553.0}} + } + }, 0.25 + + +def test_a_tool_turn_then_the_answer_tool(): + model = FakeModel( + [response([calculate()]), response([submit({"snap": 2553, VARIABLES[1]: 0})])] + ) + spend = pilot.Spend(50.0) + record = pilot.run_household(make_scenario(), VARIABLES, fake_tool, model, spend) + assert record["status"] == "complete" + assert record["predictions"] == {"snap": 2553.0, VARIABLES[1]: 0.0} + assert record["tool_turns"] == 1 and not record["cap_hit"] + first, second = model.requests + assert first["tool_choice"] == "required" + assert [t["function"]["name"] for t in first["tools"]] == [ + pilot.TOOL_NAME, + "submit_outputs", + ] + assert first["messages"][0]["content"].endswith(pilot.TOOL_NOTICE) + tool_messages = [m for m in second["messages"] if m["role"] == "tool"] + assert json.loads(tool_messages[0]["content"])["results"]["snap"]["values"] == { + "spm": 2553.0 + } + totals = record["totals"] + assert totals["model_turns"] == 2 and totals["tool_calls"] == 1 + assert totals["tool_seconds"] == pytest.approx(0.25) + # Invariant: the household's cost is the sum of its turns' costs, and the + # shared spend counter saw exactly that. + assert totals["cost_usd"] == pytest.approx( + sum(t["total_cost_usd"] for t in record["turns"]) + ) + assert spend.spent_usd == pytest.approx(totals["cost_usd"]) + # $2 / $10 per 1M at 1,000 input and 200 output tokens a turn. + assert totals["cost_usd"] == pytest.approx(2 * (1000 * 2 + 200 * 10) / 1e6) + + +def test_the_prompt_is_the_board_prompt_plus_the_notice(): + from policybench.eval_no_tools import _chat_completion_request_kwargs + + scenario = make_scenario() + board_messages, _ = _chat_completion_request_kwargs( + scenario, VARIABLES, "gemini/gemini-4-argon" + ) + messages, kwargs = pilot.build_initial_request(scenario, VARIABLES) + assert messages[0]["content"] == board_messages[0]["content"] + pilot.TOOL_NOTICE + assert "caching" not in kwargs + + +def test_the_cap_forces_the_answer_tool(): + cap = 3 + model = FakeModel( + [response([calculate(f"c{i}")]) for i in range(cap)] + + [response([submit({"snap": 1, VARIABLES[1]: 2})])] + ) + record = pilot.run_household( + make_scenario(), VARIABLES, fake_tool, model, pilot.Spend(50.0), turn_cap=cap + ) + assert record["status"] == "complete" and record["cap_hit"] + assert [r["tool_choice"] for r in model.requests[:cap]] == ["required"] * cap + assert model.requests[cap]["tool_choice"] == { + "type": "function", + "function": {"name": "submit_outputs"}, + } + + +def test_extra_calculations_in_one_turn_are_not_run(): + calls = [calculate(f"c{i}") for i in range(pilot.MAX_CALCULATIONS_PER_TURN + 2)] + ran = [] + + def counting_tool(args): + ran.append(args) + return fake_tool(args) + + model = FakeModel( + [response(calls), response([submit({"snap": 1, VARIABLES[1]: 2})])] + ) + record = pilot.run_household( + make_scenario(), VARIABLES, counting_tool, model, pilot.Spend(50.0) + ) + assert len(ran) == pilot.MAX_CALCULATIONS_PER_TURN + results = [c["result"] for c in record["turns"][0]["calls"]] + assert sum(r.get("error") == "not_run" for r in results) == 2 + + +def test_a_missing_output_gets_a_repair_turn_for_it_alone(): + model = FakeModel( + [response([submit({"snap": 10})]), response([submit({VARIABLES[1]: 5}, "s2")])] + ) + record = pilot.run_household( + make_scenario(), VARIABLES, fake_tool, model, pilot.Spend(50.0) + ) + assert record["status"] == "complete" and record["repair_rounds"] == 1 + repair = model.requests[1] + answer_tool = repair["tools"][1]["function"]["parameters"] + assert list(answer_tool["properties"]["outputs"]["properties"]) == [VARIABLES[1]] + assert record["predictions"] == {"snap": 10.0, VARIABLES[1]: 5.0} + + +def test_an_exhausted_budget_stops_before_calling(): + model = FakeModel([]) + record = pilot.run_household( + make_scenario(), VARIABLES, fake_tool, model, pilot.Spend(1.0, spent_usd=2.0) + ) + assert record["status"] == "budget_stop" and model.requests == [] + + +def test_a_provider_error_ends_the_household(monkeypatch): + monkeypatch.setattr(pilot, "MAX_PROVIDER_ATTEMPTS", 1) + + def failing(**kwargs): + raise RuntimeError("boom") + + record = pilot.run_household( + make_scenario(), VARIABLES, fake_tool, failing, pilot.Spend(50.0) + ) + assert record["status"] == "provider_error" and "boom" in record["error"] + + +# --------------------------------------------------------------------------- +# Scoring and diagnosis. +# --------------------------------------------------------------------------- + + +def test_hit_is_the_boards_exact_rule(): + assert pilot.hit("snap", 100.9, 100.0) + assert not pilot.hit("snap", 101.5, 100.0) + assert pilot.hit("head_medicaid_eligible", 1, 1.0) + assert not pilot.hit("head_medicaid_eligible", 0, 1.0) + assert not pilot.hit("snap", None, 1.0) + assert not pilot.hit("snap", float("nan"), 1.0) + + +def call_record(values: dict, period="2026", people=("head", "kid")): + sit = situation() + sit["people"] = {p: {"age": 30} for p in people} + return { + "name": pilot.TOOL_NAME, + "arguments": json.dumps( + {"situation": json.dumps(sit), "variables": list(values)} + ), + "result": { + "results": { + name: {"entity": "x", "period": period, "values": vals} + for name, vals in values.items() + } + }, + } + + +def household(calls, cap_hit=False): + return {"turns": [{"calls": calls}], "cap_hit": cap_hit} + + +def test_tool_value_maps_people_by_position_and_aggregates_groups(): + scenario = make_scenario() + call = call_record( + { + "is_medicaid_eligible": {"head": False, "kid": True}, + "free_school_meals": {"spm": 0.0, "spm2": 12.0}, + } + ) + assert ( + pilot.tool_value_for_output(call, scenario, "child1_medicaid_eligible") == 1.0 + ) + assert pilot.tool_value_for_output(call, scenario, "head_medicaid_eligible") == 0.0 + assert ( + pilot.tool_value_for_output(call, scenario, "free_school_meals_eligible") == 1.0 + ) + monthly = call_record({"snap": {"spm": 200.0}}, period="2026-01") + assert pilot.tool_value_for_output(monthly, scenario, "snap") is None + + +@pytest.mark.parametrize( + "calls, submitted, category", + [ + ([], 100.0, "tool_error_or_cap"), + ([call_record({"snap": {"spm": 100.0}})], 100.0, "wrong_input_mapping"), + ([call_record({"snap": {"spm": 100.0}})], 250.0, "answer_differs_from_tool"), + ( + [call_record({"snap": {"spm": 9.0}}, period="2026-01")], + 9.0, + "wrong_variable_or_period", + ), + ([call_record({"tanf": {"spm": 9.0}})], 9.0, "wrong_variable_or_period"), + ([call_record({"snap": {"spm": 100.0}})], None, "no_answer"), + ], +) +def test_each_miss_gets_its_category(calls, submitted, category): + verdict = pilot.classify_miss( + household(calls), make_scenario(), "snap", submitted, 400.0 + ) + assert verdict["category"] == category + + +@settings(max_examples=200, deadline=None) +@given( + tool_values=st.lists( + st.one_of(st.none(), st.floats(min_value=-1e5, max_value=1e5)), max_size=4 + ), + periods=st.lists( + st.sampled_from(["2026", "2026-01", "2025"]), min_size=4, max_size=4 + ), + submitted=st.one_of(st.none(), st.floats(min_value=-1e5, max_value=1e5)), +) +def test_every_miss_lands_in_exactly_one_category(tool_values, periods, submitted): + """Invariant: classify_miss returns one of MISS_CATEGORIES for any transcript.""" + calls = [ + call_record( + {"snap": {"spm": v}} if v is not None else {"tanf": {"spm": 1.0}}, p + ) + for v, p in zip(tool_values, periods) + ] + verdict = pilot.classify_miss( + household(calls), make_scenario(), "snap", submitted, 1e9 + ) + assert verdict["category"] in pilot.MISS_CATEGORIES + + +def test_situation_diff_matches_people_by_position(): + scenario = make_scenario() + reference = scenario.to_pe_household() + model = json.loads(json.dumps(reference)) + model["people"] = { + "me": dict(model["people"]["head"]), + "child": dict(model["people"]["child1"]), + } + assert pilot.situation_diff(model, reference, scenario) == [] + model["people"]["me"]["employment_income"] = {"2026": 21000} + diff = pilot.situation_diff(model, reference, scenario) + assert [(d["variable"], d["model"]) for d in diff] == [ + ("employment_income", {"2026": 21000}) + ] + + +def test_default_only_differences_are_dropped(): + defaults = { + "takes_up_snap_if_eligible": True, + "would_file_taxes_voluntarily": False, + } + assert pilot._is_default_only( + { + "variable": "takes_up_snap_if_eligible", + "reference": {"2026": True}, + "model": {}, + }, + defaults, + ) + assert not pilot._is_default_only( + { + "variable": "would_file_taxes_voluntarily", + "reference": {"2026": True}, + "model": {}, + }, + defaults, + ) + + +def test_arm1_rescoring_reproduces_the_published_board_row(): + """Differential: the board's analysis on the frozen predictions gives the + published weighted_exact for Argon (summary_by_model.csv).""" + from policybench.analysis import analyze_no_tools + + board = pilot.load_board() + arm1 = pilot.arm1_predictions() + analysis = analyze_no_tools( + board["scored"], arm1, scenarios=board["scenarios_frame"] + ) + row = analysis["model_summary"].set_index("model").loc[pilot.ARM1_LABEL] + published = pd.read_csv(pilot.SNAPSHOT / "analysis/summary_by_model.csv").set_index( + "model" + ) + assert row["weighted_exact"] == pytest.approx( + published.loc[pilot.ARM1_LABEL, "weighted_exact"], abs=1e-12 + ) + assert len(board["scored"]) == 1935 + + +def test_the_fix_module_is_the_pinned_one(): + assert pilot._sha256(pilot.FIX_MODULE) == pilot.FIX_MODULE_SHA256 + meta = json.loads((pilot.SNAPSHOT / "reference_outputs.csv.meta.json").read_text()) + revision = [ + r + for r in meta["revisions"] + if r.get("engine_version") == "policyengine-us 2.41.0" + ][0] + pins = {m["module"]: m["sha256"] for m in revision["fix_modules"]} + assert pins["latest_final.py"] == pilot.FIX_MODULE_SHA256 + + +def _engine_is_reference() -> bool: + try: + return importlib.metadata.version("policyengine-us") == pilot.REFERENCE_ENGINE + except importlib.metadata.PackageNotFoundError: + return False + + +@pytest.mark.slow +@pytest.mark.skipif(not _engine_is_reference(), reason="needs policyengine-us 2.41.0") +def test_the_real_tool_runs_a_household_through_the_reference_engine(): + system = pilot.build_tool_system(reform=True) + registry = pilot.build_registry(system) + request, error = pilot.normalize_request( + { + "situation": json.dumps( + { + "people": { + "you": { + "age": {"2026": 30}, + "employment_income": {"2026": 20000}, + }, + "kid": {"age": {"2026": 3}}, + }, + "households": { + "h": {"members": ["you", "kid"], "state_code": "CO"} + }, + } + ), + "variables": ["snap", "is_medicaid_eligible", "filing_status"], + }, + registry, + ) + assert error is None + result = pilot.execute_request(system, registry, request)["results"] + assert result["snap"]["entity"] == "spm_unit" + assert set(result["is_medicaid_eligible"]["values"]) == {"you", "kid"} + assert isinstance(result["filing_status"]["values"]["tax_unit"], str) + assert math.isfinite(result["snap"]["values"]["spm_unit"]) From 2f8023b53b84326d736aad7ca1d9eac3a2a7723b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 07:56:44 -0400 Subject: [PATCH 2/6] Attribute each tool-arm miss to its inputs and bootstrap the household difference diagnose now applies each differing input as the reference sets it and recomputes the output, so a mapping miss names the input that causes it; it reads payroll_tax from employee_payroll_tax too (spm_unit_payroll_tax sums it over the SPM unit's tax units). score adds household weighted exact scores, checked against analyze_no_tools before use, a paired household bootstrap, and hit/miss transitions by program. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 309 +++++++++++++++++++++++++++++++-- tests/test_tool_pilot_argon.py | 108 ++++++++++++ 2 files changed, 403 insertions(+), 14 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index f3ec8b32..b0e97a7f 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1414,6 +1414,38 @@ def command_score(args) -> None: per_program.to_csv(out / "per_program.csv", index=False) print(per_program.to_string(index=False)) + # Household scores, checked against the board's headline before use. + both = pd.concat([arm1, arm2], ignore_index=True) + by_household = household_weighted_exact(scored, both, scenarios) + means = by_household.groupby("model")["weighted_exact"].mean() + published = summary.set_index("model")["weighted_exact"] + for model in (ARM1_LABEL, ARM2_LABEL): + if abs(means[model] - published[model]) > 1e-9: + raise SystemExit( + f"household scores for {model} average {means[model]}, not the " + f"board's {published[model]}" + ) + wide = by_household.pivot( + index="scenario_id", columns="model", values="weighted_exact" + ) + wide.to_csv(out / "household_scores.csv") + bootstrap = paired_bootstrap(wide[ARM1_LABEL], wide[ARM2_LABEL]) + print(json.dumps(bootstrap, indent=2)) + + flips = ( + rows.assign( + transition=[ + ("hit" if a else "miss") + "->" + ("hit" if b else "miss") + for a, b in zip(rows.arm1_hit, rows.arm2_hit) + ] + ) + .groupby(["program", "transition"]) + .size() + .unstack(fill_value=0) + ) + flips.to_csv(out / "flips_by_program.csv") + print(flips.to_string()) + usage = usage_by_household(records, arm1, households) usage.to_csv(out / "usage_by_household.csv", index=False) usage_summary = summarize_usage(usage) @@ -1429,6 +1461,10 @@ def command_score(args) -> None: "arm1": int(rows.arm1_hit.sum()), "arm2": int(rows.arm2_hit.sum()), }, + "transitions": { + k: int(v) for k, v in flips.sum(axis=0).to_dict().items() + }, + "household_bootstrap": bootstrap, "usage": usage_summary, }, indent=2, @@ -1437,6 +1473,90 @@ def command_score(args) -> None: ) +def household_weighted_exact( + scored: pd.DataFrame, predictions: pd.DataFrame, scenarios_frame: pd.DataFrame +) -> pd.DataFrame: + """Each household's weighted exact score. + + The steps of policybench.analysis.weighted_hit_rate_scores_by_model up to + its household table, which it then averages with equal household weight; + command_score checks that the mean of these equals its weighted_exact. + """ + from policybench.analysis import ( + _country_from_scenarios, + _drop_zero_weight_households, + _row_weights_for_ground_truth, + bounded_global_variable_weights, + row_hit_scores, + ) + + market_income = dict( + zip( + scenarios_frame["scenario_id"].astype(str), + pd.to_numeric(scenarios_frame["total_income"], errors="coerce").fillna(0.0), + ) + ) + weights = bounded_global_variable_weights( + scored, market_income, country=_country_from_scenarios(scenarios_frame) + ) + row_weights = _row_weights_for_ground_truth(scored, weights) + models = sorted(predictions["model"].dropna().unique()) + grid = ( + scored.assign(_k=1) + .merge(pd.DataFrame({"model": models, "_k": 1}), on="_k") + .drop(columns="_k") + ) + merged = grid.merge( + predictions[["model", "scenario_id", "variable", "prediction"]], + on=["model", "scenario_id", "variable"], + how="left", + ) + merged["exact"] = [ + row_hit_scores(v, t, p)["exact"] + for v, t, p in zip(merged["variable"], merged["value"], merged["prediction"]) + ] + merged = merged.merge(row_weights, on=["scenario_id", "variable"], how="left") + merged["weight"] = merged["weight"].fillna(0.0) + merged = _drop_zero_weight_households(merged, row_weights) + merged["weighted_exact"] = merged["exact"] * merged["weight"] + return ( + merged.groupby(["model", "scenario_id"])["weighted_exact"].sum().reset_index() + ) + + +def paired_bootstrap( + first: pd.Series, second: pd.Series, draws: int = 10_000, seed: int = 20261014 +) -> dict: + """Household-resampled 95% intervals for two arms' means and their difference.""" + import numpy as np + + a = first.to_numpy(dtype=float) + b = second.to_numpy(dtype=float) + rng = np.random.default_rng(seed) + index = rng.integers(0, len(a), size=(draws, len(a))) + mean_a = a[index].mean(axis=1) + mean_b = b[index].mean(axis=1) + + def interval(values): + low, high = np.percentile(values, [2.5, 97.5]) + return [float(low), float(high)] + + return { + "households": int(len(a)), + "draws": draws, + "seed": seed, + "arm1_mean": float(a.mean()), + "arm2_mean": float(b.mean()), + "difference": float(b.mean() - a.mean()), + "arm1_interval": interval(mean_a), + "arm2_interval": interval(mean_b), + "difference_interval": interval(mean_b - mean_a), + "households_better": int((b > a + 1e-12).sum()), + "households_worse": int((b < a - 1e-12).sum()), + "households_same": int((abs(b - a) <= 1e-12).sum()), + } + + def usage_by_household(records, arm1: pd.DataFrame, households) -> pd.DataFrame: a1 = arm1.copy() for column in ( @@ -1509,6 +1629,12 @@ def stats(column): # Miss diagnosis. # --------------------------------------------------------------------------- +# Engine variables equal to an output's own variable on these households. In +# policyengine-us 2.41.0 spm_unit_payroll_tax is sum_contained_tax_units of +# employee_payroll_tax, and an SPM unit's total is the same whichever of the +# two the model asks for. +EQUIVALENT_ENGINE_VARIABLES = {"payroll_tax": ("employee_payroll_tax",)} + MISS_CATEGORIES = ( "wrong_input_mapping", "wrong_variable_or_period", @@ -1542,6 +1668,15 @@ def _person_names(arguments) -> list[str]: return [] +def engine_variables_for_output(output: str) -> tuple[str, ...]: + from policybench.ground_truth import _pe_variable_for_output + + return ( + _pe_variable_for_output(output, "us"), + *EQUIVALENT_ENGINE_VARIABLES.get(output, ()), + ) + + def tool_value_for_output(call: dict, scenario, output: str): """The value one successful call returned for a benchmark output, or None. @@ -1550,16 +1685,20 @@ def tool_value_for_output(call: dict, scenario, output: str): position (the prompt lists people in scenario order); group outputs take the board's aggregation over the returned units. """ - from policybench.ground_truth import ( - _aggregation_for_output, - _pe_variable_for_output, - ) - - pe_variable = _pe_variable_for_output(output, "us") - entry = call["result"]["results"].get(pe_variable) - if not isinstance(entry, dict) or "values" not in entry: - return None - if str(entry.get("period")) != str(YEAR): + from policybench.ground_truth import _aggregation_for_output + + results = call["result"]["results"] + entry = None + for candidate in engine_variables_for_output(output): + found = results.get(candidate) + if ( + isinstance(found, dict) + and "values" in found + and str(found.get("period")) == str(YEAR) + ): + entry = found + break + if entry is None: return None values = entry["values"] aggregation = _aggregation_for_output(output, "us") @@ -1586,6 +1725,25 @@ def tool_value_for_output(call: dict, scenario, output: str): return None +def submitted_matches(calls: list[dict], submitted, limit: int = 5) -> list[str]: + """Tool results within $1 of a submitted amount, as variable[member]@period.""" + if submitted is None: + return [] + matches = [] + for call in calls: + for name, entry in call["result"]["results"].items(): + if not isinstance(entry, dict): + continue + for member, value in (entry.get("values") or {}).items(): + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + if abs(float(value) - float(submitted)) <= 1.0: + label = f"{name}[{member}]@{entry.get('period')}" + if label not in matches: + matches.append(label) + return matches[:limit] + + def classify_miss(record: dict, scenario, output: str, submitted, reference) -> dict: """One of MISS_CATEGORIES, with the evidence behind it. @@ -1607,13 +1765,17 @@ def classify_miss(record: dict, scenario, output: str, submitted, reference) -> c for c in calls if tool_value_for_output(c, scenario, output) is not None ] if not computed: + names = engine_variables_for_output(output) any_period = [ c for c in calls - if pe_variable in c["result"]["results"] - and isinstance(c["result"]["results"][pe_variable], dict) - and "values" in c["result"]["results"][pe_variable] + if any( + isinstance(c["result"]["results"].get(name), dict) + and "values" in c["result"]["results"][name] + for name in names + ) ] + evidence["submitted_matches"] = submitted_matches(calls, submitted) if not calls: return { "category": "tool_error_or_cap", @@ -1718,6 +1880,104 @@ def _canon(value): return value +def output_value(system, situation: dict, scenario, output: str) -> float | None: + """One benchmark output from a situation, people read by position.""" + from policyengine_us import Simulation + + from policybench.ground_truth import ( + _aggregation_for_output, + _pe_variable_for_output, + ) + + try: + simulation = Simulation( + tax_benefit_system=system, situation=copy.deepcopy(situation) + ) + values = simulation.calculate(_pe_variable_for_output(output, "us"), YEAR) + except Exception: + return None + aggregation = _aggregation_for_output(output, "us") + if aggregation == "person": + from policybench.ground_truth import _person_index_for_output + + index = _person_index_for_output(scenario, output) + if index is None or index >= len(values): + return None + return float(values[index]) + total = float(values.sum()) + return float(total > 0) if aggregation == "any_positive" else total + + +def apply_reference_input(model: dict, reference: dict, diff: dict) -> dict | None: + """The model's situation with one differing input set as the reference sets it. + + People are matched by position. A group input goes on the model's first + unit of that group (a new unit holding everyone if the model left the group + out); a unit-count difference replaces the model's group with the + reference's, members renamed to the model's people. Returns None for a + difference that cannot be applied (a different number of people). + """ + situation = copy.deepcopy(model) + model_names = list(situation.get("people", {})) + reference_names = list(reference["people"]) + if len(model_names) != len(reference_names): + return None + rename = dict(zip(reference_names, model_names)) + variable = diff["variable"] + if diff["entity"] == "person": + person = situation["people"][rename[diff["reference_unit"]]] + if diff["reference"] in ({}, None): + person.pop(variable, None) + else: + person[variable] = copy.deepcopy(diff["reference"]) + return situation + group = diff["entity"] + if variable == "": + situation[group] = { + name: { + **copy.deepcopy(unit), + "members": [rename[m] for m in unit["members"]], + } + for name, unit in reference[group].items() + } + return situation + units = situation.setdefault(group, {}) + if not units: + units[f"{group}_1"] = {"members": list(model_names)} + unit = next(iter(units.values())) + if diff["reference"] in ({}, None): + unit.pop(variable, None) + else: + unit[variable] = copy.deepcopy(diff["reference"]) + return situation + + +def attribute_miss( + system, model: dict, reference: dict, scenario, output: str, target, diffs +) -> dict: + """Which differing inputs, set as the reference sets them, fix the output. + + Each difference is applied alone; the ones that bring the output within + the board's tolerance of the reference are its single causes. When no + single input does, every applicable difference is applied together. + """ + causes = [] + for diff in diffs: + situation = apply_reference_input(model, reference, diff) + if situation is None: + continue + if hit(output, output_value(system, situation, scenario, output), target): + causes.append(f"{diff['entity']}.{diff['variable']}") + joint = None + if not causes and diffs: + situation = model + for diff in diffs: + applied = apply_reference_input(situation, reference, diff) + situation = applied if applied is not None else situation + joint = hit(output, output_value(system, situation, scenario, output), target) + return {"single_causes": causes, "all_differences_fix_it": joint} + + def command_diagnose(args) -> None: board = load_board() out = Path(args.out_dir) @@ -1736,6 +1996,12 @@ def command_diagnose(args) -> None: name: system.variables[name].default_value for name in system.variables } out_rows = [] + if misses.empty: + print("no arm-2 misses") + pd.DataFrame(columns=["scenario_id", "variable", "category"]).to_csv( + out / "misses.csv", index=False + ) + return for row in misses.itertuples(): record = records[row.scenario_id] scenario = board["scenarios"][row.scenario_id] @@ -1744,14 +2010,26 @@ def command_diagnose(args) -> None: record, scenario, row.variable, submitted, row.reference ) diff = [] + attribution = {} arguments = verdict.get("call_arguments") if arguments is not None and registry is not None: request, _ = normalize_request(arguments, registry) if request is not None: + reference_situation = scenario.to_pe_household() diff = situation_diff( - request["situation"], scenario.to_pe_household(), scenario + request["situation"], reference_situation, scenario ) diff = [d for d in diff if not _is_default_only(d, defaults)] + if verdict["category"] == "wrong_input_mapping": + attribution = attribute_miss( + system, + request["situation"], + reference_situation, + scenario, + row.variable, + row.reference, + diff, + ) out_rows.append( { "scenario_id": row.scenario_id, @@ -1767,6 +2045,9 @@ def command_diagnose(args) -> None: "tool_value": verdict.get("tool_value"), "status": record.get("status"), "cap_hit": record.get("cap_hit"), + "submitted_matches": json.dumps(verdict.get("submitted_matches")), + "single_causes": json.dumps(attribution.get("single_causes")), + "all_differences_fix_it": attribution.get("all_differences_fix_it"), "input_diff": json.dumps(diff, default=str), } ) diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index 21e33fbd..6cf21e5a 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -692,3 +692,111 @@ def test_the_real_tool_runs_a_household_through_the_reference_engine(): assert set(result["is_medicaid_eligible"]["values"]) == {"you", "kid"} assert isinstance(result["filing_status"]["values"]["tax_unit"], str) assert math.isfinite(result["snap"]["values"]["spm_unit"]) + + +def test_payroll_reads_the_equivalent_engine_variable(): + call = call_record({"employee_payroll_tax": {"tax_unit": 1530.0}}) + assert pilot.tool_value_for_output(call, make_scenario(), "payroll_tax") == 1530.0 + assert pilot.tool_value_for_output(call, make_scenario(), "snap") is None + + +def test_submitted_matches_names_the_tool_values_an_answer_copied(): + calls = [ + call_record({"state_income_tax": {"tax_unit": 812.4}, "snap": {"spm": 5.0}}) + ] + assert pilot.submitted_matches(calls, 812.0) == ["state_income_tax[tax_unit]@2026"] + assert pilot.submitted_matches(calls, None) == [] + + +def test_apply_reference_input_sets_removes_and_replaces_units(): + scenario = make_scenario() + reference = scenario.to_pe_household() + model = { + "people": {"me": {"age": {"2026": 40}}, "kid": {"age": {"2026": 4}}}, + "households": {"h": {"members": ["me", "kid"], "state_name": {"2026": "CO"}}}, + } + set_income = pilot.apply_reference_input( + model, + reference, + { + "entity": "person", + "reference_unit": "head", + "variable": "employment_income", + "reference": {"2026": 20000}, + "model": {}, + }, + ) + assert set_income["people"]["me"]["employment_income"] == {"2026": 20000} + assert "employment_income" not in model["people"]["me"] # the input is not mutated + dropped = pilot.apply_reference_input( + model, + reference, + { + "entity": "households", + "variable": "state_name", + "reference": {}, + "model": {}, + }, + ) + assert "state_name" not in dropped["households"]["h"] + units = pilot.apply_reference_input( + model, + reference, + {"entity": "marital_units", "variable": "", "reference": 2, "model": 0}, + ) + members = [u["members"] for u in units["marital_units"].values()] + assert sorted(members) == sorted( + [ + [{"head": "me", "child1": "kid"}[m] for m in u["members"]] + for u in reference["marital_units"].values() + ] + ) + assert ( + pilot.apply_reference_input( + {"people": {"me": {}}}, + reference, + { + "entity": "person", + "reference_unit": "head", + "variable": "age", + "reference": {}, + }, + ) + is None + ) + + +@settings(max_examples=200, deadline=None) +@given( + first=st.lists(st.floats(min_value=0, max_value=1), min_size=2, max_size=40), + shift=st.floats(min_value=-0.5, max_value=0.5), +) +def test_paired_bootstrap_intervals_bracket_the_means(first, shift): + """Invariant: each interval contains its point estimate, and identical arms + give a zero difference interval.""" + a = pd.Series(first) + b = (a + shift).clip(0, 1) + out = pilot.paired_bootstrap(a, b, draws=500) + low, high = out["difference_interval"] + assert low - 1e-9 <= out["difference"] <= high + 1e-9 + assert out["households_better"] + out["households_worse"] + out[ + "households_same" + ] == len(a) + same = pilot.paired_bootstrap(a, a, draws=200) + assert same["difference_interval"] == [0.0, 0.0] + + +def test_household_scores_average_to_the_boards_weighted_exact(): + """Differential: the household table reproduces the board's headline.""" + board = pilot.load_board() + arm1 = pilot.arm1_predictions() + table = pilot.household_weighted_exact( + board["scored"], arm1, board["scenarios_frame"] + ) + published = pd.read_csv(pilot.SNAPSHOT / "analysis/summary_by_model.csv").set_index( + "model" + ) + assert table["weighted_exact"].mean() == pytest.approx( + published.loc[pilot.ARM1_LABEL, "weighted_exact"], abs=1e-12 + ) + assert table["scenario_id"].nunique() == 100 From 435f581dd468841fee4c06d7a0d0dfc7c6a0fb30 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 08:46:36 -0400 Subject: [PATCH 3/6] Summarize tool use and compare each household's main situation with the reference score adds a tool-use summary (first-call success, validation problems by kind, the names the model guessed wrong). diagnose writes situation_fidelity.csv from each household's main situation, the successful call that computed the most of its outputs, since a model often ends with a narrow probe on a skeleton situation. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 192 ++++++++++++++++++++++++++++++--- tests/test_tool_pilot_argon.py | 40 +++++++ 2 files changed, 219 insertions(+), 13 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index b0e97a7f..61a076f0 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -1450,6 +1450,8 @@ def command_score(args) -> None: usage.to_csv(out / "usage_by_household.csv", index=False) usage_summary = summarize_usage(usage) print(json.dumps(usage_summary, indent=2)) + tool_use = tool_use_summary({sid: records[sid] for sid in households}) + print(json.dumps(tool_use, indent=2)) (out / "score.json").write_text( json.dumps( { @@ -1465,6 +1467,7 @@ def command_score(args) -> None: k: int(v) for k, v in flips.sum(axis=0).to_dict().items() }, "household_bootstrap": bootstrap, + "tool_use": tool_use, "usage": usage_summary, }, indent=2, @@ -1603,6 +1606,74 @@ def usage_by_household(records, arm1: pd.DataFrame, households) -> pd.DataFrame: return frame +def tool_use_summary(records: Mapping[str, dict]) -> dict: + """How the model used the tool: calls, error kinds, and what it got wrong.""" + from collections import Counter + + problem_kinds: Counter = Counter() + unknown_names: Counter = Counter() + misplaced: Counter = Counter() + requested: Counter = Counter() + error_kinds: Counter = Counter() + first_call_ok = 0 + households = 0 + calls_before_first_success = [] + for record in records.values(): + calls = [ + c + for t in record.get("turns", []) + for c in t.get("calls", []) + if c.get("name") == TOOL_NAME + ] + if not calls: + continue + households += 1 + ok_index = next( + (i for i, c in enumerate(calls) if "results" in (c.get("result") or {})), + None, + ) + first_call_ok += ok_index == 0 + calls_before_first_success.append(ok_index) + for call in calls: + result = call.get("result") or {} + if "results" in result: + try: + args = json.loads(call["arguments"]) + requested.update(args.get("variables") or []) + except (TypeError, ValueError): + pass + continue + error_kinds[result.get("error")] += 1 + for problem in result.get("problems", []): + if "unknown variable" in problem: + problem_kinds["unknown variable"] += 1 + found = re.search(r"unknown variable '([^']+)'", problem) + if found: + unknown_names[found.group(1)] += 1 + elif " variable; set it under " in problem: + problem_kinds["variable on the wrong entity"] += 1 + misplaced[problem.split(": ", 1)[1].split(" is a ")[0]] += 1 + elif "period" in problem: + problem_kinds["bad period"] += 1 + elif "member" in problem: + problem_kinds["unit membership"] += 1 + elif "takes" in problem: + problem_kinds["wrong value type"] += 1 + else: + problem_kinds["other"] += 1 + never = sum(1 for i in calls_before_first_success if i is None) + return { + "households_using_tool": households, + "first_call_succeeded": int(first_call_ok), + "households_with_no_successful_call": int(never), + "failed_calls_by_kind": dict(error_kinds), + "validation_problems_by_kind": dict(problem_kinds.most_common()), + "most_common_unknown_names": dict(unknown_names.most_common(15)), + "most_common_misplaced_variables": dict(misplaced.most_common(10)), + "most_requested_variables": dict(requested.most_common(30)), + } + + def summarize_usage(usage: pd.DataFrame) -> dict: def stats(column): series = pd.to_numeric(usage[column], errors="coerce") @@ -1880,32 +1951,45 @@ def _canon(value): return value -def output_value(system, situation: dict, scenario, output: str) -> float | None: - """One benchmark output from a situation, people read by position.""" +def output_values(system, situation: dict, scenario, outputs) -> dict: + """Benchmark outputs from one simulation of a situation, people by position.""" from policyengine_us import Simulation from policybench.ground_truth import ( _aggregation_for_output, _pe_variable_for_output, + _person_index_for_output, ) try: simulation = Simulation( tax_benefit_system=system, situation=copy.deepcopy(situation) ) - values = simulation.calculate(_pe_variable_for_output(output, "us"), YEAR) except Exception: - return None - aggregation = _aggregation_for_output(output, "us") - if aggregation == "person": - from policybench.ground_truth import _person_index_for_output + return {output: None for output in outputs} + values = {} + for output in outputs: + try: + array = simulation.calculate(_pe_variable_for_output(output, "us"), YEAR) + except Exception: + values[output] = None + continue + aggregation = _aggregation_for_output(output, "us") + if aggregation == "person": + index = _person_index_for_output(scenario, output) + values[output] = ( + None if index is None or index >= len(array) else float(array[index]) + ) + else: + total = float(array.sum()) + values[output] = ( + float(total > 0) if aggregation == "any_positive" else total + ) + return values - index = _person_index_for_output(scenario, output) - if index is None or index >= len(values): - return None - return float(values[index]) - total = float(values.sum()) - return float(total > 0) if aggregation == "any_positive" else total + +def output_value(system, situation: dict, scenario, output: str) -> float | None: + return output_values(system, situation, scenario, [output])[output] def apply_reference_input(model: dict, reference: dict, diff: dict) -> dict | None: @@ -1978,6 +2062,86 @@ def attribute_miss( return {"single_causes": causes, "all_differences_fix_it": joint} +def main_call(record: dict, scenario, outputs: list[str]) -> dict | None: + """The successful call that computed the most of the household's outputs. + + A model often ends with a narrow probe on a skeleton situation (one + person, no state); the situation its answers rest on is the one that + computed the outputs. Ties go to the later call. + """ + best, best_count = None, -1 + for call in _successful_calls(record): + count = sum( + tool_value_for_output(call, scenario, output) is not None + for output in outputs + ) + if count >= best_count: + best, best_count = call, count + return best + + +def situation_fidelity(records, board, system, registry, defaults, out: Path) -> None: + """How each household's last computed situation differs from the reference's. + + Inputs left at the engine default on one side are dropped. A remaining + difference is reported with whether setting it as the reference does would + move any scored output beyond the board's tolerance, so a harmless + difference (another name for the same fact, a formula input set + explicitly) is told apart from one that matters. + """ + scored = board["scored"] + rows = [] + for sid, record in sorted(records.items()): + outputs = list(scored[scored.scenario_id == sid].variable) + call = main_call(record, board["scenarios"][sid], outputs) + if call is None: + rows.append({"scenario_id": sid, "differences": None}) + continue + request, _ = normalize_request(call["arguments"], registry) + if request is None: + rows.append({"scenario_id": sid, "differences": None}) + continue + scenario = board["scenarios"][sid] + reference = scenario.to_pe_household() + diffs = [ + d + for d in situation_diff(request["situation"], reference, scenario) + if not _is_default_only(d, defaults) + ] + base = output_values(system, request["situation"], scenario, outputs) + detail = [] + for diff in diffs: + applied = apply_reference_input(request["situation"], reference, diff) + moved = None + if applied is not None: + after = output_values(system, applied, scenario, outputs) + moved = [v for v in outputs if not hit(v, after[v], base[v])] + detail.append( + { + "input": f"{diff['entity']}.{diff['variable']}", + "reference": diff.get("reference"), + "model": diff.get("model"), + "outputs_it_moves": moved, + } + ) + rows.append( + { + "scenario_id": sid, + "differences": len(diffs), + "differences_that_move_an_output": sum( + 1 for d in detail if d["outputs_it_moves"] + ), + "detail": json.dumps(detail, default=str), + } + ) + frame = pd.DataFrame(rows) + frame.to_csv(out / "situation_fidelity.csv", index=False) + print( + f"situations matching the reference on every non-default input: " + f"{int((frame.differences == 0).sum())} of {len(frame)}" + ) + + def command_diagnose(args) -> None: board = load_board() out = Path(args.out_dir) @@ -1995,6 +2159,8 @@ def command_diagnose(args) -> None: defaults = { name: system.variables[name].default_value for name in system.variables } + if registry is not None: + situation_fidelity(records, board, system, registry, defaults, out) out_rows = [] if misses.empty: print("no arm-2 misses") diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index 6cf21e5a..3f0fe995 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -800,3 +800,43 @@ def test_household_scores_average_to_the_boards_weighted_exact(): published.loc[pilot.ARM1_LABEL, "weighted_exact"], abs=1e-12 ) assert table["scenario_id"].nunique() == 100 + + +def test_tool_use_summary_counts_failed_first_calls_and_unknown_names(): + bad = { + "name": pilot.TOOL_NAME, + "arguments": "{}", + "result": { + "error": "validation", + "problems": [ + "variables: unknown variable 'chip_eligible'. " + "Close matches: is_chip_eligible.", + "spm_units.s: auto_loan_balance is a household variable; " + "set it under households.", + ], + }, + } + good = call_record({"snap": {"spm": 1.0}}) + records = { + "a": {"turns": [{"calls": [bad]}, {"calls": [good]}]}, + "b": {"turns": [{"calls": [good]}]}, + "c": {"turns": [{"calls": []}]}, + } + summary = pilot.tool_use_summary(records) + assert summary["households_using_tool"] == 2 + assert summary["first_call_succeeded"] == 1 + assert summary["households_with_no_successful_call"] == 0 + assert summary["most_common_unknown_names"] == {"chip_eligible": 1} + assert summary["most_common_misplaced_variables"] == {"auto_loan_balance": 1} + assert summary["most_requested_variables"] == {"snap": 2} + + +def test_main_call_is_the_one_that_computed_the_most_outputs(): + full = call_record( + {"snap": {"spm": 1.0}, "income_tax_before_refundable_credits": {"tu": 2.0}} + ) + probe = call_record({"snap": {"spm": 3.0}}, people=("head",)) + record = household([full, probe]) + outputs = ["snap", "federal_income_tax_before_refundable_credits"] + assert pilot.main_call(record, make_scenario(), outputs) is full + assert pilot.main_call(household([]), make_scenario(), outputs) is None From 000d9059b0cab65155fe648e46b2b23a842a80cb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 12:49:20 -0400 Subject: [PATCH 4/6] Merge repairs as the board does and score every household Fixes from review of #229: - A repair now fills only cells still missing, with the board's _merge_repair_cells; a valid first value always stands. - score counts every household by default (--complete-only to opt in), and a resume reruns only infrastructure failures, keeping each replaced attempt's cost. - The validator's property test draws structured calls that reach acceptance (about 1 in 10), and checks it accepts at least 30 of 600. Also: every pinned fix module is verified before the reform loads; deeply nested JSON is refused, not raised; a failed household writes a harness_error record instead of ending the run; the pool restarts or reports lost workers; diagnose pins the engine; truncation retries are counted apart from model turns. New tests cover the merge, truncation retry, same-turn submit, turn limit, the litellm thought-signature round trip, pool timeout and crash restart, scoring without selection, and resume. Rescoring the pilot's run reproduces every reported number. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 225 +++++++++++++++++++----- tests/test_tool_pilot_argon.py | 302 +++++++++++++++++++++++++++++++-- 2 files changed, 469 insertions(+), 58 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 61a076f0..1bbd00e6 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -84,6 +84,9 @@ TOOL_TIMEOUT_SECONDS = 60 WORKER_START_TIMEOUT_SECONDS = 900 DEFAULT_BUDGET_USD = 50.0 +# Infrastructure failures a resume reruns. Every other status is the model's +# own outcome and stands, so no household is redrawn until it succeeds. +RERUNNABLE_STATUSES = frozenset({"provider_error", "budget_stop", "harness_error"}) MAX_PEOPLE = 16 MAX_UNITS_PER_GROUP = 16 @@ -301,10 +304,15 @@ def fail() -> tuple[None, dict]: } if isinstance(args, str): + if len(args) > 2 * MAX_SITUATION_CHARS: + problems.append("arguments are too large") + return fail() try: args = json.loads(args) - except json.JSONDecodeError as error: - problems.append(f"arguments are not valid JSON: {error}") + except (ValueError, RecursionError) as error: + problems.append( + f"arguments are not valid JSON: {error}"[:MAX_MESSAGE_CHARS] + ) return fail() if not isinstance(args, dict): problems.append("arguments must be an object") @@ -331,13 +339,17 @@ def fail() -> tuple[None, dict]: return fail() try: situation = json.loads(situation) - except json.JSONDecodeError as error: - problems.append(f"situation is not valid JSON: {error}") + except (ValueError, RecursionError) as error: + problems.append(f"situation is not valid JSON: {error}"[:MAX_MESSAGE_CHARS]) return fail() if not isinstance(situation, dict): problems.append("situation must be a JSON object") return fail() - if len(json.dumps(situation, default=str)) > MAX_SITUATION_CHARS: + try: + too_large = len(json.dumps(situation, default=str)) > MAX_SITUATION_CHARS + except (ValueError, RecursionError): + too_large = True + if too_large: problems.append("situation is too large") return fail() extra = sorted(set(situation) - {"people", *GROUPS}) @@ -381,6 +393,11 @@ def fail() -> tuple[None, dict]: normalized_units = {} for unit_name, unit in units.items(): where = f"{group}.{unit_name}" + if not isinstance(unit_name, str) or not unit_name or len(unit_name) > 64: + problems.append( + f"{group} unit name {unit_name!r} must be a short string" + ) + continue if not isinstance(unit, dict): problems.append(f"{where} must be an object with members") continue @@ -483,9 +500,44 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic return {"results": results} +def pinned_fix_modules(snapshot: Path = SNAPSHOT) -> list[dict]: + """The fix modules the 2.41.0 references were built with, as the 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"] + + +def verify_fix_pins(fixes_dir: Path = FIX_MODULE.parent) -> 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(): + 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}") + if problems: + raise SystemExit("; ".join(problems)) + + def load_fix_reform(path: Path = FIX_MODULE): - if path == FIX_MODULE and _sha256(path) != FIX_MODULE_SHA256: - raise SystemExit(f"{path} does not match the pinned sha256 {FIX_MODULE_SHA256}") + if path == FIX_MODULE: + verify_fix_pins() spec = importlib.util.spec_from_file_location("tool_pilot_latest_final", path) module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module @@ -530,14 +582,18 @@ def __init__( fix_path: Path = FIX_MODULE, timeout: float = TOOL_TIMEOUT_SECONDS, reform: bool = True, + worker_target: Callable | None = None, ): self._ctx = mp.get_context("spawn") self._fix_path = str(fix_path) self._reform = reform + self._target = worker_target or _worker_main self.timeout = timeout self._idle: queue.Queue = queue.Queue() + self._lock = threading.Lock() self.timeouts = 0 self.crashes = 0 + self.lost = 0 starting = [self._start() for _ in range(size)] for worker in starting: self._idle.put(self._await_ready(worker)) @@ -545,7 +601,7 @@ def __init__( def _start(self): parent, child = self._ctx.Pipe() process = self._ctx.Process( - target=_worker_main, + target=self._target, args=(child, self._fix_path, self._reform), daemon=True, ) @@ -563,27 +619,39 @@ def _await_ready(self, worker): return worker def call(self, args) -> tuple[dict, float]: - worker = self._idle.get() - process, conn = worker started = time.perf_counter() + try: + worker = self._idle.get(timeout=WORKER_START_TIMEOUT_SECONDS + self.timeout) + except queue.Empty: + return ( + {"error": "engine", "message": "no tool worker is available"}, + time.perf_counter() - started, + ) + process, conn = worker try: conn.send(args) if conn.poll(self.timeout): _, result = conn.recv() self._idle.put(worker) return result, time.perf_counter() - started - self.timeouts += 1 + with self._lock: + self.timeouts += 1 result = { "error": "timeout", "message": f"the calculation exceeded {self.timeout:.0f} seconds", } except (EOFError, BrokenPipeError, OSError) as error: - self.crashes += 1 + with self._lock: + self.crashes += 1 result = {"error": "engine", "message": f"tool worker failed: {error}"} elapsed = time.perf_counter() - started process.kill() process.join(timeout=10) - self._idle.put(self._await_ready(self._start())) + try: + self._idle.put(self._await_ready(self._start())) + except RuntimeError: + with self._lock: + self.lost += 1 return result, elapsed def close(self) -> None: @@ -766,6 +834,7 @@ def run_household( _build_answer_tool, _enforce_explanation_value_contract, _extract_usage_metadata, + _merge_repair_cells, _missing_explanations, _response_finish_reason, _serialize_response_payload, @@ -860,7 +929,12 @@ def run_household( record["error"] = f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS] break - message = response.choices[0].message + choices = getattr(response, "choices", None) or [] + if not choices: + record["status"] = "provider_error" + record["error"] = "response has no choices" + break + message = choices[0].message content = getattr(message, "content", None) calls = _tool_calls(message) raw = _serialize_response_payload( @@ -905,11 +979,21 @@ def run_household( got, got_explanations = _enforce_explanation_value_contract( got, got_explanations, answer_variables ) - for variable in answer_variables: - if got.get(variable) is not None: - predictions[variable] = got[variable] - if got_explanations.get(variable): - explanations[variable] = got_explanations[variable] + # The board's merge: a repair fills only cells still missing, and a + # valid earlier value or explanation always stands. + disagreements = record.setdefault("repair_disagreements", []) + predictions = _merge_repair_cells( + predictions, + {v: got.get(v) for v in answer_variables}, + field="prediction", + disagreements=disagreements, + ) + explanations = _merge_repair_cells( + explanations, + {v: got_explanations.get(v) for v in answer_variables}, + field="explanation", + disagreements=disagreements, + ) turn["calls"].append( { "name": "submit_outputs", @@ -1051,14 +1135,26 @@ def _num(value) -> float: return number if math.isfinite(number) else 0.0 +def record_cost_with_prior_attempts(record: dict) -> float: + """A record's spend, including attempts a resume replaced.""" + prior = sum( + _num(a.get("totals", {}).get("cost_usd")) + for a in record.get("prior_attempts", []) + ) + return _num(record.get("totals", {}).get("cost_usd")) + prior + + def household_totals(record: dict) -> dict: + """Tokens, cost and time over every request, truncated ones included; + ``model_turns`` counts the turns that reached the conversation.""" turns = record.get("turns", []) calls = [c for t in turns for c in t.get("calls", []) if c.get("name") == TOOL_NAME] errors = [ c for c in calls if isinstance(c.get("result"), dict) and "error" in c["result"] ] return { - "model_turns": len(turns), + "model_turns": sum(1 for t in turns if not t.get("truncated_and_retried")), + "truncation_retries": sum(1 for t in turns if t.get("truncated_and_retried")), "tool_calls": len(calls), "tool_errors": len(errors), "tool_error_kinds": sorted({c["result"]["error"] for c in errors}), @@ -1226,6 +1322,25 @@ def hit(variable: str, prediction, reference) -> bool: return row_hit_scores(variable, float(reference), float(prediction))["exact"] == 1.0 +def resume_state(households_dir: Path) -> tuple[float, set[str], dict[str, dict]]: + """Spend so far, the households that keep their record, and every record. + + A model outcome (complete, incomplete, turn_limit) stands; only an + infrastructure failure (RERUNNABLE_STATUSES) is rerun, so a resume cannot + redraw a household until it succeeds. Spend includes replaced attempts. + """ + spent = 0.0 + done: set[str] = set() + previous: dict[str, dict] = {} + for path in sorted(households_dir.glob("*.json")): + record = json.loads(path.read_text()) + spent += record_cost_with_prior_attempts(record) + previous[record["scenario_id"]] = record + if record.get("status") not in RERUNNABLE_STATUSES: + done.add(record["scenario_id"]) + return spent, done, previous + + def command_run(args) -> None: _require_engine() if not os.environ.get("GEMINI_API_KEY"): @@ -1242,13 +1357,7 @@ def command_run(args) -> None: if args.limit: ids = ids[: args.limit] - spent_before = 0.0 - done = set() - for path in households_dir.glob("*.json"): - record = json.loads(path.read_text()) - spent_before += _num(record.get("totals", {}).get("cost_usd")) - if record.get("status") == "complete": - done.add(record["scenario_id"]) + spent_before, done, previous = resume_state(households_dir) spend = Spend(args.budget_usd, spent_before) todo = [sid for sid in ids if sid not in done] manifest = { @@ -1282,13 +1391,30 @@ def command_run(args) -> None: def one(sid: str) -> dict: if spend.exhausted(args.reserve_usd): return {"scenario_id": sid, "status": "budget_stop"} - record = run_household( - board["scenarios"][sid], - board["variables"][sid], - pool.call, - completion, - spend, - ) + try: + record = run_household( + board["scenarios"][sid], + board["variables"][sid], + pool.call, + completion, + spend, + ) + except Exception as error: # one household must not end the run + record = { + "scenario_id": sid, + "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) + if sid in previous: + prior = previous[sid] + record["prior_attempts"] = prior.get("prior_attempts", []) + [ + {"status": prior.get("status"), "totals": prior.get("totals", {})} + ] (households_dir / f"{sid}.json").write_text( json.dumps(record, indent=1, default=str) ) @@ -1355,7 +1481,9 @@ def command_score(args) -> None: complete = sorted( sid for sid, r in records.items() if r.get("status") == "complete" ) - households = sorted(records) if args.all_records else complete + # The board's rule: every household counts, and an unparsed output is a + # miss. Scoring only completed households would select on the outcome. + households = complete if args.complete_only else sorted(records) if not households: raise SystemExit("no household records to score") scored = board["scored"][board["scored"].scenario_id.isin(households)] @@ -1782,9 +1910,11 @@ def tool_value_for_output(call: dict, scenario, output: str): return None value = values[names[index]] else: - numbers = [ - float(v) if not isinstance(v, str) else math.nan for v in values.values() - ] + numbers = [] + for item in values.values(): + if item is None or isinstance(item, str): + return None # a non-finite or non-numeric result has no value + numbers.append(float(item)) value = sum(numbers) if aggregation == "any_positive": value = float(value > 0) @@ -1879,9 +2009,9 @@ def situation_diff( """Inputs set differently in the model's situation and the reference's. People are matched by position (the prompt lists them in scenario order); - each group's single unit is matched to the reference's. An input absent on - one side compares as the other side's value against the engine default, - which the caller supplies through ``defaults``. + each group's first unit is matched to the reference's. An input absent on + one side shows as an empty dict; callers drop differences that are only an + input left at the engine default (``_is_default_only``). """ diffs = [] model_people = list(model_situation.get("people", {}).items()) @@ -2081,7 +2211,11 @@ def main_call(record: dict, scenario, outputs: list[str]) -> dict | None: def situation_fidelity(records, board, system, registry, defaults, out: Path) -> None: - """How each household's last computed situation differs from the reference's. + """How each household's main situation differs from the reference's. + + The main situation is ``main_call``'s: the successful call that computed + the most of the household's outputs. (``classify_miss`` reads the last call + that computed the missed output, which is what the answer rests on.) Inputs left at the engine default on one side are dropped. A remaining difference is reported with whether setting it as the reference does would @@ -2154,6 +2288,7 @@ def command_diagnose(args) -> None: registry = None defaults = {} if not args.no_engine_defaults: + _require_engine() system = build_tool_system(reform=True) registry = build_registry(system) defaults = { @@ -2265,7 +2400,11 @@ def main(argv: list[str] | None = None) -> None: help="stop starting households when this much budget remains", ) if name == "score": - p.add_argument("--all-records", action="store_true") + p.add_argument( + "--complete-only", + action="store_true", + help="score only households whose conversation completed", + ) if name == "diagnose": p.add_argument("--no-engine-defaults", action="store_true") args = parser.parse_args(argv) diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index 3f0fe995..edc03bb3 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -207,23 +207,99 @@ def test_too_many_variables_are_refused(): @settings(max_examples=300, deadline=None) @given(args=json_values) -def test_validator_never_raises_and_accepts_only_known_names(args): - """Invariant: any input gets exactly one of a request or an error, and an - accepted request names only registry variables on their own entities.""" +def test_validator_never_raises_on_arbitrary_json(args): + """Invariant: any JSON value gets exactly one of a request or an error.""" request, error = pilot.normalize_request(args, REGISTRY) assert (request is None) != (error is None) - if request is None: - return - for name in request["variables"]: - assert name in REGISTRY - for person in request["situation"]["people"].values(): - for name, periods in person.items(): - assert REGISTRY[name].entity == "person" - assert all(pilot.PERIOD_RE.match(k) for k in periods) - for group, entity in pilot.GROUPS.items(): - for unit in request["situation"].get(group, {}).values(): - for name in unit: - assert name == "members" or REGISTRY[name].entity == entity + + +PERSON_NAMES = st.sampled_from(["head", "kid", "x"]) +VARIABLE_NAMES = st.one_of( + st.sampled_from([*REGISTRY]), st.sampled_from(["nope", "members"]) +) +# Mostly plausible values, so a good share of calls is accepted, plus +# arbitrary JSON scalars so the type checks are exercised too. +PLAUSIBLE = st.one_of( + st.integers(min_value=0, max_value=120), + st.floats(min_value=0, max_value=1e6, allow_nan=False), + st.booleans(), + st.sampled_from(["CO", "TX", "SINGLE", "JOINT"]), +) +SCALARS = st.one_of(PLAUSIBLE, PLAUSIBLE, PLAUSIBLE, json_scalars) +PERIOD_VALUES = st.one_of( + SCALARS, + st.dictionaries(st.sampled_from(["2026", "2026-01", "26"]), SCALARS, max_size=2), +) +INPUTS = st.dictionaries(VARIABLE_NAMES, PERIOD_VALUES, max_size=3) +UNITS = st.dictionaries( + st.sampled_from(["u1", "u2"]), + st.builds( + lambda members, inputs: {"members": members, **inputs}, + st.lists(PERSON_NAMES, max_size=3), + INPUTS, + ), + max_size=2, +) +SITUATIONS = st.builds( + lambda people, groups: {"people": people, **groups}, + st.dictionaries(PERSON_NAMES, INPUTS, min_size=1, max_size=3), + st.fixed_dictionaries( + {}, optional={g: UNITS for g in ("tax_units", "spm_units", "households")} + ), +) +CALLS = st.fixed_dictionaries( + { + "situation": st.one_of(SITUATIONS, SITUATIONS.map(json.dumps)), + "variables": st.lists(VARIABLE_NAMES, min_size=1, max_size=4), + }, + optional={"period": st.sampled_from(["2026", "2026-03", "bad"])}, +) + + +def test_validator_accepts_only_known_names_on_their_entities(): + """Invariant: an accepted call names only registry variables, each on its + own entity, with valid periods, and every unit member is a listed person. + The structured strategy reaches acceptance often enough to check it.""" + accepted = [] + + @settings(max_examples=600, deadline=None, derandomize=True, database=None) + @given(args=CALLS) + def check(args): + request, error = pilot.normalize_request(args, REGISTRY) + assert (request is None) != (error is None) + if request is None: + return + accepted.append(1) + situation = request["situation"] + for name in request["variables"]: + assert name in REGISTRY + for person in situation["people"].values(): + for name, periods in person.items(): + assert REGISTRY[name].entity == "person" + assert periods and all(pilot.PERIOD_RE.match(k) for k in periods) + for group, entity in pilot.GROUPS.items(): + seen = set() + for unit in situation.get(group, {}).values(): + for member in unit["members"]: + assert member in situation["people"] and member not in seen + seen.add(member) + for name, periods in unit.items(): + if name != "members": + assert REGISTRY[name].entity == entity + assert all(pilot.PERIOD_RE.match(k) for k in periods) + + check() + assert len(accepted) >= 30, f"only {len(accepted)} of 600 calls were accepted" + + +def test_deeply_nested_arguments_are_refused_not_raised(): + for text in ("[" * 5000, '{"situation": ' + "[" * 5000 + "}"): + request, error = pilot.normalize_request(text, REGISTRY) + assert request is None and error["error"] == "validation" + request, error = pilot.normalize_request( + {"situation": "[" * 5000, "variables": ["snap"]}, REGISTRY + ) + assert request is None and error["error"] == "validation" @settings(max_examples=200, deadline=None) @@ -840,3 +916,199 @@ def test_main_call_is_the_one_that_computed_the_most_outputs(): outputs = ["snap", "federal_income_tax_before_refundable_credits"] assert pilot.main_call(record, make_scenario(), outputs) is full assert pilot.main_call(household([]), make_scenario(), outputs) is None + + +def test_a_repair_never_overwrites_a_valid_first_value(): + """The board's merge: the first valid value stands even when its + explanation needed a repair and the repair answers differently.""" + first = tool_call( + "submit_outputs", + { + "outputs": { + "snap": {"value": 2553, "explanation": "No trailer here."}, + VARIABLES[1]: {"value": 7, "explanation": "Tax. value = 7"}, + } + }, + "s1", + ) + repair = tool_call( + "submit_outputs", + {"outputs": {"snap": {"value": 0, "explanation": "Recomputed. value = 0"}}}, + "s2", + ) + model = FakeModel([response([first]), response([repair])]) + record = pilot.run_household( + make_scenario(), VARIABLES, fake_tool, model, pilot.Spend(50.0) + ) + assert record["repair_rounds"] == 1 and record["status"] == "complete" + assert record["predictions"]["snap"] == 2553.0 + assert record["explanations"]["snap"] == "Recomputed. value = 0" + assert record["repair_disagreements"][0]["variable"] == "snap" + + +def test_a_truncated_turn_is_retried_at_the_ceiling_and_paid_for(): + truncated = response([calculate()], prompt=1000, completion=16384) + truncated.choices[0].finish_reason = "length" + model = FakeModel( + [ + truncated, + response([calculate()]), + response([submit({"snap": 1, VARIABLES[1]: 2})]), + ] + ) + spend = pilot.Spend(50.0) + record = pilot.run_household(make_scenario(), VARIABLES, fake_tool, model, spend) + budgets = [r.get("max_completion_tokens") for r in model.requests] + assert budgets[0] < budgets[1] == 32_768 + totals = record["totals"] + assert totals["truncation_retries"] == 1 and totals["model_turns"] == 2 + assert spend.spent_usd == pytest.approx(totals["cost_usd"]) + assert totals["cost_usd"] == pytest.approx( + (3000 * 2 + (16384 + 200 + 200) * 10) / 1e6 + ) + + +def test_submit_in_the_same_turn_as_a_calculation_ends_it_without_running_it(): + ran = [] + + def counting_tool(args): + ran.append(args) + return fake_tool(args) + + model = FakeModel([response([calculate(), submit({"snap": 1, VARIABLES[1]: 2})])]) + record = pilot.run_household( + make_scenario(), VARIABLES, counting_tool, model, pilot.Spend(50.0) + ) + assert record["status"] == "complete" and ran == [] + tool_messages = [m for m in record["messages"] if m.get("role") == "tool"] + assert len(tool_messages) == 2 + assert "not run" in tool_messages[0]["content"] + + +def test_replies_without_a_function_call_end_at_the_turn_limit(): + model = FakeModel([response([], content="thinking aloud")] * pilot.MAX_TOTAL_TURNS) + record = pilot.run_household( + make_scenario(), VARIABLES, fake_tool, model, pilot.Spend(50.0) + ) + assert record["status"] == "turn_limit" + assert len(model.requests) == pilot.MAX_TOTAL_TURNS + assert all(v is None for v in record["predictions"].values()) + + +def test_a_litellm_message_keeps_its_thought_signature_round_trip(): + """Gemini refuses a function-call history without thought signatures; the + assistant message goes back as litellm returned it.""" + from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message + + call = ChatCompletionMessageToolCall( + id="call_1__thought__c2lnbmF0dXJl", + type="function", + function=Function(name=pilot.TOOL_NAME, arguments='{"variables": ["snap"]}'), + ) + message = Message(role="assistant", content=None, tool_calls=[call]) + assert pilot._has_thought_signature(message) + data = pilot._message_dict(message) + assert data["tool_calls"][0]["id"] == "call_1__thought__c2lnbmF0dXJl" + assert json.loads(json.dumps(data))["role"] == "assistant" + + +def _scripted_worker(conn, fix_path, reform): + conn.send(("ready", None)) + while True: + args = conn.recv() + if args is None: + return + if args == "sleep": + import time as _time + + _time.sleep(30) + if args == "die": + import os as _os + + _os._exit(1) + conn.send(("ok", {"results": {"echo": args}})) + + +def test_the_pool_times_out_and_restarts_a_worker_and_survives_a_crash(): + pool = pilot.ToolPool(size=1, timeout=1.0, worker_target=_scripted_worker) + try: + assert pool.call("hi")[0] == {"results": {"echo": "hi"}} + result, elapsed = pool.call("sleep") + assert result["error"] == "timeout" and elapsed < 10 + assert pool.call("after-timeout")[0] == {"results": {"echo": "after-timeout"}} + result, _ = pool.call("die") + assert result["error"] == "engine" + assert pool.call("after-crash")[0] == {"results": {"echo": "after-crash"}} + assert (pool.timeouts, pool.crashes, pool.lost) == (1, 1, 0) + finally: + pool.close() + + +def _record(sid, variables, predictions, status="complete", cost=0.01): + record = { + "scenario_id": sid, + "variables": variables, + "predictions": predictions, + "explanations": {v: "x" for v in variables}, + "status": status, + "turns": [{"total_cost_usd": cost, "calls": []}], + "cap_hit": False, + "wall_seconds": 1.0, + } + record["totals"] = pilot.household_totals(record) + return record + + +def test_score_counts_every_household_and_unparsed_outputs_as_misses(tmp_path): + """No selection on outcome: an incomplete household stays in both arms, + its unparsed outputs scored as misses, as the board scores them.""" + import argparse + + board = pilot.load_board() + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + households = tmp_path / "households" + households.mkdir() + for sid, status in (("scenario_000", "complete"), ("scenario_001", "incomplete")): + variables = board["variables"][sid] + predictions = {v: float(reference[(sid, v)]) for v in variables} + if status == "incomplete": + predictions[variables[0]] = None + record = _record(sid, variables, predictions, status) + (households / f"{sid}.json").write_text(json.dumps(record)) + pilot.command_score(argparse.Namespace(out_dir=str(tmp_path), complete_only=False)) + score = json.loads((tmp_path / "score.json").read_text()) + assert score["households"] == 2 + scored = board["scored"][ + board["scored"].scenario_id.isin(["scenario_000", "scenario_001"]) + ] + missed = int( + ( + (scored.scenario_id == "scenario_001") + & (scored.variable == board["variables"]["scenario_001"][0]) + ).sum() + ) + assert score["unweighted_hits"]["arm2"] == len(scored) - missed + pilot.command_score(argparse.Namespace(out_dir=str(tmp_path), complete_only=True)) + assert json.loads((tmp_path / "score.json").read_text())["households"] == 1 + + +def test_a_resume_reruns_only_infrastructure_failures_and_keeps_their_cost(tmp_path): + statuses = { + "a": "complete", + "b": "incomplete", + "c": "turn_limit", + "d": "provider_error", + "e": "budget_stop", + "f": "harness_error", + } + for sid, status in statuses.items(): + record = _record(sid, ["snap"], {"snap": None}, status, cost=1.0) + if sid == "d": + record["prior_attempts"] = [ + {"status": "provider_error", "totals": {"cost_usd": 0.5}} + ] + (tmp_path / f"{sid}.json").write_text(json.dumps(record)) + spent, done, previous = pilot.resume_state(tmp_path) + assert done == {"a", "b", "c"} + assert set(previous) == set(statuses) + assert spent == pytest.approx(6.5) From dc38fd62af556bc5a0d059a3a31e74d09fb428bd Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 13:59:54 -0400 Subject: [PATCH 5/6] Stop on a lost tool pool, keep partial spend, and fuzz the validator's groups Fixes from round 2 of the #229 review: - A replacement worker that dies before ready is killed and counted as lost (any failure, not only RuntimeError), and a pool with no worker returns pool_unavailable, which ends the household as harness_error, so an infrastructure failure is never scored as the model's answer. The run writes the pool's counters to run-summary-*.json. - The structured property test draws each group's unit over a unique, non-empty subset of the drawn people, weighted to the group's own inputs: of 1,000 calls 273 are accepted and 23 carry a group input (floors 50 and 10). It found the validator accepted an input with an empty period dict; that is now refused (no call in the pilot's run sent one). Also: an exception in the turn loop keeps the turns and spend so far; an explanation-only repair names the kept value; score warns when households still need a rerun; the nested-JSON test reaches the RecursionError path. Rescoring the pilot's run reproduces score.json. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 539 +++++++++++++++++++-------------- tests/test_tool_pilot_argon.py | 187 ++++++++++-- 2 files changed, 473 insertions(+), 253 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 1bbd00e6..9f8c6eee 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -50,6 +50,7 @@ import sys import threading import time +from collections import Counter from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, Mapping @@ -63,8 +64,8 @@ ) FIX_MODULE = ROOT / "reference_audit/2026-09-28/fixes/latest_final.py" # The bytes the 2.41.0 references were built with: the sidecar's 2026-10-13 -# engine_upgrade revision pins latest_final.py to this sha256. The modules it -# composes are checked end to end by the oracle subcommand. +# engine_upgrade revision pins latest_final.py to this sha256, and +# verify_fix_pins checks it and every module it composes against that revision. FIX_MODULE_SHA256 = "dbbdd228b99933c6f336af97970821c25c863a18606285e8d6d9ffb3bc9af78e" MODEL_NAME = "gemini-4-argon" @@ -87,6 +88,7 @@ # Infrastructure failures a resume reruns. Every other status is the model's # own outcome and stands, so no household is redrawn until it succeeds. RERUNNABLE_STATUSES = frozenset({"provider_error", "budget_stop", "harness_error"}) +POOL_UNAVAILABLE = "pool_unavailable" MAX_PEOPLE = 16 MAX_UNITS_PER_GROUP = 16 @@ -261,6 +263,11 @@ def _normalize_inputs( ) 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} @@ -583,12 +590,18 @@ def __init__( timeout: float = TOOL_TIMEOUT_SECONDS, reform: bool = True, worker_target: Callable | None = None, + acquire_timeout: float | None = None, ): self._ctx = mp.get_context("spawn") self._fix_path = str(fix_path) self._reform = reform self._target = worker_target or _worker_main self.timeout = timeout + self.acquire_timeout = ( + WORKER_START_TIMEOUT_SECONDS + timeout + if acquire_timeout is None + else acquire_timeout + ) self._idle: queue.Queue = queue.Queue() self._lock = threading.Lock() self.timeouts = 0 @@ -609,22 +622,38 @@ def _start(self): return process, parent def _await_ready(self, worker): + """The worker once it reports ready; RuntimeError (worker killed) if it + does not start, however it fails.""" process, conn = worker - if not conn.poll(WORKER_START_TIMEOUT_SECONDS): + try: + if not conn.poll(WORKER_START_TIMEOUT_SECONDS): + raise RuntimeError("tool worker did not start") + status, _ = conn.recv() + if status != "ready": + raise RuntimeError(f"tool worker failed to start: {status}") + except (RuntimeError, EOFError, OSError, ValueError) as error: process.kill() - raise RuntimeError("tool worker did not start") - status, _ = conn.recv() - if status != "ready": - raise RuntimeError(f"tool worker failed to start: {status}") + process.join(timeout=10) + raise RuntimeError(f"tool worker failed to start: {error}") from error return worker + def counters(self) -> dict: + with self._lock: + return { + "timeouts": self.timeouts, + "crashes": self.crashes, + "lost": self.lost, + } + def call(self, args) -> tuple[dict, float]: started = time.perf_counter() try: - worker = self._idle.get(timeout=WORKER_START_TIMEOUT_SECONDS + self.timeout) + worker = self._idle.get(timeout=self.acquire_timeout) except queue.Empty: + # Infrastructure, not the model: run_household stops the household + # with a rerunnable status rather than show the model this result. return ( - {"error": "engine", "message": "no tool worker is available"}, + {"error": POOL_UNAVAILABLE, "message": "no tool worker is available"}, time.perf_counter() - started, ) process, conn = worker @@ -640,7 +669,7 @@ def call(self, args) -> tuple[dict, float]: "error": "timeout", "message": f"the calculation exceeded {self.timeout:.0f} seconds", } - except (EOFError, BrokenPipeError, OSError) as error: + except (EOFError, BrokenPipeError, OSError, ValueError) as error: with self._lock: self.crashes += 1 result = {"error": "engine", "message": f"tool worker failed: {error}"} @@ -864,217 +893,268 @@ def run_household( answer_variables = list(variables) household_started = time.perf_counter() - for turn_index in range(MAX_TOTAL_TURNS): - if spend.exhausted(): - record["status"] = "budget_stop" - break - if repair_rounds: - kind = "repair" - tool_choice = forced_answer_choice() - tools = [ - calculate_tool, - _build_answer_tool(answer_variables, country=scenario.country), - ] - elif forced or tool_turns >= turn_cap: - kind = "forced_answer" - forced = True - tool_choice = forced_answer_choice() - tools = base["tools"] - else: - kind = "open" - tool_choice = "required" - tools = base["tools"] - kwargs = {k: v for k, v in base.items() if k not in {"tools", "tool_choice"}} - kwargs.update( - { - "model": model_id, - "messages": messages, - "tools": tools, - "tool_choice": tool_choice, + try: + for turn_index in range(MAX_TOTAL_TURNS): + if spend.exhausted(): + record["status"] = "budget_stop" + break + if repair_rounds: + kind = "repair" + tool_choice = forced_answer_choice() + tools = [ + calculate_tool, + _build_answer_tool(answer_variables, country=scenario.country), + ] + elif forced or tool_turns >= turn_cap: + kind = "forced_answer" + forced = True + tool_choice = forced_answer_choice() + tools = base["tools"] + else: + kind = "open" + tool_choice = "required" + tools = base["tools"] + kwargs = { + k: v for k, v in base.items() if k not in {"tools", "tool_choice"} } - ) - _apply_thinking_configuration(kwargs, model_id) - budget_key = ( - "max_completion_tokens" - if "max_completion_tokens" in kwargs - else "max_tokens" - ) - escalated_from = None - try: - response, started_at, completed_at, latency = _request_with_retries( - completion_fn, kwargs + kwargs.update( + { + "model": model_id, + "messages": messages, + "tools": tools, + "tool_choice": tool_choice, + } ) - finish_reason = _response_finish_reason(response, responses_api=False) - if finish_reason == "length" and kwargs.get(budget_key, 0) < ceiling: - usage = _extract_usage_metadata(response, model_id, messages, "") - spend.add(usage.get("total_cost_usd")) - record["turns"].append( - { - "index": turn_index, - "kind": kind, - "truncated_and_retried": True, - "completion_budget": kwargs.get(budget_key), - "latency_seconds": latency, - **usage, - } - ) - escalated_from = kwargs.get(budget_key) - kwargs[budget_key] = ceiling + _apply_thinking_configuration(kwargs, model_id) + budget_key = ( + "max_completion_tokens" + if "max_completion_tokens" in kwargs + else "max_tokens" + ) + escalated_from = None + try: response, started_at, completed_at, latency = _request_with_retries( completion_fn, kwargs ) finish_reason = _response_finish_reason(response, responses_api=False) - except Exception as error: - record["status"] = "provider_error" - record["error"] = f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS] - break - - choices = getattr(response, "choices", None) or [] - if not choices: - record["status"] = "provider_error" - record["error"] = "response has no choices" - break - message = choices[0].message - content = getattr(message, "content", None) - calls = _tool_calls(message) - raw = _serialize_response_payload( - content=content, tool_calls=calls, function_call=None - ) - usage = _extract_usage_metadata( - response, model_id, messages, raw or content or "" - ) - spend.add(usage.get("total_cost_usd")) - turn = { - "index": turn_index, - "kind": kind, - "tool_choice": tool_choice, - "completion_budget": kwargs.get(budget_key), - "escalated_from": escalated_from, - "request_started_at": started_at, - "request_completed_at": completed_at, - "latency_seconds": latency, - "finish_reason": finish_reason, - "content": content, - "reasoning_content": getattr(message, "reasoning_content", None), - "has_thought_signature": _has_thought_signature(message), - "calls": [], - **usage, - } - record["turns"].append(turn) - messages.append(_message_dict(message)) - - answer_calls = [c for c in calls if _call_name(c) == "submit_outputs"] - calculate_calls = [c for c in calls if _call_name(c) == TOOL_NAME] - other_calls = [ - c for c in calls if c not in answer_calls and c not in calculate_calls - ] + if finish_reason == "length" and kwargs.get(budget_key, 0) < ceiling: + usage = _extract_usage_metadata(response, model_id, messages, "") + spend.add(usage.get("total_cost_usd")) + record["turns"].append( + { + "index": turn_index, + "kind": kind, + "truncated_and_retried": True, + "completion_budget": kwargs.get(budget_key), + "latency_seconds": latency, + **usage, + } + ) + escalated_from = kwargs.get(budget_key) + kwargs[budget_key] = ceiling + response, started_at, completed_at, latency = _request_with_retries( + completion_fn, kwargs + ) + finish_reason = _response_finish_reason( + response, responses_api=False + ) + except Exception as error: + record["status"] = "provider_error" + record["error"] = f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS] + break - if answer_calls: - got = extract_predictions( - content, answer_variables, tool_calls=answer_calls - ) - got_explanations = extract_explanations( - content, answer_variables, tool_calls=answer_calls - ) - got, got_explanations = _enforce_explanation_value_contract( - got, got_explanations, answer_variables - ) - # The board's merge: a repair fills only cells still missing, and a - # valid earlier value or explanation always stands. - disagreements = record.setdefault("repair_disagreements", []) - predictions = _merge_repair_cells( - predictions, - {v: got.get(v) for v in answer_variables}, - field="prediction", - disagreements=disagreements, - ) - explanations = _merge_repair_cells( - explanations, - {v: got_explanations.get(v) for v in answer_variables}, - field="explanation", - disagreements=disagreements, + choices = getattr(response, "choices", None) or [] + if not choices: + record["status"] = "provider_error" + record["error"] = "response has no choices" + break + message = choices[0].message + content = getattr(message, "content", None) + calls = _tool_calls(message) + raw = _serialize_response_payload( + content=content, tool_calls=calls, function_call=None ) - turn["calls"].append( - { - "name": "submit_outputs", - "arguments": _call_arguments(answer_calls[0]), - "ignored_calls": len(calculate_calls) - + len(other_calls) - + len(answer_calls) - - 1, - } + usage = _extract_usage_metadata( + response, model_id, messages, raw or content or "" ) - pending = [v for v in variables if predictions[v] is None] - missing_explanation = _missing_explanations(explanations, variables) - still = [ - v for v in variables if v in set(pending) | set(missing_explanation) + spend.add(usage.get("total_cost_usd")) + turn = { + "index": turn_index, + "kind": kind, + "tool_choice": tool_choice, + "completion_budget": kwargs.get(budget_key), + "escalated_from": escalated_from, + "request_started_at": started_at, + "request_completed_at": completed_at, + "latency_seconds": latency, + "finish_reason": finish_reason, + "content": content, + "reasoning_content": getattr(message, "reasoning_content", None), + "has_thought_signature": _has_thought_signature(message), + "calls": [], + **usage, + } + record["turns"].append(turn) + messages.append(_message_dict(message)) + + answer_calls = [c for c in calls if _call_name(c) == "submit_outputs"] + calculate_calls = [c for c in calls if _call_name(c) == TOOL_NAME] + other_calls = [ + c for c in calls if c not in answer_calls and c not in calculate_calls ] - for call in calls: - messages.append( + + if answer_calls: + got = extract_predictions( + content, answer_variables, tool_calls=answer_calls + ) + got_explanations = extract_explanations( + content, answer_variables, tool_calls=answer_calls + ) + got, got_explanations = _enforce_explanation_value_contract( + got, got_explanations, answer_variables + ) + # The board's merge: a repair fills only cells still missing, and a + # valid earlier value or explanation always stands. + disagreements = record.setdefault("repair_disagreements", []) + predictions = _merge_repair_cells( + predictions, + {v: got.get(v) for v in answer_variables}, + field="prediction", + disagreements=disagreements, + ) + explanations = _merge_repair_cells( + explanations, + {v: got_explanations.get(v) for v in answer_variables}, + field="explanation", + disagreements=disagreements, + ) + turn["calls"].append( { - "role": "tool", - "tool_call_id": _call_id(call), - "name": _call_name(call), - "content": json.dumps( - {"accepted": not still, "missing": still} - if call is answer_calls[0] - else {"error": "not run: submit_outputs ended the turn"} - ), + "name": "submit_outputs", + "arguments": _call_arguments(answer_calls[0]), + "ignored_calls": len(calculate_calls) + + len(other_calls) + + len(answer_calls) + - 1, } ) - if not still: - record["status"] = "complete" - break - if repair_rounds >= MAX_REPAIR_ROUNDS: - record["status"] = "incomplete" - break - repair_rounds += 1 - answer_variables = still - messages.append( - { - "role": "user", - "content": ( - "Your submit_outputs call omitted or left unusable these " - f"outputs: {', '.join(still)}. Call submit_outputs again " - "with only these outputs, each with a numeric value and an " - "explanation ending with `value = X`." - ), - } - ) - continue - - if calculate_calls: - tool_turns += 1 - for position, call in enumerate(calculate_calls): - if position >= MAX_CALCULATIONS_PER_TURN: - result, tool_latency = ( + pending = [v for v in variables if predictions[v] is None] + missing_explanation = _missing_explanations(explanations, variables) + still = [ + v for v in variables if v in set(pending) | set(missing_explanation) + ] + for call in calls: + messages.append( { - "error": "not_run", - "message": ( - f"at most {MAX_CALCULATIONS_PER_TURN} " - "calculations per turn" + "role": "tool", + "tool_call_id": _call_id(call), + "name": _call_name(call), + "content": json.dumps( + {"accepted": not still, "missing": still} + if call is answer_calls[0] + else {"error": "not run: submit_outputs ended the turn"} ), - }, - 0.0, + } ) - else: - result, tool_latency = tool_call(_call_arguments(call)) - turn["calls"].append( - { - "name": TOOL_NAME, - "arguments": _call_arguments(call), - "result": result, - "tool_latency_seconds": tool_latency, - } - ) + if not still: + record["status"] = "complete" + break + if repair_rounds >= MAX_REPAIR_ROUNDS: + record["status"] = "incomplete" + break + repair_rounds += 1 + answer_variables = still + kept = [ + f"{v} (keep your value {predictions[v]:g}; its explanation must " + f"end with `value = {predictions[v]:g}`)" + for v in still + if predictions[v] is not None + ] messages.append( { - "role": "tool", - "tool_call_id": _call_id(call), - "name": TOOL_NAME, - "content": json.dumps(result), + "role": "user", + "content": ( + "Your submit_outputs call omitted or left unusable these " + f"outputs: {', '.join(still)}. Call submit_outputs again " + "with only these outputs, each with a numeric value and an " + "explanation ending with `value = X`." + + ( + " Explanations only are missing for: " + + "; ".join(kept) + + "." + if kept + else "" + ) + ), } ) + continue + + if calculate_calls: + tool_turns += 1 + pool_down = False + for position, call in enumerate(calculate_calls): + if position >= MAX_CALCULATIONS_PER_TURN: + result, tool_latency = ( + { + "error": "not_run", + "message": ( + f"at most {MAX_CALCULATIONS_PER_TURN} " + "calculations per turn" + ), + }, + 0.0, + ) + else: + result, tool_latency = tool_call(_call_arguments(call)) + if ( + isinstance(result, dict) + and result.get("error") == POOL_UNAVAILABLE + ): + pool_down = True + turn["calls"].append( + { + "name": TOOL_NAME, + "arguments": _call_arguments(call), + "result": result, + "tool_latency_seconds": tool_latency, + } + ) + messages.append( + { + "role": "tool", + "tool_call_id": _call_id(call), + "name": TOOL_NAME, + "content": json.dumps(result), + } + ) + for call in other_calls: + messages.append( + { + "role": "tool", + "tool_call_id": _call_id(call), + "name": _call_name(call) or "unknown", + "content": json.dumps({"error": "unknown function"}), + } + ) + if pool_down: + # Infrastructure, not the model: a rerunnable status, so the + # household is not scored on an answer the pool prevented. + record["status"] = "harness_error" + record["error"] = "the tool pool had no worker" + break + if tool_turns >= turn_cap: + messages.append( + { + "role": "user", + "content": ( + f"You have used all {turn_cap} calculation turns. Call " + "submit_outputs now with your final answers." + ), + } + ) + continue + for call in other_calls: messages.append( { @@ -1084,35 +1164,17 @@ def run_household( "content": json.dumps({"error": "unknown function"}), } ) - if tool_turns >= turn_cap: - messages.append( - { - "role": "user", - "content": ( - f"You have used all {turn_cap} calculation turns. Call " - "submit_outputs now with your final answers." - ), - } - ) - continue - - for call in other_calls: messages.append( { - "role": "tool", - "tool_call_id": _call_id(call), - "name": _call_name(call) or "unknown", - "content": json.dumps({"error": "unknown function"}), + "role": "user", + "content": "Call calculate_household or submit_outputs.", } ) - messages.append( - { - "role": "user", - "content": "Call calculate_household or submit_outputs.", - } - ) - else: - record["status"] = "turn_limit" + else: + record["status"] = "turn_limit" + except Exception as error: # keep the turns, and their spend, so far + record["status"] = "harness_error" + record["error"] = f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS] if record["status"] == "running": record["status"] = "incomplete" @@ -1432,8 +1494,24 @@ def one(sid: str) -> dict: ) as executor: results = list(executor.map(one, todo)) finally: + counters = pool.counters() pool.close() stopped = [r["scenario_id"] for r in results if r.get("status") == "budget_stop"] + rerunnable = sorted( + r["scenario_id"] for r in results if r.get("status") in RERUNNABLE_STATUSES + ) + summary = { + "finished_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "spent_usd": spend.spent_usd, + "budget_usd": args.budget_usd, + "statuses": dict(sorted(Counter(r.get("status") for r in results).items())), + "rerunnable_households": rerunnable, + "tool_pool": counters, + } + ( + out / f"run-summary-{summary['finished_at_utc'].replace(':', '')}.json" + ).write_text(json.dumps(summary, indent=2)) + print(f"tool pool: {counters}; rerunnable households: {rerunnable}") print( f"spent ${spend.spent_usd:.4f} of ${args.budget_usd:.2f}; " f"budget stops: {stopped}" @@ -1486,6 +1564,14 @@ def command_score(args) -> None: households = complete if args.complete_only else sorted(records) if not households: raise SystemExit("no household records to score") + leftover = sorted( + sid for sid in households if records[sid].get("status") in RERUNNABLE_STATUSES + ) + if leftover: + print( + f"WARNING: {len(leftover)} households ended on an infrastructure " + f"failure and score as misses until a resume reruns them: {leftover}" + ) scored = board["scored"][board["scored"].scenario_id.isin(households)] arm1 = arm1_predictions() arm1 = arm1[arm1.scenario_id.isin(households)] @@ -1585,6 +1671,7 @@ def command_score(args) -> None: { "households": len(households), "incomplete_households": sorted(set(records) - set(complete)), + "rerunnable_households": leftover, "scored_outputs": int(len(scored)), "summary": summary[columns].to_dict(orient="records"), "unweighted_hits": { diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index edc03bb3..4c8c27f9 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -214,9 +214,7 @@ def test_validator_never_raises_on_arbitrary_json(args): PERSON_NAMES = st.sampled_from(["head", "kid", "x"]) -VARIABLE_NAMES = st.one_of( - st.sampled_from([*REGISTRY]), st.sampled_from(["nope", "members"]) -) +OTHER_NAMES = st.sampled_from([*REGISTRY, "nope", "members"]) # Mostly plausible values, so a good share of calls is accepted, plus # arbitrary JSON scalars so the type checks are exercised too. PLAUSIBLE = st.one_of( @@ -230,27 +228,61 @@ def test_validator_never_raises_on_arbitrary_json(args): SCALARS, st.dictionaries(st.sampled_from(["2026", "2026-01", "26"]), SCALARS, max_size=2), ) -INPUTS = st.dictionaries(VARIABLE_NAMES, PERIOD_VALUES, max_size=3) -UNITS = st.dictionaries( - st.sampled_from(["u1", "u2"]), - st.builds( - lambda members, inputs: {"members": members, **inputs}, - st.lists(PERSON_NAMES, max_size=3), - INPUTS, - ), - max_size=2, -) -SITUATIONS = st.builds( - lambda people, groups: {"people": people, **groups}, - st.dictionaries(PERSON_NAMES, INPUTS, min_size=1, max_size=3), - st.fixed_dictionaries( - {}, optional={g: UNITS for g in ("tax_units", "spm_units", "households")} - ), -) + + +def names_on(entity: str) -> list[str]: + return [n for n, info in REGISTRY.items() if info.entity == entity] + + +def inputs_for(entity: str): + """Inputs weighted toward the entity's own variables, with some misplaced + or unknown names so the rejections are exercised as well.""" + own = names_on(entity) + if not own: # families and marital units carry no inputs in REGISTRY + return st.one_of( + st.just({}), + st.just({}), + st.dictionaries(OTHER_NAMES, PERIOD_VALUES, max_size=1), + ) + names = st.one_of(*[st.sampled_from(own)] * 4, OTHER_NAMES) + return st.dictionaries(names, PERIOD_VALUES, max_size=2) + + +def situations(): + """People first, then each group's single unit over a non-empty, unique + subset of those people, so accepted calls carry groups.""" + + def with_groups(people): + names = sorted(people) + + def unit(entity): + return st.builds( + lambda members, inputs: {"members": members, **inputs}, + st.lists( + st.sampled_from(names), min_size=1, max_size=len(names), unique=True + ), + inputs_for(entity), + ) + + groups = st.fixed_dictionaries( + {}, + optional={ + group: st.dictionaries(st.just("u1"), unit(entity), min_size=1) + for group, entity in pilot.GROUPS.items() + }, + ) + return groups.map(lambda g: {"people": people, **g}) + + return st.dictionaries( + PERSON_NAMES, inputs_for("person"), min_size=1, max_size=3 + ).flatmap(with_groups) + + +SITUATIONS = situations() CALLS = st.fixed_dictionaries( { "situation": st.one_of(SITUATIONS, SITUATIONS.map(json.dumps)), - "variables": st.lists(VARIABLE_NAMES, min_size=1, max_size=4), + "variables": st.lists(OTHER_NAMES, min_size=1, max_size=4), }, optional={"period": st.sampled_from(["2026", "2026-03", "bad"])}, ) @@ -258,11 +290,13 @@ def test_validator_never_raises_on_arbitrary_json(args): def test_validator_accepts_only_known_names_on_their_entities(): """Invariant: an accepted call names only registry variables, each on its - own entity, with valid periods, and every unit member is a listed person. - The structured strategy reaches acceptance often enough to check it.""" + own entity, with valid periods, and every unit member is a listed person + in at most one unit of its group. The strategy reaches acceptance, with + groups and group inputs, often enough to check it.""" accepted = [] + with_group_inputs = [] - @settings(max_examples=600, deadline=None, derandomize=True, database=None) + @settings(max_examples=1000, deadline=None, derandomize=True, database=None) @given(args=CALLS) def check(args): request, error = pilot.normalize_request(args, REGISTRY) @@ -277,6 +311,7 @@ def check(args): for name, periods in person.items(): assert REGISTRY[name].entity == "person" assert periods and all(pilot.PERIOD_RE.match(k) for k in periods) + group_input = False for group, entity in pilot.GROUPS.items(): seen = set() for unit in situation.get(group, {}).values(): @@ -285,19 +320,30 @@ def check(args): seen.add(member) for name, periods in unit.items(): if name != "members": + group_input = True assert REGISTRY[name].entity == entity assert all(pilot.PERIOD_RE.match(k) for k in periods) + if group_input: + with_group_inputs.append(1) check() - assert len(accepted) >= 30, f"only {len(accepted)} of 600 calls were accepted" + assert len(accepted) >= 50, f"only {len(accepted)} of 1000 calls were accepted" + assert len(with_group_inputs) >= 10, ( + f"only {len(with_group_inputs)} accepted calls carried a group input" + ) def test_deeply_nested_arguments_are_refused_not_raised(): - for text in ("[" * 5000, '{"situation": ' + "[" * 5000 + "}"): + # Deep enough to raise RecursionError (not JSONDecodeError) in json.loads, + # yet under the size cap, so the RecursionError path is what refuses it. + deep = "[" * 100_000 + "]" * 100_000 + with pytest.raises(RecursionError): + json.loads(deep) + for text in (deep, '{"situation": ' + deep + "}"): request, error = pilot.normalize_request(text, REGISTRY) assert request is None and error["error"] == "validation" request, error = pilot.normalize_request( - {"situation": "[" * 5000, "variables": ["snap"]}, REGISTRY + {"situation": "[" * 100_000 + "]" * 100_000, "variables": ["snap"]}, REGISTRY ) assert request is None and error["error"] == "validation" @@ -1112,3 +1158,90 @@ def test_a_resume_reruns_only_infrastructure_failures_and_keeps_their_cost(tmp_p assert done == {"a", "b", "c"} assert set(previous) == set(statuses) assert spent == pytest.approx(6.5) + + +def _dies_on_restart_worker(conn, marker, reform): + """Starts once; a replacement (the marker exists) exits before ready.""" + from pathlib import Path as _Path + + if _Path(marker).exists(): + import os as _os + + _os._exit(3) + _Path(marker).touch() + conn.send(("ready", None)) + while True: + args = conn.recv() + if args is None: + return + if args == "die": + import os as _os + + _os._exit(1) + conn.send(("ok", {"results": {"echo": args}})) + + +def test_a_replacement_that_fails_to_start_is_counted_lost_not_raised(tmp_path): + pool = pilot.ToolPool( + size=1, + fix_path=tmp_path / "started", + timeout=1.0, + worker_target=_dies_on_restart_worker, + acquire_timeout=0.5, + ) + try: + result, _ = pool.call("die") + assert result["error"] == "engine" + assert pool.counters() == {"timeouts": 0, "crashes": 1, "lost": 1} + result, _ = pool.call("anything") + assert result["error"] == pilot.POOL_UNAVAILABLE + finally: + pool.close() + + +def test_an_unavailable_pool_stops_the_household_as_rerunnable(): + def no_pool(args): + return {"error": pilot.POOL_UNAVAILABLE, "message": "none"}, 0.5 + + model = FakeModel([response([calculate()])]) + record = pilot.run_household( + make_scenario(), VARIABLES, no_pool, model, pilot.Spend(50.0) + ) + assert record["status"] == "harness_error" + assert record["status"] in pilot.RERUNNABLE_STATUSES + assert len(model.requests) == 1 + + +def test_a_harness_exception_keeps_the_turns_and_their_spend(): + def broken_tool(args): + raise RuntimeError("tool bug") + + model = FakeModel([response([calculate()])]) + spend = pilot.Spend(50.0) + record = pilot.run_household(make_scenario(), VARIABLES, broken_tool, model, spend) + assert record["status"] == "harness_error" and "tool bug" in record["error"] + assert len(record["turns"]) == 1 + assert record["totals"]["cost_usd"] > 0 + assert record["totals"]["cost_usd"] == pytest.approx(spend.spent_usd) + + +def test_an_explanation_only_repair_names_the_kept_value(): + first = tool_call( + "submit_outputs", + { + "outputs": { + "snap": {"value": 2553, "explanation": "No trailer."}, + VARIABLES[1]: {"value": 7, "explanation": "Tax. value = 7"}, + } + }, + "s1", + ) + repair = tool_call( + "submit_outputs", + {"outputs": {"snap": {"value": 2553, "explanation": "SNAP. value = 2553"}}}, + "s2", + ) + model = FakeModel([response([first]), response([repair])]) + pilot.run_household(make_scenario(), VARIABLES, fake_tool, model, pilot.Spend(50.0)) + prompt = model.requests[1]["messages"][-1]["content"] + assert "keep your value 2553" in prompt and "value = 2553" in prompt From bed93f45c7ba9dc01e34da7c1235d75b3b614100 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 15:16:15 -0400 Subject: [PATCH 6/6] Address the round-3 review nits on #229 - Kept values in a repair prompt are plain decimals (plain_number), which the board's value = X trailer pattern reads back; :g gave 12345.7 for 12345.67 and 1.23457e+06 for 1234567. - A household stops at the first pool_unavailable result, and no new household starts once every worker is lost. - The run summary lists each rerunnable household's error. - The validator fuzz draws typed group inputs and some invalid memberships (a ghost, a person twice or in two units). Over four seeds it accepts 322-346 of 1,000 calls, 58-71 with units and 14-38 with group inputs, and rejects 510-557 on membership; the floors (100, 20, 5) sit under the worst seed. Co-Authored-By: Claude Opus 5.5 --- scripts/tool_pilot_argon.py | 37 ++++++++++++--- tests/test_tool_pilot_argon.py | 86 ++++++++++++++++++++++++++++------ 2 files changed, 101 insertions(+), 22 deletions(-) diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index 9f8c6eee..def1829e 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -36,6 +36,7 @@ import argparse import concurrent.futures import copy +import decimal import difflib import hashlib import importlib.metadata @@ -1065,10 +1066,11 @@ def run_household( repair_rounds += 1 answer_variables = still kept = [ - f"{v} (keep your value {predictions[v]:g}; its explanation must " - f"end with `value = {predictions[v]:g}`)" + f"{v} (keep your value {number}; its explanation must end " + f"with `value = {number}`)" for v in still if predictions[v] is not None + for number in [plain_number(predictions[v])] ] messages.append( { @@ -1107,11 +1109,6 @@ def run_household( ) else: result, tool_latency = tool_call(_call_arguments(call)) - if ( - isinstance(result, dict) - and result.get("error") == POOL_UNAVAILABLE - ): - pool_down = True turn["calls"].append( { "name": TOOL_NAME, @@ -1120,6 +1117,12 @@ def run_household( "tool_latency_seconds": tool_latency, } ) + if ( + isinstance(result, dict) + and result.get("error") == POOL_UNAVAILABLE + ): + pool_down = True + break messages.append( { "role": "tool", @@ -1189,6 +1192,13 @@ def run_household( return record +def plain_number(value) -> str: + """A number as plain decimal digits, never scientific notation, so the + board's `value = X` trailer pattern can read it back.""" + text = format(decimal.Decimal(repr(float(value))), "f") + return text.rstrip("0").rstrip(".") if "." in text else text + + def _num(value) -> float: try: number = float(value) @@ -1453,6 +1463,14 @@ def command_run(args) -> None: def one(sid: str) -> dict: if spend.exhausted(args.reserve_usd): return {"scenario_id": sid, "status": "budget_stop"} + if pool.counters()["lost"] >= args.workers: + # Every worker is gone: start nothing more. The household keeps + # no record, so a resume runs it. + return { + "scenario_id": sid, + "status": "harness_error", + "error": "every tool worker was lost", + } try: record = run_household( board["scenarios"][sid], @@ -1506,6 +1524,11 @@ def one(sid: str) -> dict: "budget_usd": args.budget_usd, "statuses": dict(sorted(Counter(r.get("status") for r in results).items())), "rerunnable_households": rerunnable, + "errors": { + r["scenario_id"]: r.get("error") + for r in results + if r.get("status") in RERUNNABLE_STATUSES and r.get("error") + }, "tool_pool": counters, } ( diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index 4c8c27f9..9a9d87d1 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -234,40 +234,71 @@ def names_on(entity: str) -> list[str]: return [n for n, info in REGISTRY.items() if info.entity == entity] +def typed_value(name: str): + """A value of the variable's own type, so a group input can be valid.""" + info = REGISTRY.get(name) + if info is None: + return SCALARS + if info.value_type == "Enum": + return st.sampled_from(list(info.possible_values)) + if info.value_type in {"float", "int"}: + return st.one_of( + st.integers(min_value=0, max_value=120), + st.floats(min_value=0, max_value=1e6, allow_nan=False), + ) + if info.value_type == "bool": + return st.booleans() + return SCALARS + + def inputs_for(entity: str): - """Inputs weighted toward the entity's own variables, with some misplaced - or unknown names so the rejections are exercised as well.""" + """Mostly typed inputs on the entity's own variables, with some misplaced + or unknown names and arbitrary values so the rejections are exercised.""" own = names_on(entity) + wrong = st.tuples(OTHER_NAMES, PERIOD_VALUES) if not own: # families and marital units carry no inputs in REGISTRY return st.one_of( - st.just({}), - st.just({}), - st.dictionaries(OTHER_NAMES, PERIOD_VALUES, max_size=1), + st.just({}), st.just({}), st.lists(wrong, max_size=1).map(dict) ) - names = st.one_of(*[st.sampled_from(own)] * 4, OTHER_NAMES) - return st.dictionaries(names, PERIOD_VALUES, max_size=2) + right = st.sampled_from(own).flatmap( + lambda name: st.one_of( + typed_value(name), + st.dictionaries( + st.sampled_from(["2026", "2026-01"]), + typed_value(name), + min_size=1, + max_size=2, + ), + ).map(lambda value: (name, value)) + ) + return st.lists(st.one_of(right, right, right, wrong), max_size=2).map(dict) def situations(): - """People first, then each group's single unit over a non-empty, unique - subset of those people, so accepted calls carry groups.""" + """People first, then each group's units, mostly over unique, non-empty + subsets of those people; some memberships are invalid (a ghost, a person + twice or in two units), so the membership checks can fail.""" def with_groups(people): names = sorted(people) + valid = st.lists( + st.sampled_from(names), min_size=1, max_size=len(names), unique=True + ) + invalid = st.lists(st.sampled_from([*names, "ghost"]), max_size=3) def unit(entity): return st.builds( lambda members, inputs: {"members": members, **inputs}, - st.lists( - st.sampled_from(names), min_size=1, max_size=len(names), unique=True - ), + st.one_of(valid, valid, valid, invalid), inputs_for(entity), ) groups = st.fixed_dictionaries( {}, optional={ - group: st.dictionaries(st.just("u1"), unit(entity), min_size=1) + group: st.dictionaries( + st.sampled_from(["u1", "u1", "u2"]), unit(entity), min_size=1 + ) for group, entity in pilot.GROUPS.items() }, ) @@ -294,6 +325,7 @@ def test_validator_accepts_only_known_names_on_their_entities(): in at most one unit of its group. The strategy reaches acceptance, with groups and group inputs, often enough to check it.""" accepted = [] + with_units = [] with_group_inputs = [] @settings(max_examples=1000, deadline=None, derandomize=True, database=None) @@ -312,6 +344,8 @@ def check(args): assert REGISTRY[name].entity == "person" assert periods and all(pilot.PERIOD_RE.match(k) for k in periods) group_input = False + if any(situation.get(group) for group in pilot.GROUPS): + with_units.append(1) for group, entity in pilot.GROUPS.items(): seen = set() for unit in situation.get(group, {}).values(): @@ -327,8 +361,12 @@ def check(args): with_group_inputs.append(1) check() - assert len(accepted) >= 50, f"only {len(accepted)} of 1000 calls were accepted" - assert len(with_group_inputs) >= 10, ( + # Floors well under the worst of four seeds measured (322 accepted, 58 + # with units, 14 with group inputs), so a Hypothesis upgrade or an edit to + # this function cannot quietly make the checks vacuous or flaky. + assert len(accepted) >= 100, f"only {len(accepted)} of 1000 calls were accepted" + assert len(with_units) >= 20, f"only {len(with_units)} accepted calls had units" + assert len(with_group_inputs) >= 5, ( f"only {len(with_group_inputs)} accepted calls carried a group input" ) @@ -1245,3 +1283,21 @@ def test_an_explanation_only_repair_names_the_kept_value(): pilot.run_household(make_scenario(), VARIABLES, fake_tool, model, pilot.Spend(50.0)) prompt = model.requests[1]["messages"][-1]["content"] assert "keep your value 2553" in prompt and "value = 2553" in prompt + + +@pytest.mark.parametrize( + "value, text", + [ + (2553.0, "2553"), + (12345.67, "12345.67"), + (1234567.0, "1234567"), + (1e-05, "0.00001"), + (-12.5, "-12.5"), + (0.0, "0"), + ], +) +def test_plain_number_round_trips_through_the_boards_trailer_pattern(value, text): + from policybench.eval_no_tools import _extract_terminal_explanation_value + + assert pilot.plain_number(value) == text + assert _extract_terminal_explanation_value(f"Reason. value = {text}") == value