Skip to content

Ready the tool pilot's v2-pilot run: prompt fixes, a no-tools arm, web timeout #239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MaxGhenis merged 2 commits into from
Oct 22, 2026
Merged

Ready the tool pilot's v2-pilot run: prompt fixes, a no-tools arm, web timeout #239

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
01b077f
Ready the tool pilot's v2-pilot run: prompt fixes, a no-tools arm, we…
MaxGhenis Oct 22, 2026
274527b
Address review r1 of the v2-pilot run prep
MaxGhenis Oct 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 137 additions & 21 deletions scripts/tool_pilot_argon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,9 +147,25 @@

# The arms ``run --arm`` takes. "pe" is arm 2 as run on 2026-10-14 and
# 2026-10-15. "web" swaps the PolicyEngine tool for the provider's own web
# search, and "web+pe" offers both.
ARMS = ("pe", "web", "web+pe")
ARM_SUFFIX = {"pe": "+pe-tool", "web": "+web", "web+pe": "+web+pe-tool"}
# search, and "web+pe" offers both. "none" offers neither: its first request is
# the board's own no-tools request (the forced answer function, or the JSON reply
# 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")
ARM_SUFFIX = {
"none": "+no-tools",
"pe": "+pe-tool",
"web": "+web",
"web+pe": "+web+pe-tool",
}
# A web arm's request timeout. A search-heavy turn can outlast the board's
# (300 s for Gemini 4 Argon, the harness default for its card; its web-only arm
# never finished scenario_005 within it in six attempts on 2026-10-21), so the
# web arms allow this much. A timed-out web request is not retried: three
# 900 s attempts would hold a worker for 45 minutes, so the household ends as
# provider_error and a resume pass reruns it.
WEB_REQUEST_TIMEOUT_SECONDS = 900

_CALCULATE_OPENING = "You also have a `calculate_household` function."
WEB_SENTENCE = "You also have web search, which searches the public web."
Expand Down Expand Up @@ -275,6 +291,60 @@
"{who}. Where a credit depends on income tax liability, compute that "
"liability from the section 1 tax rate schedules, not the Tax Table",
),
(
# policyengine-us 2.41.0's state_income_tax_before_refundable_credits
# list adds id_pbf for Idaho ("Form 40 permanent building fund tax,
# added before refundable credits"); no other state adds such a tax.
# GPT-6 Sol's tool answers for scenario_007 were $10 short without it.
re.compile(
r"(state individual income tax after nonrefundable credits and before "
r"refundable credits, excluding local income and payroll taxes)"
),
"{who}; include Idaho's $10 permanent building fund tax where it applies",
),
(
# The reference rule (reference sidecar, quoted in each
# reference_audit/2026-09-28/fixes/latest_c_* module): "a scored
# reference follows from the stated facts and from law published before
# the 2026-07-03 reference freeze". Where no 2026 amount was published by
# then, the reference holds the last one published before it: California's
# 2025 indexed amounts, Idaho's zero-rate thresholds, the IRS 2025
# optional sales tax tables, and some Maryland, Michigan and Minnesota
# amounts. SNAP holds the FY2026 schedule (allotments, deductions,
# utility allowances, asset limits) for October to December; "Not held:
# the poverty guideline", so income limits move to the 2026 guideline in
# October. Web arms had applied SNAP's FY2027 schedule, published in
# August 2026 (GPT-6 Sol, scenarios 008 and 012).
re.compile(
r"(Do not infer unlisted income, expenses, assets, benefit receipt, "
r"rent, or health coverage\.)"
),
"{who} Apply the law and program parameters published before July 3, "
"2026. Where a 2026 amount had not been published by then, use the latest "
"amount published before that date (for example, a state's 2025 "
"inflation-adjusted amounts, or the IRS's 2025 optional sales tax "
"tables). For SNAP, keep the fiscal year 2026 allotments, deductions, "
"utility allowances and asset limits through December 2026; the income "
"limits use the 2026 HHS poverty guidelines from October 2026.",
),
(
# policyengine-us documents employer_sponsored_insurance_premiums as
# "Annual employer-paid health insurance premiums". Three web arms read
# the bare label as premiums the household paid (scenario_000).
re.compile(r"employer sponsored insurance premiums: "),
"employer-paid health insurance premiums (paid by the employer, not by "
"the household): ",
),
(
# spm_unit_energy_subsidy, the SPM's energy subsidy (home energy
# assistance). No scored output reads it. policyengine-us never reads the
# amount for SNAP: it deems a utility allowance from the state's
# always-SUA flag, and in 2026 only for elderly or disabled households
# (heat_and_eat_requires_elderly_disabled, P.L. 119-21 section 10103).
# The engine label hid what the amount is.
re.compile(r"spm unit energy subsidy: "),
"home energy assistance received: ",
),
)

# What a JSON-contract web arm hears once its search-only and calculation turns
Expand Down Expand Up @@ -1055,6 +1125,8 @@ class ModelProfile:
@property
def notice(self) -> str:
tool = self.contract == "tool"
if self.arm == "none":
return ""
if self.arm == "web":
notice = WEB_NOTICE if tool else JSON_WEB_NOTICE
elif self.arm == "web+pe":
Expand All @@ -1065,11 +1137,11 @@ def notice(self) -> str:

@property
def uses_calculator(self) -> bool:
return self.arm != "web"
return self.arm in ("pe", "web+pe")

@property
def uses_web(self) -> bool:
return self.arm != "pe"
return self.arm in ("web", "web+pe")

@property
def provider(self) -> str:
Expand Down Expand Up @@ -1116,7 +1188,7 @@ def model_profile(
f"unknown prompt {prompt_version!r}; the prompts are "
f"{', '.join(PROMPT_VERSIONS)}"
)
if benchmark_search_rule and arm == "pe":
if benchmark_search_rule and arm not in ("web", "web+pe"):
raise SystemExit("the benchmark-search rule needs a web arm")
model_id = MODELS[name]
return ModelProfile(
Expand Down Expand Up @@ -1338,7 +1410,11 @@ def request_kwargs(self, base: dict, tools: list[dict], tool_choice) -> dict:
# then rebuilds the message in another block order.
if self.cache_breakpoint and messages[-1].get("role") != "assistant":
messages = [*messages[:-1], _with_cache_breakpoint(messages[-1])]
kwargs.update({"messages": messages, "tools": tools})
kwargs["messages"] = messages
# A request with no tools sends no tools key at all, as the board's
# JSON-contract request does.
if tools:
kwargs["tools"] = tools
if tool_choice is not None:
kwargs["tool_choice"] = tool_choice
if self.budget_key not in kwargs and "max_tokens" in kwargs:
Expand Down Expand Up @@ -1674,8 +1750,13 @@ class BudgetStop(RuntimeError):
pass


def is_timeout(error: Exception) -> bool:
"""A provider or wall-clock timeout (litellm.Timeout and its kin)."""
return any("Timeout" in cls.__name__ for cls in type(error).__mro__)


def _request_with_retries(
completion_fn, kwargs: dict
completion_fn, kwargs: dict, retry_timeouts: bool = True
) -> tuple[Any, float, float, float]:
from policybench.eval_no_tools import _run_request_with_wall_timeout, _should_retry

Expand All @@ -1686,7 +1767,11 @@ def _request_with_retries(
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):
if (
attempt == MAX_PROVIDER_ATTEMPTS - 1
or not _should_retry(error)
or (not retry_timeouts and is_timeout(error))
):
raise
time.sleep(2 * (2**attempt))
raise RuntimeError("unreachable")
Expand Down Expand Up @@ -1735,6 +1820,10 @@ def run_household(
offered = ([web_search_tool(profile)] if profile.uses_web else []) + (
[calculate_tool] if profile.uses_calculator else []
)
if profile.uses_web:
base["timeout"] = max(
float(base.get("timeout") or 0), WEB_REQUEST_TIMEOUT_SECONDS
)
if profile.uses_web and profile.provider == "gemini":
base["include_server_side_tool_invocations"] = True
if profile.uses_web and profile.provider == "openai":
Expand All @@ -1746,7 +1835,11 @@ def run_household(
ResponsesConversation(prompt)
if profile.api == "responses"
else ChatConversation(
prompt, cache_breakpoint=profile.model_id.startswith("claude-")
# The board's own request marks no prompt-cache breakpoint, so the
# no-tools arm, its stand-in, marks none either.
prompt,
cache_breakpoint=profile.model_id.startswith("claude-")
and profile.arm != "none",
)
)
json_contract = profile.contract == "json"
Expand Down Expand Up @@ -1779,7 +1872,9 @@ def run_household(

def tools_and_choice(kind: str):
if json_contract:
return list(offered), "auto"
# The board's JSON contract sends no tools, so the no-tools arm
# sends none and no tool_choice either.
return list(offered), ("auto" if offered else None)
if kind == "repair":
return (
[
Expand All @@ -1788,7 +1883,9 @@ def tools_and_choice(kind: str):
],
forced_answer_choice(profile),
)
if kind == "forced_answer":
if kind == "forced_answer" or not offered:
# With no tools offered, the board forces its answer function from
# the first turn; so does the no-tools arm.
return [*offered, answer_tool], forced_answer_choice(profile)
return [*offered, answer_tool], "required"

Expand Down Expand Up @@ -1913,7 +2010,7 @@ def run_calculations(turn_record: dict, calls: list[Call]) -> bool:
if repair_rounds:
kind = "repair"
elif not json_contract and (
forced or tool_turns + search_only_turns >= turn_cap
forced or tool_turns + search_only_turns >= turn_cap or not offered
):
kind, forced = "forced_answer", True
else:
Expand All @@ -1930,7 +2027,7 @@ def run_calculations(turn_record: dict, calls: list[Call]) -> bool:
escalated_from = None
try:
response, started_at, completed_at, latency = _request_with_retries(
completion_fn, kwargs
completion_fn, kwargs, retry_timeouts=not profile.uses_web
)
turn = conversation.read(response, kwargs.get(budget_key))
if (
Expand All @@ -1953,7 +2050,7 @@ def run_calculations(turn_record: dict, calls: list[Call]) -> bool:
escalated_from = kwargs.get(budget_key)
kwargs[budget_key] = ceiling
response, started_at, completed_at, latency = _request_with_retries(
completion_fn, kwargs
completion_fn, kwargs, retry_timeouts=not profile.uses_web
)
turn = conversation.read(response, kwargs.get(budget_key))
except Exception as error:
Expand Down Expand Up @@ -2540,6 +2637,13 @@ def command_run(args) -> None:
"tool_notice": profile.notice,
"calculate_tool": calculate_tool_schema() if profile.uses_calculator else None,
"web_search_tool": web_search_tool(profile) if profile.uses_web else None,
# The timeout actually sent: the board's for the model, or the web arms'.
"request_timeout_seconds": build_initial_request(
board["scenarios"][ids[0]], board["variables"][ids[0]], profile
)[1].get("timeout")
if ids and not profile.uses_web
else (WEB_REQUEST_TIMEOUT_SECONDS if profile.uses_web else None),
"web_timeouts_retried": not profile.uses_web,
"search_usd_per_1k": (
SEARCH_USD_PER_1K[profile.provider] if profile.uses_web else None
),
Expand All @@ -2557,13 +2661,18 @@ def command_run(args) -> None:
print(f"{len(todo)} households to run; ${spent_before:.4f} already spent")
if not todo:
return
pool = ToolPool(size=args.workers)
# Arms without the calculator start no engine workers.
pool = ToolPool(size=args.workers) if profile.uses_calculator else None

def no_calculator(arguments) -> tuple[dict, float]:
return {"error": "not_run", "message": "no calculator in this arm"}, 0.0

try:

def one(sid: str) -> dict:
if spend.exhausted(args.reserve_usd):
return {"scenario_id": sid, "status": "budget_stop"}
if pool.counters()["lost"] >= args.workers:
if pool is not None and pool.counters()["lost"] >= args.workers:
# Every worker is gone: start nothing more. The household keeps
# no record, so a resume runs it.
return {
Expand All @@ -2575,7 +2684,7 @@ def one(sid: str) -> dict:
record = run_household(
board["scenarios"][sid],
board["variables"][sid],
pool.call,
pool.call if pool is not None else no_calculator,
completion_fn,
spend,
profile=profile,
Expand All @@ -2584,6 +2693,12 @@ def one(sid: str) -> dict:
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],
Expand Down Expand Up @@ -2614,8 +2729,9 @@ def one(sid: str) -> dict:
) as executor:
results = list(executor.map(one, todo))
finally:
counters = pool.counters()
pool.close()
counters = pool.counters() if pool is not None else {}
if pool is not None:
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
Expand Down Expand Up @@ -4040,7 +4156,7 @@ def main(argv: list[str] | None = None) -> None:
"--arm",
choices=ARMS,
default="pe",
help="the tools: PolicyEngine (pe), web search (web), or both",
help="the tools: none, PolicyEngine (pe), web search (web), or both",
)
p.add_argument(
"--prompt",
Expand Down
Loading
Loading