diff --git a/scripts/tool_pilot_argon.py b/scripts/tool_pilot_argon.py index c8a1559b..f09e0b6f 100644 --- a/scripts/tool_pilot_argon.py +++ b/scripts/tool_pilot_argon.py @@ -48,9 +48,11 @@ import argparse import concurrent.futures +import contextlib import copy import decimal import difflib +import functools import hashlib import importlib.metadata import importlib.util @@ -152,18 +154,83 @@ # with no tools), so a prompt version other than the board's has its own # no-tools baseline. Repairs, retries and truncation then follow this pilot's # loop, not the board's; they matter only when a first answer does not parse. -ARMS = ("none", "pe", "web", "web+pe", "pe2") +ARMS = ("none", "pe", "web", "web+pe", "pe2", "pe3") ARM_SUFFIX = { "none": "+no-tools", "pe": "+pe-tool", "web": "+web", "web+pe": "+web+pe-tool", "pe2": "+pe-tool-v2", + "pe3": "+pe-tool-v3", } -# "pe2" is the calculator, tool version 2 (Max, 2026-10-23: "yes go"), alone: -# it reads the names the prompt uses, places inputs set on the wrong entity, -# flags inputs that replace a formula, and explains a value on request. -TOOL_VERSION_BY_ARM = {"pe": 1, "web+pe": 1, "pe2": 2} +# The calculator's features, each on or off, so a study can change one at a +# time (Max, 2026-10-23: "measuring this as an experiment could be +# publishable"; "do all these tests with a cheap model"): +# - names: a name written with spaces, or a benchmark output key, is read as +# its engine variable; +# - labels: a fact written under the label the prompt gives it (a prompt +# labels some facts in words, e.g. employment_income as "gross wages and +# salaries"; policybench.prompts.INPUT_LABEL_OVERRIDES) is read as its +# variable, and a fact the reference also sets under another name +# (policybench.scenarios.PE_INPUT_ALIASES: usual weekly hours) is set there +# too; +# - placement: an input set on another entity moves to the one unit holding +# its people; +# - overrides: a realized contribution is read as its _desired input, and an +# input that replaces a formula is flagged; +# - explain: `explain` traces up to three variables; +# - concise: an explanation is brief unless `detail` is "full"; +# - flat: the household is described as people with roles and facts +# (`household`), and the tool builds its units as the benchmark does. +# "pe2" is tool version 2 (names, placement, overrides and explain); "pe3" is +# all seven. Any other set +# runs as an arm named "pe:" plus its features joined by "+", e.g. +# "pe:names+flat". +NAMES, LABELS, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT = ( + "names", + "labels", + "placement", + "overrides", + "explain", + "concise", + "flat", +) +TOOL_FEATURES = (NAMES, LABELS, PLACEMENT, OVERRIDES, EXPLAIN, CONCISE, FLAT) +V2_FEATURES = frozenset({NAMES, PLACEMENT, OVERRIDES, EXPLAIN}) +FEATURES_BY_ARM = { + "pe": frozenset(), + "web+pe": frozenset(), + "pe2": V2_FEATURES, + "pe3": frozenset(TOOL_FEATURES), +} +_FEATURE_ARM = re.compile(r"^pe:([a-z]+(?:\+[a-z]+)*)$") + + +def arm_features(arm: str) -> frozenset | None: + """The calculator features an arm offers, or None for an arm without the + calculator. Refuses an unknown feature, or concise without explain.""" + if arm in FEATURES_BY_ARM: + return FEATURES_BY_ARM[arm] + match = _FEATURE_ARM.match(arm or "") + if match is None: + return None + features = frozenset(match.group(1).split("+")) + unknown = sorted(features - set(TOOL_FEATURES)) + if unknown: + raise SystemExit( + f"unknown tool features {unknown}; the features are {TOOL_FEATURES}" + ) + if CONCISE in features and EXPLAIN not in features: + raise SystemExit("the concise feature needs explain") + return features + + +def arm_suffix(arm: str) -> str: + if arm in ARM_SUFFIX: + return ARM_SUFFIX[arm] + features = arm_features(arm) + return f"+pe-tool[{'+'.join(f for f in TOOL_FEATURES if f in features)}]" + TOOL_NOTICE_V2 = ( "\n\nYou also have a `calculate_household` function. It runs PolicyEngine US " @@ -183,6 +250,90 @@ "instructed above.", "Then give your final answers as the single JSON object described above.", ) + + +def _listing(items: list[str]) -> str: + if len(items) <= 2: + return " and ".join(items) + return ", ".join(items[:-1]) + ", and " + items[-1] + + +def compose_tool_notice(features: frozenset, tool_contract: bool) -> str: + """The calculator notice for a set of features: the opening, a sentence + per feature, then the turn limit and the answer instruction. No feature + gives TOOL_NOTICE and the version 2 features TOOL_NOTICE_V2, as those arms + ran. TOOL_NOTICE_V2 says the facts are labeled with the engine's variable + names, which holds for most facts but not those the prompt labels in words + ("gross wages and salaries"); every other set says only what it reads.""" + if not features: + return TOOL_NOTICE if tool_contract else JSON_TOOL_NOTICE + if features == V2_FEATURES: + return TOOL_NOTICE_V2 if tool_contract else JSON_TOOL_NOTICE_V2 + what = ( + "household you describe" + if FLAT in features + else "household situation you build" + ) + text = ( + "\n\nYou also have a `calculate_household` function. It runs PolicyEngine US " + "(policyengine-us 2.41.0), an open-source rules engine for US federal and " + f"state taxes and benefits, on a {what} from the facts above." + ) + if FLAT in features: + text += ( + " Describe it as `household`: each person with their role (head, spouse " + "or dependent) and facts, and the facts about the household as a " + "whole; it builds the household's units as the facts imply." + ) + if LABELS in features and NAMES in features: + text += ( + " It accepts the facts above under the labels they have there, its " + "variable names written with spaces, and the output keys you are asked " + "for. Set every listed fact." + ) + elif LABELS in features: + text += ( + " It accepts the facts above under the labels they have there. Set " + "every listed fact." + ) + elif NAMES in features: + text += ( + " It accepts its variable names written with spaces, and the output " + "keys you are asked for. Set every listed fact." + ) + notes = [] + if NAMES in features or LABELS in features: + notes.append("which names it read") + if PLACEMENT in features: + notes.append("which inputs it moved to the entity they belong to") + if OVERRIDES in features: + notes.append("which inputs replace a formula") + explain = "" + if EXPLAIN in features: + explain = ( + "ask it to `explain` up to three variables to see in brief how it " + "computed them, with the parameters and laws it used, or the fuller " + 'trace with `detail` set to "full"' + if CONCISE in features + else "ask it to `explain` up to three variables to see how it computed " + "them, with the parameters and laws it used" + ) + if notes and explain: + text += f" Its result says {_listing(notes)}; {explain}." + elif notes: + text += f" Its result says {_listing(notes)}." + elif explain: + text += f" You may {explain}." + text += " You may call it in up to 8 turns. " + text += ( + "Then call `submit_outputs` exactly once with your final answers, as " + "instructed above." + if tool_contract + else "Then give your final answers as the single JSON object described above." + ) + return text + + # The engine version the notices and the tool description state: the board's # (set by command_run), so a household set on a newer engine says so. The # frozen board's 2.41.0 leaves every text byte for byte as before. @@ -193,6 +344,18 @@ def set_stated_engine(version: str) -> None: _STATED_ENGINE["version"] = version +@contextlib.contextmanager +def stating_engine(version: str): + """The texts as a run on `version` sent them, for a check that rebuilds a + notice; the stated engine is restored afterwards.""" + before = _STATED_ENGINE["version"] + set_stated_engine(version) + try: + yield + finally: + set_stated_engine(before) + + def engine_text(text: str) -> str: return text.replace( f"policyengine-us {REFERENCE_ENGINE}", @@ -511,19 +674,106 @@ def refuse_policyengine(label: str, value) -> None: MAX_EXPLAIN = 3 -def calculate_tool_schema(version: int = 1) -> dict: - if version == 2: +FLAT_DESCRIPTION = ( + "Run PolicyEngine US (policyengine-us 2.41.0) on one household and return " + "the requested variables. `household` is a JSON object: `state` (the " + "two-letter state code), `people` (a list, each with `role` -- head, " + 'spouse or dependent --, `facts`, for example {"age": 40, ' + '"employment_income": 30000}, and optionally `name`, by default the role, ' + "numbered for dependents), and `household_facts` (facts about the " + "household as a whole: its tax unit, family, SPM unit or home). The head " + "and spouse are one couple and everyone else is on their own; everyone " + "shares one tax unit, family, SPM unit and household. A fact is a value " + 'for `period`, or {period: value} with a period such as "2026" or ' + '"2026-01". Facts you do not give take the engine\'s defaults. `variables` ' + 'lists the variables to compute for `period` (default "2026"). Each ' + "result gives the variable's entity, label and value for each member of " + "that entity; a monthly variable computed for a year is summed over its " + "months. An unknown name or a malformed household returns an error with " + "suggestions instead of results." +) +# The version 1 description's last sentence, which placement amends. +_ENTITY_ERROR_SENTENCE = ( + "An unknown name, a variable set on the wrong entity, or a malformed " + "situation returns an error with suggestions instead of results." +) + + +def calculate_tool_schema(version: int = 1, features: frozenset | None = None) -> dict: + if features is None: + features = V2_FEATURES if version == 2 else frozenset() + if features: schema = calculate_tool_schema(1) function = schema["function"] - function["description"] = engine_text(CALCULATE_DESCRIPTION_V2) - function["parameters"]["properties"]["explain"] = { - "type": "array", - "items": {"type": "string"}, - "description": ( - f"Up to {MAX_EXPLAIN} variables to explain: how the engine " - "computed them." - ), - } + properties = function["parameters"]["properties"] + if features == V2_FEATURES: + function["description"] = engine_text(CALCULATE_DESCRIPTION_V2) + else: + description = ( + FLAT_DESCRIPTION if FLAT in features else CALCULATE_DESCRIPTION + ) + if PLACEMENT in features and FLAT not in features: + description = description.replace( + _ENTITY_ERROR_SENTENCE, + "An unknown name or a malformed situation returns an error " + "with suggestions instead of results.", + ) + notes = [] + if LABELS in features: + notes.append( + "A fact may be written under the label it has in the facts " + "you were given." + ) + if NAMES in features: + notes.append( + "A name may also be written with spaces or be one of the " + "output keys you are asked for." + ) + if PLACEMENT in features and FLAT not in features: + notes.append( + "An input set on the wrong entity is moved to the one unit " + "holding its people; where no single unit holds them, it " + "returns an error." + ) + if OVERRIDES in features: + notes.append( + "Setting a variable the engine computes replaces its formula; a " + "contribution set as the amount contributed is read as the " + "amount asked for before limits (`..._desired`)." + ) + if notes: + notes.append("The result's `notes` say what it did.") + function["description"] = engine_text(" ".join([description, *notes])) + if FLAT in features: + del properties["situation"] + properties["household"] = { + "type": "string", + "description": ( + "The household as a JSON object (state, people with role and " + "facts, household_facts), written as a JSON string." + ), + } + function["parameters"]["required"] = ["household", "variables"] + if EXPLAIN in features: + properties["explain"] = { + "type": "array", + "items": {"type": "string"}, + "description": ( + f"Up to {MAX_EXPLAIN} variables to explain: how the engine " + "computed them." + ), + } + if CONCISE in features: + properties["detail"] = { + "type": "string", + "enum": ["brief", "full"], + "description": ( + 'How much an explanation shows: "brief" (the default: the ' + 'value, what it came from, and the main parameters) or "full" ' + f"(the trace, up to {EXPLAIN_MAX_LINES} steps and " + f"{EXPLAIN_MAX_PARAMETERS} parameters)." + ), + } return schema return { "type": "function", @@ -598,13 +848,9 @@ class VariableInfo: def stated_inputs() -> frozenset[str]: """Every variable a household fact the benchmark states can set.""" - from policybench.scenarios import ( - PE_INPUT_ALIASES, - PE_INPUT_RENAMES, - get_promptable_input_specs, - ) + from policybench.scenarios import PE_INPUT_ALIASES, PE_INPUT_RENAMES - names = {spec.source_name for spec in get_promptable_input_specs()} + names = {spec.source_name for spec in _promptable_specs()} names |= BASE_STATED_INPUTS names |= set(PE_INPUT_ALIASES.values()) | set(PE_INPUT_RENAMES.values()) return frozenset(names) @@ -734,10 +980,13 @@ def _normalize_inputs( if allow_members and name == "members": continue if v2 is not None and isinstance(name, str) and name not in registry: - resolved, note = resolve_name(name, registry, output_names=False) + resolved, note = read_fact_name(name, registry, v2["features"]) if resolved is not None: v2["notes"].append(f"{where}: {note}") name = resolved + elif (field := dropped_fact(name, registry, v2["features"])) is not None: + v2["notes"].append(_dropped_note(where, name, field)) + continue if not isinstance(name, str) or name not in registry: problems.append( f"{where}: unknown variable {name!r}." @@ -745,7 +994,7 @@ def _normalize_inputs( ) continue info = registry[name] - if v2 is not None: + if v2 is not None and OVERRIDES in v2["features"]: desired = f"{name}{DESIRED_SUFFIX}" if ( info.computed @@ -763,7 +1012,7 @@ def _normalize_inputs( "replaces its formula" ) if info.entity != entity: - if v2 is not None: + if v2 is not None and PLACEMENT in v2["features"]: v2["misplaced"].append((where, entity, name, value)) continue problems.append( @@ -782,7 +1031,10 @@ def _normalize_inputs( def normalize_request( - args, registry: Mapping[str, VariableInfo], version: int = 1 + args, + registry: Mapping[str, VariableInfo], + version: int = 1, + features: frozenset | None = None, ) -> tuple[dict | None, dict | None]: """Check a calculate_household call; return (request, None) or (None, error). @@ -795,7 +1047,9 @@ def normalize_request( ``explain``; the request's ``notes`` say what it did. """ problems: list[str] = [] - v2 = {"notes": [], "misplaced": []} if version == 2 else None + if features is None: + features = V2_FEATURES if version == 2 else frozenset() + v2 = {"notes": [], "misplaced": [], "features": features} if features else None def fail() -> tuple[None, dict]: return None, { @@ -818,9 +1072,10 @@ def fail() -> tuple[None, dict]: if not isinstance(args, dict): problems.append("arguments must be an object") return fail() - allowed = {"situation", "variables", "period"} | ( - {"explain"} if version == 2 else set() - ) + allowed = {"variables", "period"} + allowed |= {"household"} if FLAT in features else {"situation"} + allowed |= {"explain"} if EXPLAIN in features else set() + allowed |= {"detail"} if CONCISE in features else set() unknown = sorted(set(args) - allowed) if unknown: problems.append(f"unknown arguments: {', '.join(map(str, unknown))}") @@ -836,6 +1091,13 @@ def fail() -> tuple[None, dict]: ) period = str(YEAR) + if FLAT in features: + args = dict(args) + args["situation"] = flat_situation( + args.get("household"), registry, problems, v2, period + ) + if args["situation"] is None: + return fail() situation = args.get("situation") if isinstance(situation, str): if len(situation) > MAX_SITUATION_CHARS: @@ -943,11 +1205,18 @@ def fail() -> tuple[None, dict]: variables = [] elif len(variables) > MAX_VARIABLES: problems.append(f"at most {MAX_VARIABLES} variables per call") - if v2 is not None: + if v2 is not None and PLACEMENT in features: _place_misplaced(v2, normalized, situation, registry, period, problems) + if v2 is not None and LABELS in features: + _set_aliases(v2, normalized["people"], registry) requested = [] for name in variables: - if v2 is not None and isinstance(name, str) and name not in registry: + if ( + v2 is not None + and NAMES in features + and isinstance(name, str) + and name not in registry + ): resolved, note = resolve_name(name, registry, output_names=True) if resolved is not None: v2["notes"].append(f"variables: {note}") @@ -959,7 +1228,7 @@ def fail() -> tuple[None, dict]: elif name not in requested: requested.append(name) explain = [] - if v2 is not None and args.get("explain") is not None: + if EXPLAIN in features and args.get("explain") is not None: names = args["explain"] if isinstance(names, str): names = [names] @@ -969,7 +1238,7 @@ def fail() -> tuple[None, dict]: elif len(names) > MAX_EXPLAIN: problems.append(f"explain at most {MAX_EXPLAIN} variables per call") for name in names: - if name not in registry: + if NAMES in features and name not in registry: resolved, note = resolve_name(name, registry, output_names=True) if resolved is not None: v2["notes"].append(f"explain: {note}") @@ -981,15 +1250,263 @@ def fail() -> tuple[None, dict]: elif name not in explain: explain.append(name) + detail = None + if CONCISE in features: + detail = args.get("detail") or "brief" + if detail not in ("brief", "full"): + problems.append('detail must be "brief" or "full"') + if problems: return fail() request = {"situation": normalized, "variables": requested, "period": period} if v2 is not None: request["explain"] = explain request["notes"] = v2["notes"] + if detail is not None: + request["detail"] = detail return request, None +FLAT_ROLES = ("head", "spouse", "dependent") +FLAT_PERSON_KEYS = ("name", "role", "facts") +# The relationship inputs a role sets (as Scenario.to_pe_household's people +# carry them). +FLAT_RELATIONS = {"is_tax_unit_head": "head", "is_tax_unit_spouse": "spouse"} + + +def flat_names(people) -> list: + """The name each listed person goes by: the one given, or, for a person + given none, the role (head, spouse, dependent_1, ...: dependents are + numbered by their place among the dependents listed).""" + names, dependents = [], 0 + for person in people if isinstance(people, list) else []: + if not isinstance(person, dict): + names.append(None) + continue + role = person.get("role") + if role == "dependent": + dependents += 1 + name = person.get("name") + if name is None and role in FLAT_ROLES: + name = role if role != "dependent" else f"dependent_{dependents}" + names.append(name) + return names + + +def _same_value( + registry, + name: str, + a, + b, + period: str, + problems: list[str] | None = None, + where: str = "", +) -> bool | None: + """Whether two values of an input are the same once normalized, or None + when either is malformed. ``a`` stays in the situation, so normalization + reports it later; a caller that drops ``b`` passes ``problems`` so a + malformed ``b`` is reported, not lost.""" + info = registry[name] + first_problems: list[str] = [] + second_problems: list[str] = [] + first = _normalize_value(where, name, a, info, period, first_problems) + second = _normalize_value(where, name, b, info, period, second_problems) + if problems is not None: + problems.extend(second_problems) + if first_problems or second_problems: + return None + return first == second + + +def flat_situation( + household, + registry, + problems: list[str], + v2: dict | None, + period: str = str(YEAR), +): + """The flat feature's household -- state, people with role and facts, and + household facts -- as a situation, with the units the benchmark's own + reference situations have (policybench.scenarios.Scenario.to_pe_household): + the head and spouse one marital unit and everyone else alone, and one tax + unit, SPM unit, family and household holding everyone. Roles set the + relationship inputs; a stated one that contradicts its role is refused. + A household fact goes to the unit of its variable's entity, under the + name it was written with, so normalization reads and compares it as any + other input. A person given no name is named by role (flat_names). + Returns None after recording problems.""" + features = v2["features"] if v2 is not None else frozenset() + + def engine_name(name): + if isinstance(name, str) and name in registry: + return name + if isinstance(name, str): + return read_fact_name(name, registry, features)[0] + return None + + if isinstance(household, str): + if len(household) > MAX_SITUATION_CHARS: + problems.append("household is too large") + return None + try: + household = json.loads(household) + except (ValueError, RecursionError) as error: + problems.append(f"household is not valid JSON: {error}"[:MAX_MESSAGE_CHARS]) + return None + if not isinstance(household, dict): + problems.append( + "household must be a JSON object with state, people and household_facts" + ) + return None + extra = sorted(set(household) - {"state", "people", "household_facts"}) + if extra: + problems.append( + "household accepts state, people and household_facts; remove " + + ", ".join(map(str, extra)) + ) + people = household.get("people") + if not isinstance(people, list) or not people: + problems.append("household.people must be a non-empty list of people") + return None + if len(people) > MAX_PEOPLE: + problems.append(f"at most {MAX_PEOPLE} people") + return None + situation: dict[str, Any] = {"people": {}} + roles: dict[str, str] = {} + for index, (person, name) in enumerate(zip(people, flat_names(people))): + if not isinstance(person, dict): + problems.append( + f"people[{index}] must be an object with name, role and facts" + ) + continue + stray = sorted(map(str, set(person) - set(FLAT_PERSON_KEYS))) + if stray: + problems.append( + f"people[{index}] takes only name, role and facts; give " + f"{', '.join(stray)} under facts" + ) + continue + role, facts = person.get("role"), person.get("facts", {}) + if role not in FLAT_ROLES: + problems.append( + f"people[{index}].role must be one of {', '.join(FLAT_ROLES)}" + ) + continue + if person.get("name") is None and v2 is not None: + v2["notes"].append(f"people[{index}]: named {name}") + if not isinstance(name, str) or not name or len(name) > 64: + problems.append(f"people[{index}].name must be a short string") + continue + if name in roles: + problems.append(f"people[{index}]: {name!r} is listed twice") + continue + if not isinstance(facts, dict): + problems.append(f"people[{index}].facts must be an object") + continue + roles[name] = role + situation["people"][name] = dict(facts) + heads = [n for n, r in roles.items() if r == "head"] + spouses = [n for n, r in roles.items() if r == "spouse"] + if len(heads) != 1: + problems.append("the household needs exactly one head") + if len(spouses) > 1: + problems.append("the household has at most one spouse") + if problems: + return None + for name, role in roles.items(): + facts = situation["people"][name] + for relation, relation_role in FLAT_RELATIONS.items(): + if relation not in registry: + continue + implied = role == relation_role + stated = [k for k in facts if engine_name(k) == relation] + if not stated: + facts[relation] = implied + elif any( + _same_value(registry, relation, facts[k], implied, period) is False + for k in stated + ): + problems.append( + f"people.{name}: {relation} contradicts the role {role}; the " + "role sets it" + ) + everyone = list(situation["people"]) + couple = heads + spouses + groups = [couple] + [[n] for n in everyone if n not in couple] + situation["marital_units"] = { + f"marital_unit_{i}": {"members": members} + for i, members in enumerate(groups, start=1) + } + situation["tax_units"] = {"tax_unit": {"members": list(everyone)}} + situation["spm_units"] = {"spm_unit": {"members": list(everyone)}} + situation["families"] = {"family": {"members": list(everyone)}} + situation["households"] = {"household": {"members": list(everyone)}} + state = household.get("state") + if state is not None: + situation["households"]["household"]["state_code"] = state + unit_of = { + "tax_unit": situation["tax_units"]["tax_unit"], + "spm_unit": situation["spm_units"]["spm_unit"], + "family": situation["families"]["family"], + "household": situation["households"]["household"], + "marital_unit": situation["marital_units"]["marital_unit_1"], + } + facts = household.get("household_facts") or {} + if not isinstance(facts, dict): + problems.append("household_facts must be an object") + return None + for name, value in facts.items(): + resolved = engine_name(name) + if resolved is None: + field = ( + dropped_fact(name, registry, features) + if isinstance(name, str) + else None + ) + if field is not None: + v2["notes"].append(_dropped_note("household_facts", name, field)) + continue + problems.append( + f"household_facts: unknown variable {name!r}." + + (_suggest(name, registry) if isinstance(name, str) else "") + ) + continue + entity = registry[resolved].entity + if entity == "person": + # As placement moves a person variable off a one-person unit: only + # a household of one says whose fact it is. + if v2 is None or PLACEMENT not in features or len(everyone) != 1: + problems.append( + f"household_facts: {resolved} is a person's fact; give it " + "under the person whose facts list it" + ) + continue + target = situation["people"][everyone[0]] + v2["notes"].append( + f"household_facts: moved {resolved}, a person's fact, to {everyone[0]}" + ) + else: + target = unit_of[entity] + if name in target: + # The repeat is dropped, so only one validated as equal is kept quiet. + same = _same_value( + registry, + resolved, + target[name], + value, + period, + problems, + "household_facts", + ) + if same is False: + problems.append( + f"household_facts: {resolved} is set twice with different values" + ) + continue + target[name] = value + return None if problems else situation + + DESIRED_SUFFIX = "_desired" _LABEL_SEPARATORS = re.compile(r"[\s\-]+") @@ -1019,6 +1536,145 @@ def output_reading(name: str) -> tuple[str, str] | None: return engine, how +def _label_key(text: str) -> str: + """A label as compared: lower case, letters and digits only, so "Roth + 401(k) contributions" and "roth 401k contributions" read alike.""" + return re.sub(r"[^0-9a-z]", "", text.casefold()) + + +# The prompt's household block gives the state as "state: CT". +PROMPT_EXTRA_LABELS = {"state": "state_code"} + + +@functools.lru_cache(maxsize=1) +def _promptable_specs(): + from policybench.scenarios import get_promptable_input_specs + + return get_promptable_input_specs() + + +def _sample_value(value_type: str): + if value_type == "bool": + return True + if value_type in ("str", "Enum"): + return "x" + return 1.0 + + +def rendered_labels(field: str, value_type: str = "float") -> tuple[str, ...]: + """The labels a fact is shown under: in the board's prompt + (policybench.prompts._format_input_line, special renderers included) and + in the v2-pilot prompt (v2_pilot_text applied to that line).""" + from policybench.prompts import _format_input_line, _humanize_input_label + + line = _format_input_line(field, _sample_value(value_type)) + label = _humanize_input_label(field) + if not (line == f"- {label}" or line.startswith(f"- {label}: ")): + label = line[2:].split(": ", 1)[0] + labels = [label] + suffix = line[2 + len(label) :] + rewritten = v2_pilot_text(line) + if rewritten != line and rewritten.startswith("- ") and rewritten.endswith(suffix): + labels.append(rewritten[2 : len(rewritten) - len(suffix)]) + return tuple(labels) + + +@functools.lru_cache(maxsize=1) +def prompt_labels() -> dict[str, frozenset[tuple[str, bool]]]: + """Every label the prompts show a fact under, keyed by _label_key, each + with the facts it may name as (engine name, dropped): the promptable + inputs, every labeled input (policybench.prompts.INPUT_LABEL_OVERRIDES), + the inputs the reference sets besides them (BASE_STATED_INPUTS), inputs + the engine has renamed and inputs the reference leaves out + (policybench.scenarios.PE_INPUT_RENAMES, PE_DROPPED_INPUTS), each in the + board's and the v2-pilot prompt's wording. Dropped facts are shown but are + not engine inputs. A label may name facts of another country; a lookup + reads only the one its engine has.""" + from policybench.prompts import INPUT_LABEL_OVERRIDES + from policybench.scenarios import PE_DROPPED_INPUTS, PE_INPUT_RENAMES + + fields = {spec.output_name: spec.value_type for spec in _promptable_specs()} + dropped = {f for names in PE_DROPPED_INPUTS.values() for f in names} + for source in ( + "age", + *BASE_STATED_INPUTS, + *INPUT_LABEL_OVERRIDES, + *PE_INPUT_RENAMES, + *dropped, + ): + fields.setdefault(source, "int" if source == "age" else "float") + labels: dict[str, set[tuple[str, bool]]] = {} + for source, value_type in fields.items(): + target = (PE_INPUT_RENAMES.get(source, source), source in dropped) + keys = {_label_key(label) for label in rendered_labels(source, value_type)} + if source in PE_INPUT_RENAMES or source in dropped: + keys.add(_label_key(source)) + for key in keys: + labels.setdefault(key, set()).add(target) + for label, name in PROMPT_EXTRA_LABELS.items(): + labels.setdefault(_label_key(label), set()).add((name, False)) + return {key: frozenset(targets) for key, targets in labels.items()} + + +def read_fact_name( + name: str, registry: Mapping[str, VariableInfo], features: frozenset +) -> tuple[str | None, str | None]: + """A fact's engine variable under the labels feature (any label the + prompt shows it under, when exactly one fact it may name is the engine's) + and then the names feature (spaces for underscores). Returns (engine + name, note) or (None, None).""" + if LABELS in features: + found = [ + engine + for engine, dropped in prompt_labels().get(_label_key(name), ()) + if not dropped and engine in registry + ] + if len(found) == 1: + return found[0], f"read {name!r} as {found[0]}" + if NAMES in features: + return resolve_name(name, registry, output_names=False) + return None, None + + +def dropped_fact( + name: str, registry: Mapping[str, VariableInfo], features: frozenset +) -> str | None: + """Under the labels feature, the input a fact names when the prompt shows + it but the reference leaves it out (policybench.scenarios.PE_DROPPED_INPUTS: + no longer an engine input); otherwise None.""" + if LABELS not in features: + return None + found = [ + engine + for engine, dropped in prompt_labels().get(_label_key(name), ()) + if dropped + ] + return found[0] if len(found) == 1 else None + + +def _dropped_note(where: str, name: str, field: str) -> str: + return ( + f"{where}: {name!r} ({field}) is not an input of this engine; the " + "benchmark's reference leaves it out too, so it was ignored" + ) + + +def _set_aliases(v2: dict, people: dict, registry: Mapping[str, VariableInfo]): + """Set each fact the reference also sets under another name there too + (policybench.scenarios.PE_INPUT_ALIASES: the SNAP work rules read stated + usual weekly hours as weekly_hours_worked_before_lsr), unless it is set.""" + from policybench.scenarios import PE_INPUT_ALIASES + + for person, inputs in people.items(): + for source, target in PE_INPUT_ALIASES.items(): + if source in inputs and target in registry and target not in inputs: + inputs[target] = dict(inputs[source]) + v2["notes"].append( + f"people.{person}: also set {target} from {source}, as the " + "benchmark's reference does" + ) + + def resolve_name( name: str, registry: Mapping[str, VariableInfo], output_names: bool ) -> tuple[str | None, str | None]: @@ -1152,12 +1808,40 @@ def execute_request(system, registry, request: dict, simulation_cls=None) -> dic out["notes"] = list(request["notes"]) if request.get("explain"): out["explanations"] = { - name: explain_variable(system, request, name, simulation_cls) + name: ( + brief_explanation( + explain_variable(system, request, name, simulation_cls) + ) + if request.get("detail") == "brief" + else explain_variable(system, request, name, simulation_cls) + ) for name in request["explain"] } return out +BRIEF_INPUTS = 8 +BRIEF_PARAMETERS = 5 + + +def brief_explanation(full: dict) -> dict: + """The concise feature's explanation: the value, the variables it came + from directly (one level down, the first BRIEF_INPUTS), and the first + parameters it read with their sources.""" + if "error" in full: + return full + steps = full["steps"] + direct = [ + s.strip() for s in steps[1:] if s.startswith(" ") and not s.startswith(" ") + ] + return { + "value": steps[0].strip() if steps else "", + "from": direct[:BRIEF_INPUTS], + "parameters": full["parameters"][:BRIEF_PARAMETERS], + "more": 'set detail to "full" for the trace', + } + + EXPLAIN_MAX_LINES = 60 EXPLAIN_MAX_DEPTH = 6 EXPLAIN_MAX_PARAMETERS = 15 @@ -1342,6 +2026,46 @@ def build_tool_system( return CountryTaxBenefitSystem(reform=load_fix_reform(path, snapshot)) +# How a pool tells its worker the calculator's features: a tuple, which a +# model's JSON arguments can never be, so a model cannot turn on a feature its +# arm does not offer. A pool without features sends the arguments as they are. +CALL_FRAME = "calculate" + + +def tool_call_frame(features: frozenset, arguments) -> tuple: + return (CALL_FRAME, tuple(sorted(features)), arguments) + + +def frame_features(names) -> frozenset: + """A frame's features, refused (ValueError) unless they are a tuple of + known feature names with concise only beside explain.""" + if not isinstance(names, tuple) or not all(isinstance(n, str) for n in names): + raise ValueError("the features must be a tuple of names") + features = frozenset(names) + if features - set(TOOL_FEATURES): + raise ValueError( + f"unknown tool features {sorted(features - set(TOOL_FEATURES))}" + ) + if CONCISE in features and EXPLAIN not in features: + raise ValueError("the concise feature needs explain") + return features + + +def worker_request(message, registry) -> tuple[dict | None, dict | None]: + """A worker's reading of one pool message: a call frame's arguments under + its features, or anything else as a model's arguments with no features + (keys such as "__tool_features__" or "__tool_version__" in them are the + model's, and refused as unknown arguments).""" + features = frozenset() + if isinstance(message, tuple) and len(message) == 3 and message[0] == CALL_FRAME: + try: + features = frame_features(message[1]) + except ValueError as error: + return None, {"error": "configuration", "message": str(error)} + message = message[2] + return normalize_request(message, registry, features=features) + + def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> None: for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"): os.environ[name] = "1" @@ -1356,10 +2080,7 @@ def _worker_main(conn, fix_path: str, reform: bool, board: str | None = None) -> return if args is None: return - version = 1 - if isinstance(args, dict) and args.get("__tool_version__") == 2: - version, args = 2, args.get("arguments") - request, error = normalize_request(args, registry, version) + request, error = worker_request(args, registry) if error is not None: conn.send(("ok", error)) continue @@ -1379,9 +2100,15 @@ def __init__( acquire_timeout: float | None = None, board: Path | None = None, tool_version: int = 1, + tool_features: frozenset | None = None, ): self._ctx = mp.get_context("spawn") self._tool_version = tool_version + self._tool_features = ( + tool_features + if tool_features is not None + else (V2_FEATURES if tool_version == 2 else frozenset()) + ) self._fix_path = str(fix_path) # A worker gets the board only when it is not the frozen one, so a # three-argument worker target still starts on the default. @@ -1448,8 +2175,8 @@ def call(self, args) -> tuple[dict, float]: time.perf_counter() - started, ) process, conn = worker - if self._tool_version == 2: - args = {"__tool_version__": 2, "arguments": args} + if self._tool_features: + args = tool_call_frame(self._tool_features, args) try: conn.send(args) if conn.poll(self.timeout): @@ -1607,6 +2334,8 @@ def notice(self) -> str: notice = WEB_TOOL_NOTICE if tool else JSON_WEB_TOOL_NOTICE elif self.arm == "pe2": return engine_text(TOOL_NOTICE_V2 if tool else JSON_TOOL_NOTICE_V2) + elif self.arm != "pe" and arm_features(self.arm) is not None: + return engine_text(compose_tool_notice(arm_features(self.arm), tool)) else: return engine_text(TOOL_NOTICE if tool else JSON_TOOL_NOTICE) return engine_text(notice) + ( @@ -1615,11 +2344,15 @@ def notice(self) -> str: @property def uses_calculator(self) -> bool: - return self.arm in TOOL_VERSION_BY_ARM + return arm_features(self.arm) is not None + + @property + def tool_features(self) -> frozenset: + return arm_features(self.arm) or frozenset() @property def tool_version(self) -> int: - return TOOL_VERSION_BY_ARM.get(self.arm, 1) + return 2 if self.tool_features == V2_FEATURES else 1 @property def uses_web(self) -> bool: @@ -1637,7 +2370,7 @@ def arm1_label(self) -> str: def arm2_label(self) -> str: version = "" if self.prompt_version == "v1" else f"@{self.prompt_version}" rule = "+no-benchmark-search" if self.benchmark_search_rule else "" - return f"{self.name}{ARM_SUFFIX[self.arm]}{version}{rule}" + return f"{self.name}{arm_suffix(self.arm)}{version}{rule}" def provider_family(model_id: str) -> str: @@ -1663,8 +2396,11 @@ def model_profile( _uses_responses_api, ) - if arm not in ARMS: - raise SystemExit(f"unknown arm {arm!r}; the arms are {', '.join(ARMS)}") + if arm not in ARMS and arm_features(arm) is None: + raise SystemExit( + f"unknown arm {arm!r}; the arms are {', '.join(ARMS)}, or pe: plus " + f"calculator features joined by + ({', '.join(TOOL_FEATURES)})" + ) if prompt_version not in PROMPT_VERSIONS: raise SystemExit( f"unknown prompt {prompt_version!r}; the prompts are " @@ -2298,7 +3034,7 @@ def run_household( sent_board_prompt_sha256 = base.pop("_sent_board_prompt_sha256") ceiling = base.pop("_budget_ceiling") answer_tool = base.pop("_answer_tool") - calculate_tool = calculate_tool_schema(profile.tool_version) + calculate_tool = calculate_tool_schema(features=profile.tool_features) offered = ([web_search_tool(profile)] if profile.uses_web else []) + ( [calculate_tool] if profile.uses_calculator else [] ) @@ -2339,6 +3075,7 @@ def run_household( "benchmark_search_rule": profile.benchmark_search_rule, "sent_board_prompt_sha256": sent_board_prompt_sha256, "tool_notice_sha256": hashlib.sha256(profile.notice.encode()).hexdigest(), + **profile_treatment(profile), "turns": [], "status": "running", } @@ -3047,6 +3784,115 @@ def records_wording(records: Mapping[str, dict]) -> str | None: return hashes.pop() if hashes else None +# Arms whose calculator features never changed, so a record written before +# households recorded their features has these. +LEGACY_TOOL_ARMS = {"pe": frozenset(), "web+pe": frozenset(), "pe2": V2_FEATURES} + + +def profile_treatment(profile: "ModelProfile") -> dict: + """What a household record keeps of the treatment it ran under, beyond + the notice hash: the calculator's features and its schema's hash.""" + if not profile.uses_calculator: + return {"tool_features": None, "calculate_tool_sha256": None} + schema = calculate_tool_schema(features=profile.tool_features) + return { + "tool_features": sorted(profile.tool_features), + "calculate_tool_sha256": hashlib.sha256( + json.dumps(schema, sort_keys=True).encode() + ).hexdigest(), + } + + +def record_treatment_problem(record: dict, profile: "ModelProfile") -> str | None: + """Why a household record ran under another treatment than the profile's + (features, notice or calculator schema), or None. A record written before + a field existed counts as matching only for an arm whose treatment never + changed (none, web, and LEGACY_TOOL_ARMS); otherwise it is unrecorded.""" + sid = record.get("scenario_id") + want = profile_treatment(profile) + legacy = profile.arm in ("none", "web") or profile.arm in LEGACY_TOOL_ARMS + if "tool_features" in record: + if record["tool_features"] != want["tool_features"]: + return f"{sid} ran with tool features {record['tool_features']}" + elif not legacy: + return f"{sid} has no recorded tool features" + notice = record.get("tool_notice_sha256") + if notice is None and not legacy: + return f"{sid} has no recorded tool notice" + if ( + notice is not None + and notice != hashlib.sha256(profile.notice.encode()).hexdigest() + ): + return f"{sid} ran under another tool notice" + schema = record.get("calculate_tool_sha256") + if schema is not None and schema != want["calculate_tool_sha256"]: + return f"{sid} ran with another calculator schema" + return None + + +def treatment_problems( + records: Mapping[str, dict], profile: "ModelProfile" +) -> list[str]: + """record_treatment_problem for every record in a run directory, a + rerunnable one (RERUNNABLE_STATUSES) too. A rerun is not sure to replace + it: a budget stop, lost workers or a narrower --scenarios leave it in + place, and scoring reads whatever is there (review r3). So a resume and + a score both refuse a directory holding any record from another treatment.""" + return [ + problem + for record in records.values() + if (problem := record_treatment_problem(record, profile)) is not None + ] + + +def scored_treatment_problems( + records: Mapping[str, dict], profile: "ModelProfile", engine: str +) -> tuple[list[str], str | None]: + """treatment_problems as score applies it, with the note score prints. + The records' notices state the board's engine, as command_run set it. + Runs before the texts stated the board's engine (6e1578fc) named the + reference engine whatever the board's: a directory written wholly so is + one treatment and scores as it ran. A mix of the two texts does not. A + resume stays strict, since its new records would state the board's.""" + with stating_engine(engine): + problems = treatment_problems(records, profile) + if not problems or engine == REFERENCE_ENGINE: + return problems, None + with stating_engine(REFERENCE_ENGINE): + if treatment_problems(records, profile): + return problems, None + return [], ( + f"NOTE: these records' notices name policyengine-us {REFERENCE_ENGINE}, " + f"as runs before 6e1578fc stated it; the board's engine is {engine}" + ) + + +def harness_error_record( + sid: str, variables: list[str], profile: "ModelProfile", error: Exception +) -> dict: + """The record of a household the harness itself failed on: rerunnable, + with every setup and treatment key a resume and score check, so a folder + holding it still resumes and scores.""" + record = { + "scenario_id": sid, + "model": profile.name, + "arm": profile.arm, + "prompt_version": profile.prompt_version, + "prompt_wording_sha256": prompt_wording_sha256(profile), + "benchmark_search_rule": profile.benchmark_search_rule, + "tool_notice_sha256": hashlib.sha256(profile.notice.encode()).hexdigest(), + **profile_treatment(profile), + "variables": variables, + "status": "harness_error", + "error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS], + "turns": [], + "predictions": {v: None for v in variables}, + "explanations": {v: None for v in variables}, + } + record["totals"] = household_totals(record) + return record + + def records_arm(records: Mapping[str, dict]) -> str: """The one arm a run directory's records belong to; records written before the web arms are the PE tool's.""" @@ -3105,6 +3951,13 @@ def command_run(args) -> None: f"{households_dir} holds records sent under other v2-pilot wording " f"({records_wording(previous)}); use a fresh --out-dir" ) + mixed = treatment_problems(previous, profile) + if mixed: + raise SystemExit( + f"{households_dir} holds records from another treatment (" + + "; ".join(mixed[:3]) + + "); use a fresh --out-dir" + ) spend = Spend(args.budget_usd, spent_before) todo = [sid for sid in ids if sid not in done] manifest = { @@ -3133,7 +3986,10 @@ def command_run(args) -> None: "api": profile.api, "answer_contract": profile.contract, "tool_notice": profile.notice, - "calculate_tool": calculate_tool_schema(profile.tool_version) + "calculate_tool": calculate_tool_schema(features=profile.tool_features) + if profile.uses_calculator + else None, + "tool_features": sorted(profile.tool_features) if profile.uses_calculator else None, "tool_version": profile.tool_version if profile.uses_calculator else None, @@ -3164,7 +4020,12 @@ def command_run(args) -> None: return # Arms without the calculator start no engine workers. pool = ( - ToolPool(size=args.workers, board=snapshot, tool_version=profile.tool_version) + ToolPool( + size=args.workers, + board=snapshot, + tool_version=profile.tool_version, + tool_features=profile.tool_features, + ) if profile.uses_calculator else None ) @@ -3195,23 +4056,9 @@ def one(sid: str) -> dict: profile=profile, ) except Exception as error: # one household must not end the run - record = { - "scenario_id": sid, - "model": profile.name, - # The setup keys a resume and score check, so a folder with - # this record still resumes and scores. - "arm": profile.arm, - "prompt_version": profile.prompt_version, - "prompt_wording_sha256": prompt_wording_sha256(profile), - "benchmark_search_rule": profile.benchmark_search_rule, - "variables": board["variables"][sid], - "status": "harness_error", - "error": f"{type(error).__name__}: {error}"[:MAX_MESSAGE_CHARS], - "turns": [], - "predictions": {v: None for v in board["variables"][sid]}, - "explanations": {v: None for v in board["variables"][sid]}, - } - record["totals"] = household_totals(record) + record = harness_error_record( + sid, board["variables"][sid], profile, error + ) if sid in previous: prior = previous[sid] record["prior_attempts"] = prior.get("prior_attempts", []) + [ @@ -3337,6 +4184,15 @@ def command_score(args) -> None: profile = model_profile( records_model(records), records_arm(records), *records_prompt(records) ) + mixed, note = scored_treatment_problems(records, profile, board_engine(snapshot)) + if note: + print(note) + if mixed: + raise SystemExit( + f"{out / 'households'} holds records from another treatment (" + + "; ".join(mixed[:3]) + + "); scoring would mix them" + ) arm1_label, arm2_label = profile.arm1_label, profile.arm2_label complete = sorted( sid for sid, r in records.items() if r.get("status") == "complete" @@ -3879,6 +4735,12 @@ def _successful_calls(record: dict) -> list[dict]: def _person_names(arguments) -> list[str]: try: args = json.loads(arguments) if isinstance(arguments, str) else arguments + if "household" in args: # the flat feature's people list + household = args["household"] + household = ( + json.loads(household) if isinstance(household, str) else household + ) + return [name or "" for name in flat_names(household["people"])] situation = args["situation"] situation = json.loads(situation) if isinstance(situation, str) else situation return list(situation["people"]) @@ -4486,7 +5348,7 @@ def main_call(record: dict, scenario, outputs: list[str]) -> dict | None: def situation_fidelity( - records, board, system, registry, defaults, out: Path, version: int = 1 + records, board, system, registry, defaults, out: Path, features=frozenset() ) -> None: """How each household's main situation differs from the reference's. @@ -4508,7 +5370,7 @@ def situation_fidelity( if call is None: rows.append({"scenario_id": sid, "differences": None}) continue - request, _ = normalize_request(call["arguments"], registry, version) + request, _ = normalize_request(call["arguments"], registry, features=features) if request is None: rows.append({"scenario_id": sid, "differences": None}) continue @@ -4561,7 +5423,7 @@ def command_diagnose(args) -> None: out = Path(args.out_dir) records = load_records(out) # A run's calls are read as its own tool version read them. - version = TOOL_VERSION_BY_ARM.get(records_arm(records), 1) + features = arm_features(records_arm(records)) or frozenset() rows_path = out / "scored_rows.csv" if not rows_path.exists(): raise SystemExit("run score first") @@ -4577,7 +5439,7 @@ def command_diagnose(args) -> None: name: system.variables[name].default_value for name in system.variables } if registry is not None: - situation_fidelity(records, board, system, registry, defaults, out, version) + situation_fidelity(records, board, system, registry, defaults, out, features) out_rows = [] if misses.empty: print("no arm-2 misses") @@ -4595,7 +5457,7 @@ def command_diagnose(args) -> None: if registry is not None: def board_value_for(arguments, scenario=scenario, output=row.variable): - request, _ = normalize_request(arguments, registry, version) + request, _ = normalize_request(arguments, registry, features=features) if request is None: return None return output_value(system, request["situation"], scenario, output) @@ -4619,7 +5481,7 @@ def board_value_for(arguments, scenario=scenario, output=row.variable): attribution = {} arguments = verdict.get("call_arguments") if arguments is not None and registry is not None: - request, _ = normalize_request(arguments, registry, version) + request, _ = normalize_request(arguments, registry, features=features) if request is not None: reference_situation = scenario.to_pe_household() diff = situation_diff( @@ -4716,9 +5578,10 @@ def main(argv: list[str] | None = None) -> None: ) p.add_argument( "--arm", - choices=ARMS, default="pe", - help="the tools: none, PolicyEngine (pe), web search (web), or both", + help="the tools: none, PolicyEngine (pe, pe2, pe3, or pe: plus " + "calculator features joined by +), web search (web), or both " + "(web+pe)", ) p.add_argument( "--prompt", diff --git a/tests/test_tool_pilot_argon.py b/tests/test_tool_pilot_argon.py index 986a89d9..4e586374 100644 --- a/tests/test_tool_pilot_argon.py +++ b/tests/test_tool_pilot_argon.py @@ -2294,6 +2294,7 @@ def test_web_notices_name_search_and_keep_the_answer_instruction(): "web": "gpt-6-sol+web", "web+pe": "gpt-6-sol+web+pe-tool", "pe2": "gpt-6-sol+pe-tool-v2", + "pe3": "gpt-6-sol+pe-tool-v3", } with pytest.raises(SystemExit): pilot.model_profile(SOL, "search") diff --git a/tests/test_tool_v2.py b/tests/test_tool_v2.py index d6e381b5..68248090 100644 --- a/tests/test_tool_v2.py +++ b/tests/test_tool_v2.py @@ -261,8 +261,13 @@ def test_the_pool_tells_its_workers_the_tool_version(): ) try: result, _ = pool.call({"situation": "{}"}) + # A frame no model's JSON arguments can be (tool v3 review r1). assert result == { - "received": {"__tool_version__": 2, "arguments": {"situation": "{}"}} + "received": ( + "calculate", + ("explain", "names", "overrides", "placement"), + {"situation": "{}"}, + ) } finally: pool.close() diff --git a/tests/test_tool_v3.py b/tests/test_tool_v3.py new file mode 100644 index 00000000..71bc73c3 --- /dev/null +++ b/tests/test_tool_v3.py @@ -0,0 +1,1015 @@ +"""The calculator's features one at a time (tool v3 and the ablation arms). + +Each feature (names, placement, overrides, explain, concise, flat) must work +on its own and change nothing else, so a study can turn one on at a time. The +version 2 feature set must reproduce tool v2's notice byte for byte, and the +flat household format must build the units the benchmark's own reference +situations have. +""" + +from __future__ import annotations + +import itertools +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) + +import tool_pilot_argon as pilot # noqa: E402 + +V = pilot.VariableInfo +REGISTRY = { + "age": V("person", "int", "age", "year"), + "employment_income": V("person", "float", "employment income", "year"), + "is_tax_unit_head": V("person", "bool", "head", "year"), + "is_tax_unit_spouse": V("person", "bool", "spouse", "year"), + "traditional_401k_contributions": V( + "person", "float", "401(k)", "year", computed=True + ), + "traditional_401k_contributions_desired": V("person", "float", "asked", "year"), + "snap": V("spm_unit", "float", "SNAP", "month", computed=True), + "pre_subsidy_rent": V("person", "float", "rent", "year"), + "spm_unit_energy_subsidy": V("spm_unit", "float", "energy", "year"), + "qualified_passenger_vehicle_loan_interest": V( + "household", "float", "car loan", "year" + ), + "state_code": V("household", "Enum", "state", "year", ("CA", "LA")), + "tax_unit_childcare_expenses": V("tax_unit", "float", "childcare", "year"), + "marital_unit_flag": V("marital_unit", "bool", "flag", "year"), + "has_esi": V("person", "bool", "ESI", "year"), + "hours_worked_last_week": V("person", "float", "hours", "year"), + "weekly_hours_worked_before_lsr": V("person", "float", "hours", "year"), + "partnership_self_employment_net_earnings": V("person", "float", "p", "year"), +} +F = pilot + + +def _situation(): + return { + "people": {"head": {"age": 40}, "kid": {"age": 8}}, + "spm_units": {"spm": {"members": ["head", "kid"]}}, + "households": {"h": {"members": ["head", "kid"], "state_code": "CA"}}, + } + + +def _norm(args, *features): + return pilot.normalize_request(args, REGISTRY, features=frozenset(features)) + + +def test_the_version_2_set_reproduces_tool_v2s_notice(): + assert F.compose_tool_notice(F.V2_FEATURES, True) == F.TOOL_NOTICE_V2 + assert F.compose_tool_notice(F.V2_FEATURES, False) == F.JSON_TOOL_NOTICE_V2 + assert F.compose_tool_notice(frozenset(), True) == F.TOOL_NOTICE + assert F.compose_tool_notice(frozenset(), False) == F.JSON_TOOL_NOTICE + + +def test_each_feature_adds_its_own_sentence_only(): + base = F.compose_tool_notice(frozenset({F.EXPLAIN}), True) + assert "`explain`" in base and "names it read" not in base + assert "written with spaces" not in base + names = F.compose_tool_notice(frozenset({F.NAMES}), True) + assert "written with spaces" in names and "`explain`" not in names + flat = F.compose_tool_notice(frozenset({F.FLAT}), True) + assert "household you describe" in flat and "`household`" in flat + concise = F.compose_tool_notice(frozenset({F.EXPLAIN, F.CONCISE}), True) + assert '`detail` set to "full"' in concise + + +def test_arms_name_their_features(): + assert F.arm_features("pe") == frozenset() + assert F.arm_features("pe2") == F.V2_FEATURES + assert F.arm_features("pe3") == frozenset(F.TOOL_FEATURES) + assert len(F.TOOL_FEATURES) == 7 and F.LABELS in F.arm_features("pe3") + assert F.arm_features("pe:flat+names") == frozenset({"flat", "names"}) + assert F.arm_features("none") is None and F.arm_features("web") is None + assert F.arm_suffix("pe:flat+names") == "+pe-tool[names+flat]" + with pytest.raises(SystemExit, match="unknown tool features"): + F.arm_features("pe:names+telepathy") + with pytest.raises(SystemExit, match="concise feature needs explain"): + F.arm_features("pe:concise") + profile = F.ModelProfile("m", "gpt-x", "chat", "tool", arm="pe:names") + assert profile.uses_calculator and profile.tool_features == frozenset({"names"}) + assert profile.tool_version == 1 + assert "written with spaces" in profile.notice + assert profile.arm2_label == "m+pe-tool[names]" + + +def test_names_alone_reads_labels_but_places_and_redirects_nothing(): + situation = _situation() + situation["people"]["head"]["Employment income"] = 30000 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.NAMES) + assert error is None + assert request["situation"]["people"]["head"]["employment_income"] == { + "2026": 30000 + } + misplaced = _situation() + misplaced["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 5 + request, error = _norm({"situation": misplaced, "variables": ["snap"]}, F.NAMES) + assert request is None and any( + "set it under households" in p for p in error["problems"] + ) + realized = _situation() + realized["people"]["head"]["traditional_401k_contributions"] = 100 + request, error = _norm({"situation": realized, "variables": ["snap"]}, F.NAMES) + assert "traditional_401k_contributions" in request["situation"]["people"]["head"] + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, F.NAMES + ) + assert request is None and any( + "unknown arguments: explain" in p for p in error["problems"] + ) + + +def test_placement_alone_moves_inputs_but_reads_no_labels(): + situation = _situation() + situation["people"]["head"]["qualified_passenger_vehicle_loan_interest"] = 5 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.PLACEMENT) + assert error is None + assert request["situation"]["households"]["h"][ + "qualified_passenger_vehicle_loan_interest" + ] == {"2026": 5} + labeled = _situation() + labeled["people"]["head"]["employment income"] = 1 + request, error = _norm({"situation": labeled, "variables": ["snap"]}, F.PLACEMENT) + assert request is None + + +def test_overrides_alone_redirects_and_flags(): + situation = _situation() + situation["people"]["head"]["traditional_401k_contributions"] = 100 + situation["spm_units"]["spm"]["snap"] = 5 + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.OVERRIDES) + assert error is None + head = request["situation"]["people"]["head"] + assert head == { + "age": {"2026": 40}, + "traditional_401k_contributions_desired": {"2026": 100}, + } + assert any("snap is computed" in n for n in request["notes"]) + + +def test_concise_takes_a_detail_and_defaults_to_brief(): + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, + F.EXPLAIN, + F.CONCISE, + ) + assert error is None and request["detail"] == "brief" + request, _ = _norm( + { + "situation": _situation(), + "variables": ["snap"], + "explain": ["snap"], + "detail": "full", + }, + F.EXPLAIN, + F.CONCISE, + ) + assert request["detail"] == "full" + request, error = _norm( + {"situation": _situation(), "variables": ["snap"], "detail": "short"}, + F.EXPLAIN, + F.CONCISE, + ) + assert request is None + request, _ = _norm( + {"situation": _situation(), "variables": ["snap"], "explain": ["snap"]}, + F.EXPLAIN, + ) + assert "detail" not in request # the trace, as tool version 2 explains + + +def test_a_brief_explanation_keeps_the_value_its_direct_inputs_and_main_parameters(): + full = { + "steps": [ + "snap (2026) = 288", + " snap_normal_allotment (2026) = 288", + " snap_net_income (2026) = 1700", + " is_snap_eligible (2026) = True", + ], + "parameters": [{"name": f"p{i}", "value": i, "sources": []} for i in range(9)], + } + brief = F.brief_explanation(full) + assert brief["value"] == "snap (2026) = 288" + assert brief["from"] == [ + "snap_normal_allotment (2026) = 288", + "is_snap_eligible (2026) = True", + ] + assert [p["name"] for p in brief["parameters"]] == ["p0", "p1", "p2", "p3", "p4"] + assert F.brief_explanation({"error": "x"}) == {"error": "x"} + + +def test_labels_reads_the_facts_as_the_prompt_labels_them(): + situation = _situation() + situation["people"]["head"].update( + { + "gross wages and salaries": 30000, + "Has Employer-Sponsored Insurance": True, + "usual weekly hours worked": 40, + "self-employment partnership income": 5, + } + ) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None, error + head = request["situation"]["people"]["head"] + assert head["employment_income"] == {"2026": 30000} + assert head["has_esi"] == {"2026": True} + assert head["hours_worked_last_week"] == {"2026": 40} + assert head["partnership_self_employment_net_earnings"] == {"2026": 5} + # The reference also sets stated hours where the SNAP work rules read them. + assert head["weekly_hours_worked_before_lsr"] == {"2026": 40} + assert any("also set weekly_hours_worked_before_lsr" in n for n in request["notes"]) + # A name written with spaces is the names feature's, not this one's. + spaced = _situation() + spaced["people"]["head"]["employment income"] = 1 + request, error = _norm({"situation": spaced, "variables": ["snap"]}, F.LABELS) + assert request is None and "unknown variable 'employment income'" in str(error) + + +def test_labels_leaves_hours_already_set_where_the_rules_read_them(): + situation = _situation() + situation["people"]["head"].update( + {"hours_worked_last_week": 40, "weekly_hours_worked_before_lsr": 20} + ) + request, _ = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + head = request["situation"]["people"]["head"] + assert head["weekly_hours_worked_before_lsr"] == {"2026": 20} + + +def test_pe2_reads_no_labels_and_sets_no_alias_as_it_ran(): + situation = _situation() + situation["people"]["head"]["gross wages and salaries"] = 30000 + request, error = _norm( + {"situation": situation, "variables": ["snap"]}, *F.V2_FEATURES + ) + assert request is None and "gross wages and salaries" in str(error) + situation = _situation() + situation["people"]["head"]["hours_worked_last_week"] = 40 + request, _ = _norm({"situation": situation, "variables": ["snap"]}, *F.V2_FEATURES) + assert ( + "weekly_hours_worked_before_lsr" not in request["situation"]["people"]["head"] + ) + + +def test_only_pe2_says_the_facts_carry_variable_names(): + claim = "labeled with its variable names" + assert claim in F.TOOL_NOTICE_V2 + for size in range(1, len(F.TOOL_FEATURES) + 1): + for features in itertools.combinations(F.TOOL_FEATURES, size): + features = frozenset(features) + if F.CONCISE in features and F.EXPLAIN not in features: + continue + notice = F.compose_tool_notice(features, True) + assert (claim in notice) == (features == F.V2_FEATURES), features + assert notice.endswith("as instructed above.") + assert ("names it read" in notice) == bool(features & {F.NAMES, F.LABELS}) + + +def test_every_label_a_prompt_shows_reads_back_to_its_engine_name(): + from policybench.prompts import _format_input_line + from policybench.scenarios import PE_DROPPED_INPUTS, PE_INPUT_RENAMES + + labels = F.prompt_labels() + fields = {spec.output_name: spec.value_type for spec in F._promptable_specs()} + fields |= {"age": "int", "employment_income": "float"} + dropped = {f for names in PE_DROPPED_INPUTS.values() for f in names} + for field, value_type in {**fields, **dict.fromkeys(dropped, "float")}.items(): + for label in F.rendered_labels(field, value_type): + want = (PE_INPUT_RENAMES.get(field, field), field in dropped) + assert want in labels[F._label_key(label)], (field, label) + # Fallback labels (the name with spaces) are the prompt's too. + assert ("first_home_mortgage_balance", False) in labels[ + F._label_key("first home mortgage balance") + ] + # Special renderers: the Marketplace plan line is labeled on its own. + line = _format_input_line("selected_marketplace_plan_benchmark_ratio", 0.8) + assert line.startswith("- selected Marketplace plan: ") + assert ("selected_marketplace_plan_benchmark_ratio", False) in labels[ + F._label_key("selected Marketplace plan") + ] + # The v2-pilot prompt's relabelings, as that prompt writes them. + for field, label in ( + ("spm_unit_energy_subsidy", "home energy assistance received"), + ( + "employer_sponsored_insurance_premiums", + "employer-paid health insurance premiums (paid by the employer, " + "not by the household)", + ), + ): + assert label in F.rendered_labels(field) + assert (field, False) in labels[F._label_key(label)] + assert ("state_code", False) in labels[F._label_key("state")] + assert ("first_home_mortgage_interest", True) in labels[ + F._label_key("first home mortgage interest") + ] + assert ("roth_401k_contributions_desired", False) in labels[ + F._label_key("Roth 401k contributions desired") + ] + assert ("partnership_s_corp_income", False) in labels[ + F._label_key("partnership or S-corp income") + ] + + +def test_labels_reads_no_label_that_names_two_of_the_engines_inputs(monkeypatch): + two = frozenset({("has_esi", False), ("age", False)}) + monkeypatch.setattr(F, "prompt_labels", lambda: {F._label_key("covered"): two}) + situation = _situation() + situation["people"]["head"]["covered"] = True + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert request is None and "unknown variable 'covered'" in str(error) + one = frozenset({("has_esi", False), ("uk_only_input", False)}) + monkeypatch.setattr(F, "prompt_labels", lambda: {F._label_key("covered"): one}) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None and request["situation"]["people"]["head"]["has_esi"] + + +def test_labels_ignores_a_fact_the_reference_leaves_out(): + situation = _situation() + situation["tax_units"] = { + "t": {"members": ["head", "kid"], "first home mortgage interest": 900} + } + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.LABELS) + assert error is None, error + assert request["situation"]["tax_units"]["t"] == {"members": ["head", "kid"]} + assert any("leaves it out too" in n for n in request["notes"]) + request, error = _norm({"situation": situation, "variables": ["snap"]}, F.NAMES) + assert request is None and "first home mortgage interest" in str(error) + + +def _flat(**overrides): + household = { + "state": "LA", + "people": [ + { + "name": "Head", + "role": "head", + "facts": {"age": 40, "employment_income": 30000}, + }, + {"name": "Spouse", "role": "spouse", "facts": {"age": 38}}, + {"name": "Kid", "role": "dependent", "facts": {"age": 6}}, + ], + "household_facts": { + "spm_unit_energy_subsidy": 900, + "tax_unit_childcare_expenses": 1200, + "qualified_passenger_vehicle_loan_interest": 300, + }, + } + household.update(overrides) + return {"household": json.dumps(household), "variables": ["snap"]} + + +def test_flat_builds_the_benchmarks_units(): + request, error = _norm(_flat(), F.FLAT) + assert error is None, error + s = request["situation"] + assert s["marital_units"] == { + "marital_unit_1": {"members": ["Head", "Spouse"]}, + "marital_unit_2": {"members": ["Kid"]}, + } + for group, unit in ( + ("tax_units", "tax_unit"), + ("spm_units", "spm_unit"), + ("families", "family"), + ): + assert s[group][unit]["members"] == ["Head", "Spouse", "Kid"] + assert s["households"]["household"]["state_code"] == {"2026": "LA"} + assert s["spm_units"]["spm_unit"]["spm_unit_energy_subsidy"] == {"2026": 900} + assert s["tax_units"]["tax_unit"]["tax_unit_childcare_expenses"] == {"2026": 1200} + assert s["households"]["household"][ + "qualified_passenger_vehicle_loan_interest" + ] == {"2026": 300} + assert s["people"]["Head"]["is_tax_unit_head"] == {"2026": True} + assert s["people"]["Spouse"]["is_tax_unit_spouse"] == {"2026": True} + assert s["people"]["Kid"]["is_tax_unit_head"] == {"2026": False} + + +def test_flat_names_a_person_by_role_when_no_name_is_given(): + people = [ + {"role": "head", "facts": {"age": 40}}, + {"role": "dependent", "facts": {"age": 6}}, + {"name": "Kid", "role": "dependent", "facts": {"age": 4}}, + {"role": "dependent", "facts": {"age": 2}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert error is None, error + assert list(request["situation"]["people"]) == [ + "head", + "dependent_1", + "Kid", + "dependent_3", + ] + assert "people[0]: named head" in request["notes"] + people = [ + {"role": "head", "facts": {}}, + {"name": "head", "role": "dependent", "facts": {}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("listed twice" in p for p in error["problems"]) + people = [{"name": "", "role": "head", "facts": {}}] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("name must be a short string" in p for p in error["problems"]) + + +def test_flat_placement_gives_a_household_of_ones_person_fact_to_that_person(): + one = [{"name": "Solo", "role": "head", "facts": {"age": 70}}] + args = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, error + solo = request["situation"]["people"]["Solo"] + assert solo["pre_subsidy_rent"] == {"2026": 9000} + assert "pre_subsidy_rent" not in request["situation"]["spm_units"]["spm_unit"] + assert any("moved pre_subsidy_rent" in n for n in request["notes"]) + one[0]["facts"]["pre_subsidy_rent"] = 100 + conflict = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(conflict, F.FLAT, F.PLACEMENT) + assert any("different value" in p for p in error["problems"]) + # Two people: whose rent it is is not said, so it is refused. + request, error = _norm( + _flat(household_facts={"pre_subsidy_rent": 9000}), F.FLAT, F.PLACEMENT + ) + assert any("whose facts list it" in p for p in error["problems"]) + # Without placement, even a household of one is refused. + request, error = _norm(args, F.FLAT) + assert any("whose facts list it" in p for p in error["problems"]) + + +def test_flat_refuses_facts_given_beside_the_person(): + people = [{"role": "head", "age": 40, "employment_income": 30000}] + request, error = _norm(_flat(people=people), F.FLAT) + assert request is None + assert any( + "takes only name, role and facts; give age, employment_income under facts" in p + for p in error["problems"] + ) + + +def test_flat_reads_a_fact_written_two_ways_once_whatever_the_order(): + for facts in ( + {"spm_unit_energy_subsidy": 900, "spm unit energy subsidy": 800}, + {"spm unit energy subsidy": 800, "spm_unit_energy_subsidy": 900}, + ): + request, error = _norm(_flat(household_facts=facts), F.FLAT, F.NAMES) + assert request is None + assert any("set twice with different values" in p for p in error["problems"]) + same = {"spm_unit_energy_subsidy": 900, "spm unit energy subsidy": {"2026": 900}} + request, error = _norm(_flat(household_facts=same), F.FLAT, F.NAMES) + assert error is None, error + + +def test_flat_roles_set_the_relationship_inputs_and_refuse_a_contradiction(): + for written in ("is_tax_unit_head", "is tax unit head"): + people = [{"name": "A", "role": "head", "facts": {written: False}}] + request, error = _norm(_flat(people=people), F.FLAT, F.NAMES) + assert request is None, written + assert any("contradicts the role head" in p for p in error["problems"]) + people = [{"name": "A", "role": "head", "facts": {written: True}}] + request, error = _norm(_flat(people=people), F.FLAT, F.NAMES) + assert error is None, (written, error) + head = request["situation"]["people"]["A"] + assert head["is_tax_unit_head"] == {"2026": True} + assert head["is_tax_unit_spouse"] == {"2026": False} + + +def test_flat_placement_compares_values_once_normalized(): + one = [ + {"name": "Solo", "role": "head", "facts": {"pre_subsidy_rent": {"2026": 9000}}} + ] + args = _flat(people=one, household_facts={"pre_subsidy_rent": 9000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, error + assert request["situation"]["people"]["Solo"]["pre_subsidy_rent"] == {"2026": 9000} + + +def test_the_descriptions_promise_only_what_the_features_do(): + flat = F.calculate_tool_schema(features=frozenset({F.FLAT}))["function"] + assert '"employment_income": 30000' in flat["description"] + example = [ + {"role": "head", "facts": {"age": 40, "employment_income": 30000}}, + ] + request, error = _norm(_flat(people=example), F.FLAT) + assert error is None, error + entity_error = "a variable set on the wrong entity" + placed = F.calculate_tool_schema(features=frozenset({F.PLACEMENT}))["function"] + assert entity_error not in placed["description"] + assert "moved to the one unit holding its people" in placed["description"] + named = F.calculate_tool_schema(features=frozenset({F.NAMES}))["function"] + assert entity_error in named["description"] + concise = F.calculate_tool_schema(features=frozenset({F.EXPLAIN, F.CONCISE})) + detail = concise["function"]["parameters"]["properties"]["detail"]["description"] + assert f"up to {F.EXPLAIN_MAX_LINES} steps" in detail + for features in (frozenset({F.EXPLAIN, F.CONCISE}), frozenset(F.TOOL_FEATURES)): + assert "every step" not in F.compose_tool_notice(features, True) + + +def test_flat_refuses_what_it_cannot_place(): + people = [ + {"name": "A", "role": "head", "facts": {}}, + {"name": "B", "role": "head", "facts": {}}, + ] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("exactly one head" in p for p in error["problems"]) + request, error = _norm(_flat(household_facts={"age": 3}), F.FLAT) + assert any("person's fact" in p for p in error["problems"]) + people = [{"name": "A", "role": "boss", "facts": {}}] + request, error = _norm(_flat(people=people), F.FLAT) + assert any("role must be one of" in p for p in error["problems"]) + request, error = _norm({"situation": _situation(), "variables": ["snap"]}, F.FLAT) + assert any("unknown arguments: situation" in p for p in error["problems"]) + request, error = _norm(_flat(), F.NAMES) + assert any("unknown arguments: household" in p for p in error["problems"]) + + +def test_flat_with_names_reads_labeled_household_facts(): + request, error = _norm( + _flat(household_facts={"spm unit energy subsidy": 900}), F.FLAT, F.NAMES + ) + assert error is None + spm_unit = request["situation"]["spm_units"]["spm_unit"] + assert spm_unit["spm_unit_energy_subsidy"] == {"2026": 900} + assert any("as spm_unit_energy_subsidy" in n for n in request["notes"]) + + +def test_the_schema_offers_each_features_arguments(): + flat = F.calculate_tool_schema(features=frozenset({F.FLAT}))["function"] + props = flat["parameters"]["properties"] + assert "household" in props and "situation" not in props + assert flat["parameters"]["required"] == ["household", "variables"] + concise = F.calculate_tool_schema(features=frozenset({F.EXPLAIN, F.CONCISE})) + props = concise["function"]["parameters"]["properties"] + assert "explain" in props and props["detail"]["enum"] == ["brief", "full"] + assert F.calculate_tool_schema(features=F.V2_FEATURES) == F.calculate_tool_schema(2) + assert F.calculate_tool_schema(features=frozenset()) == F.calculate_tool_schema(1) + + +def _echo_worker(conn, fix_path, reform): + conn.send(("ready", None)) + while True: + args = conn.recv() + if args is None: + return + conn.send(("ok", {"received": args})) + + +def test_the_pool_tells_its_workers_the_features_in_a_frame(): + pool = F.ToolPool( + size=1, + timeout=5.0, + worker_target=_echo_worker, + tool_features=frozenset({F.FLAT, F.NAMES}), + ) + try: + result, _ = pool.call({"x": 1}) + assert result == {"received": ("calculate", ("flat", "names"), {"x": 1})} + finally: + pool.close() + + +def test_a_model_cannot_turn_on_a_feature_its_arm_does_not_offer(): + forged = { + "__tool_features__": ["names"], + "arguments": { + "situation": {"people": {"head": {"Employment income": 12345}}}, + "variables": ["age"], + }, + } + for message in (forged, {"__tool_version__": 2, "arguments": forged["arguments"]}): + request, error = F.worker_request(message, REGISTRY) + assert request is None + assert "unknown arguments: " in error["problems"][0] + request, error = F.worker_request( + F.tool_call_frame(frozenset({F.NAMES}), forged["arguments"]), REGISTRY + ) + assert error is None + assert request["situation"]["people"]["head"]["employment_income"] == { + "2026": 12345 + } + for names in (("telepathy",), ("concise",), ({},), ["names"]): + request, error = F.worker_request(("calculate", names, {}), REGISTRY) + assert request is None and error["error"] == "configuration", names + + +def _profile(arm): + return F.ModelProfile("m", "gpt-x", "chat", "tool", arm=arm) + + +def _record(profile, **changes): + record = { + "scenario_id": "mt_000", + "tool_notice_sha256": __import__("hashlib") + .sha256(profile.notice.encode()) + .hexdigest(), + **F.profile_treatment(profile), + } + record.update(changes) + return record + + +def test_a_resume_refuses_households_run_under_another_treatment(): + pe3 = _profile("pe3") + assert F.record_treatment_problem(_record(pe3), pe3) is None + six = sorted(set(F.TOOL_FEATURES) - {F.LABELS}) + assert "tool features" in F.record_treatment_problem( + _record(pe3, tool_features=six), pe3 + ) + assert "another tool notice" in F.record_treatment_problem( + _record(pe3, tool_notice_sha256="0" * 64), pe3 + ) + assert "calculator schema" in F.record_treatment_problem( + _record(pe3, calculate_tool_sha256="0" * 64), pe3 + ) + old = _record(pe3) + del old["tool_features"] + assert "no recorded tool features" in F.record_treatment_problem(old, pe3) + # Arms whose treatment never changed accept records from before the fields. + for arm in ("pe", "pe2", "web+pe", "none", "web"): + profile = _profile(arm) + legacy = _record(profile) + for key in ("tool_features", "calculate_tool_sha256", "tool_notice_sha256"): + legacy.pop(key) + assert F.record_treatment_problem(legacy, profile) is None, arm + pe2 = _profile("pe2") + assert "tool features" in F.record_treatment_problem(_record(_profile("pe")), pe2) + + +def test_diagnostics_read_flat_people_by_their_default_names(): + household = { + "people": [ + {"role": "head", "facts": {}}, + {"role": "spouse", "facts": {}}, + {"name": "Kid", "role": "dependent", "facts": {}}, + {"role": "dependent", "facts": {}}, + ] + } + names = F._person_names({"household": json.dumps(household)}) + assert names == ["head", "spouse", "Kid", "dependent_2"] + + +@pytest.mark.slow +def test_a_household_described_flat_computes_what_the_benchmarks_own_does(): + """Differential: on the installed engine, each frozen-board household + described in the flat format computes every board output exactly as the + benchmark's own situation (Scenario.to_pe_household) does.""" + import household_set as hs + import pandas as pd + from policyengine_us import Simulation + + from policybench.ground_truth import _extract_person_value, _pe_variable_for_output + + system = F.build_tool_system(reform=True) + registry = F.build_registry(system) + reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") + variables = reference.groupby("scenario_id").variable.agg(list) + relations = {"is_tax_unit_head", "is_tax_unit_spouse"} + checked = 0 + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:8]: + situation = hs.situation_for(scenario) + couple = scenario.marital_couple() or () + head = couple[0] if couple else scenario.adults[0].name + people = [ + { + "name": name, + "role": "head" + if name == head + else ("spouse" if name in couple else "dependent"), + "facts": {k: v for k, v in facts.items() if k not in relations}, + } + for name, facts in situation["people"].items() + ] + household_facts = {} + for group in ("tax_units", "spm_units", "families", "households"): + for unit in situation[group].values(): + household_facts.update( + { + k: v + for k, v in unit.items() + if k not in ("members", "state_code") + } + ) + household = { + "state": scenario.state, + "people": people, + "household_facts": household_facts, + } + request, error = F.normalize_request( + {"household": json.dumps(household), "variables": ["snap"]}, + registry, + features=frozenset({F.FLAT}), + ) + assert error is None, (scenario.id, error) + own = Simulation(tax_benefit_system=system, situation=situation) + flat = Simulation(tax_benefit_system=system, situation=request["situation"]) + for variable in variables[scenario.id]: + name = _pe_variable_for_output(variable, "us") + values = [ + float( + _extract_person_value( + simulation.calculate(name, 2026), scenario, variable + ) + ) + for simulation in (own, flat) + ] + assert values[0] == values[1], (scenario.id, variable, values) + checked += 1 + assert checked > 80 + + +@pytest.mark.slow +def test_a_household_described_from_its_prompt_computes_the_same(): + """Differential, with the labels feature alone: each of 30 frozen-board + households described flat from the facts its v2-pilot prompt shows, each + under the label the prompt shows it under (dropped facts included), with + the state and roles the prompt gives, computes every board output as the + benchmark's own situation does. The take-up conventions, which the prompt + states in words, are set under their engine names.""" + import household_set as hs + import pandas as pd + from policyengine_us import Simulation + + from policybench.ground_truth import _extract_person_value, _pe_variable_for_output + from policybench.scenarios import ( + DEFAULT_TAKEUP_INPUTS, + is_excluded_prompt_input_name, + ) + + system = F.build_tool_system(reform=True) + registry = F.build_registry(system) + reference = pd.read_csv(hs.FROZEN_BOARD / "reference_outputs.csv") + variables = reference.groupby("scenario_id").variable.agg(list) + shown = set() + + def label(field, value): + value_type = ( + "bool" + if isinstance(value, bool) + else ("str" if isinstance(value, str) else "float") + ) + written = F.rendered_labels(field, value_type)[-1] + shown.add((field, written)) + return written + + def prompt_facts(inputs): + return { + label(k, v): v + for k, v in inputs.items() + if not is_excluded_prompt_input_name(k) + } + + checked = 0 + for scenario in hs.load_scenarios(hs.FROZEN_BOARD)[:30]: + couple = scenario.marital_couple() or () + head = couple[0] if couple else scenario.adults[0].name + people = [] + for person in scenario.all_people: + facts = {"age": person.age} + if person.name.startswith("adult") or abs(person.employment_income) > 1e-6: + facts[label("employment_income", 1.0)] = person.employment_income + facts |= prompt_facts(person.inputs) + facts |= DEFAULT_TAKEUP_INPUTS["person"] + role = ( + "head" + if person.name == head + else ("spouse" if person.name in couple else "dependent") + ) + people.append({"name": person.name, "role": role, "facts": facts}) + household_facts = {} + for inputs in ( + scenario.tax_unit_inputs, + scenario.spm_unit_inputs, + scenario.household_inputs, + ): + household_facts |= prompt_facts(inputs) + for entity in ("tax_unit", "spm_unit", "household"): + household_facts |= DEFAULT_TAKEUP_INPUTS[entity] + household = { + "state": scenario.state, + "people": people, + "household_facts": household_facts, + } + request, error = F.normalize_request( + {"household": json.dumps(household), "variables": ["snap"]}, + registry, + features=frozenset({F.FLAT, F.LABELS}), + ) + assert error is None, (scenario.id, error) + own = Simulation( + tax_benefit_system=system, situation=hs.situation_for(scenario) + ) + flat = Simulation(tax_benefit_system=system, situation=request["situation"]) + for variable in variables[scenario.id]: + name = _pe_variable_for_output(variable, "us") + values = [ + float( + _extract_person_value( + simulation.calculate(name, 2026), scenario, variable + ) + ) + for simulation in (own, flat) + ] + assert values[0] == values[1], (scenario.id, variable, values) + checked += 1 + worded = {w for f, w in shown if w != f.replace("_", " ")} + assert checked > 300 and len(worded) > 10, (checked, sorted(worded)) + # No label a US prompt can show names two of this engine's inputs. + ambiguous = { + key: targets + for key, targets in F.prompt_labels().items() + if len({e for e, dropped in targets if not dropped and e in registry}) > 1 + } + assert not ambiguous, ambiguous + + +def test_flat_reports_a_malformed_repeat_of_a_fact_it_drops(): + # Review r2: a repeat that fails validation was dropped without a word. + one = [{"name": "Solo", "role": "head", "facts": {"employment_income": 30000}}] + for bad in ("40000", [40000], {"26": 40000}, {}): + args = _flat(people=one, household_facts={"employment_income": bad}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None, bad + assert any("employment_income" in m for m in error["problems"]), (bad, error) + # An equal repeat, scalar or by period, still merges quietly; a different + # one is still refused. + for same in (30000, {"2026": 30000}): + args = _flat(people=one, household_facts={"employment_income": same}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert error is None, (same, error) + args = _flat(people=one, household_facts={"employment_income": 40000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None + assert any("set twice with different values" in m for m in error["problems"]) + + +def test_a_harness_error_record_resumes_under_its_own_arm(): + # Review r2: the fallback record carried no treatment, so one harness + # failure blocked every custom-arm resume. + for arm in ("pe3", "pe:labels+flat", "pe", "pe2", "web+pe", "none", "web"): + profile = _profile(arm) + record = F.harness_error_record("mt_000", ["snap"], profile, RuntimeError("x")) + assert record["status"] in F.RERUNNABLE_STATUSES + assert F.record_treatment_problem(record, profile) is None, arm + assert record["predictions"] == {"snap": None} + assert record["totals"]["cost_usd"] == 0 + + +def test_a_rerunnable_record_from_another_treatment_is_refused_too(): + # Review r3: a rerun is not sure to replace a rerunnable record (a budget + # stop, lost workers, a narrower --scenarios), so every record is checked. + pe3 = _profile("pe3") + stale = _record(pe3, tool_features=["names"]) + records = { + "mt_000": dict(stale, status="complete"), + "mt_001": dict(stale, scenario_id="mt_001", status="provider_error"), + "mt_002": F.harness_error_record("mt_002", ["snap"], pe3, RuntimeError("x")), + } + problems = F.treatment_problems(records, pe3) + assert [p.split()[0] for p in problems] == ["mt_000", "mt_001"] + + +SIX = ("concise", "explain", "flat", "names", "overrides", "placement") + + +def _scored_record(sid, profile, treatment_of=None, status="complete", **changes): + """A household record as command_run writes it, with the treatment of + `treatment_of` (another arm's profile) when given.""" + import hashlib + + board = pilot.load_board() + reference = board["reference"].set_index(["scenario_id", "variable"])["value"] + variables = board["variables"][sid] + source = treatment_of or profile + record = { + "scenario_id": sid, + "model": profile.name, + "arm": profile.arm, + "prompt_version": profile.prompt_version, + "benchmark_search_rule": profile.benchmark_search_rule, + "tool_notice_sha256": hashlib.sha256(source.notice.encode()).hexdigest(), + **F.profile_treatment(source), + "variables": variables, + "predictions": {v: float(reference[(sid, v)]) for v in variables}, + "explanations": {v: "x" for v in variables}, + "status": status, + "turns": [{"total_cost_usd": 1.25, "calls": []}], + "cap_hit": False, + "wall_seconds": 1.0, + } + record.update(changes) + record["totals"] = F.household_totals(record) + return record + + +def _write(directory, *records): + (directory / "households").mkdir(parents=True, exist_ok=True) + for record in records: + path = directory / "households" / f"{record['scenario_id']}.json" + path.write_text(json.dumps(record)) + + +def test_a_resume_refuses_an_old_treatments_provider_error_before_running( + tmp_path, monkeypatch +): + # The review's case: a pe3 folder holding a provider_error written under + # six features, with the budget already spent. Before the fix the resume + # passed, `one` stopped on the budget without replacing it, and scoring + # read its predictions as pe3's. + import argparse + + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + six = F.model_profile(F.MODEL_NAME, "pe:" + "+".join(SIX)) + _write(tmp_path, _scored_record("scenario_000", pe3, six, status="provider_error")) + monkeypatch.setattr(F, "_require_engine", lambda board: None) + monkeypatch.setenv(F.provider_key(pe3.model_id), "test-key") + + def no_pool(**kwargs): + raise AssertionError("the run got as far as starting the calculator") + + monkeypatch.setattr(F, "ToolPool", no_pool) + args = argparse.Namespace( + out_dir=str(tmp_path), + board=None, + model=F.MODEL_NAME, + arm="pe3", + prompt="v1", + forbid_benchmark_search=False, + scenarios=None, + limit=None, + workers=1, + budget_usd=1.0, + reserve_usd=0.0, + ) + with pytest.raises(SystemExit, match="another treatment"): + F.command_run(args) + assert not list(tmp_path.glob("manifest-*.json")) + + +def test_score_refuses_a_folder_mixing_treatments(tmp_path): + import argparse + + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + six = F.model_profile(F.MODEL_NAME, "pe:" + "+".join(SIX)) + _write( + tmp_path, + _scored_record("scenario_000", pe3), + _scored_record("scenario_001", pe3, six, status="provider_error"), + ) + args = argparse.Namespace(out_dir=str(tmp_path), complete_only=False) + with pytest.raises(SystemExit, match="another treatment"): + F.command_score(args) + assert not (tmp_path / "score.json").exists() + # The same folder without the stray record scores. + (tmp_path / "households" / "scenario_001.json").unlink() + F.command_score(args) + assert json.loads((tmp_path / "score.json").read_text())["households"] == 1 + + +def test_score_rebuilds_notices_on_the_boards_engine_and_restores_it( + tmp_path, monkeypatch +): + # A household set's records state its engine (2.60.4) in their notices and + # schema, so a check on the reference engine's text would refuse them all. + import argparse + + monkeypatch.setattr(F, "board_engine", lambda snapshot=None: "2.60.4") + pe3 = F.model_profile(F.MODEL_NAME, "pe3") + with F.stating_engine("2.60.4"): + current = _scored_record("scenario_000", pe3) + _write(tmp_path, current) + args = argparse.Namespace(out_dir=str(tmp_path), complete_only=False) + F.command_score(args) + assert F._STATED_ENGINE["version"] == F.REFERENCE_ENGINE + # Beside it, a record sent under the reference engine's text mixes two. + _write(tmp_path, _scored_record("scenario_001", pe3)) + with pytest.raises(SystemExit, match="another tool notice"): + F.command_score(args) + assert F._STATED_ENGINE["version"] == F.REFERENCE_ENGINE + + +def test_score_takes_a_folder_written_wholly_before_the_texts_stated_the_engine( + tmp_path, monkeypatch, capsys +): + # Test batch 0's frontier calculator arms (b6d835a8) were told 2.41.0 while + # the board and calculator ran 2.60.4: one treatment, scored as it ran. A + # resume into it is still refused, since new records would say 2.60.4. + import argparse + + monkeypatch.setattr(F, "board_engine", lambda snapshot=None: "2.60.4") + pe = F.model_profile(F.MODEL_NAME, "pe") + _write( + tmp_path, _scored_record("scenario_000", pe), _scored_record("scenario_001", pe) + ) + F.command_score(argparse.Namespace(out_dir=str(tmp_path), complete_only=False)) + assert "name policyengine-us 2.41.0" in capsys.readouterr().out + assert json.loads((tmp_path / "score.json").read_text())["households"] == 2 + with F.stating_engine("2.60.4"): + assert len(F.treatment_problems(F.load_records(tmp_path), pe)) == 2 + + +def test_flat_reports_a_malformed_first_value_once_and_not_as_a_conflict(): + one = [{"name": "Solo", "role": "head", "facts": {"employment_income": "30000"}}] + args = _flat(people=one, household_facts={"employment_income": 40000}) + request, error = _norm(args, F.FLAT, F.PLACEMENT) + assert request is None + assert not any("set twice" in m for m in error["problems"]), error + assert sum("employment_income" in m for m in error["problems"]) == 1, error