From e6b5e5978281727f52af6488b9c0a7d19c929444 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Wed, 14 Oct 2026 12:45:49 -0700 Subject: [PATCH] Loop options for AIRO's instrument: floor, coverage, turn limit, debrief pi_tools.ts gains options, all off by default so the ForecastBench loop is unchanged: floor_counts ("returned" counts failed research), coverage (a recent-window search tagged with each group before submission), withhold_final (submit tools not offered until unlocked), warn_missing_keys, max_turns (then only submit_forecast, asked for in words, two more turns, then no tools; prose replies get a nudge), debrief (one more turn after the answer: merged deliveries, a Python validator via tool_cli.py validate, retries, debrief.json), and per-tool parameter schemas. Together they mirror AIRO's redlines.llm.call_tools. harness.run passes them through and returns `debrief` and `forced`. Tests drive each option through the real pi with the scripted model; each was checked by breaking it in pi_tools.ts. Co-Authored-By: Claude Opus 5.5 --- README.md | 13 +++ fri_forecaster/harness.py | 36 +++++-- fri_forecaster/pi_tools.ts | 199 +++++++++++++++++++++++++++++++++---- fri_forecaster/tool_cli.py | 17 +++- tests/test_pi_tools.py | 118 ++++++++++++++++++++-- 5 files changed, 347 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index ed08e5c..4b171e7 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,19 @@ with the tools it is given, and the answers land in the same folder, which is th ``` Every attempt gets a fresh `attempt-N` folder, so a retry never reads an older answer. + +The loop has options, all off by default (the ForecastBench driver's loop); AIRO's +runner turns them on to match its own instrument. `harness.run(..., **options)`: + +| option | what it does | +|---|---| +| `floor_counts="returned"` | failed research calls count toward the research floor too | +| `coverage=[keys]` | submission waits for a recent-window `web_search` tagged with each key | +| `withhold_final=True` | the submit tools are not offered at all until the floor and coverage are met | +| `warn_missing_keys=False` | the final answer is accepted as sent, with no warning about missing cells | +| `max_turns=N` | from turn N only `submit_forecast` is offered and asked for, for up to two more turns; a reply in prose gets a nudge | +| `debrief={...}` | after the answer, one more turn for a structured debrief, validated, with retries, saved to `debrief.json` | +| `tool_parameters={...}` | replace a tool's parameter schema, e.g. to require a `cause` on `web_search` | pi runs isolated: no built-in tools, no discovered extensions or skills, no AGENTS.md/CLAUDE.md, so the inbox is the whole prompt. diff --git a/fri_forecaster/harness.py b/fri_forecaster/harness.py index 46cda47..010550b 100644 --- a/fri_forecaster/harness.py +++ b/fri_forecaster/harness.py @@ -43,16 +43,28 @@ def key_of(cell, key_fields): return "|".join(str(cell.get(k)) for k in key_fields) +# The loop's optional settings (pi_tools.ts's header says what each does). +# Unset, a call runs the default loop: the ForecastBench driver's. +LOOP_OPTIONS = ("floor_counts", "coverage", "withhold_final", "warn_missing_keys", "max_turns", "debrief") + + def write_spec(workdir, *, cell_schema, key_fields, allowed_keys, tool_descriptions, - min_research=0, research_tools=RESEARCH_TOOLS, final_extra=None, cells_tool=True): + min_research=0, research_tools=RESEARCH_TOOLS, final_extra=None, cells_tool=True, + tool_parameters=None, **loop): """spec.json for pi_tools.ts. tool_descriptions: {tool name: description} for every tool offered (research tools, submit_forecast, and submit_cells - unless cells_tool is False).""" - tools = {name: {"description": tool_descriptions[name]} + unless cells_tool is False); tool_parameters: {tool name: JSON schema} to + replace a tool's built-in parameters; `loop`: LOOP_OPTIONS.""" + unknown = set(loop) - set(LOOP_OPTIONS) + if unknown: + raise TypeError(f"unknown loop options {sorted(unknown)}; have {LOOP_OPTIONS}") + tools = {name: {"description": tool_descriptions[name], + **({"parameters": tool_parameters[name]} if name in (tool_parameters or {}) else {})} for name in list(research_tools) + (["submit_cells"] if cells_tool else []) + ["submit_forecast"]} spec = {"tool_cli": TOOL_CLI, "research_tools": list(research_tools), "min_research": min_research, "tools": tools, "cell_schema": cell_schema, "key_fields": list(key_fields), - "allowed_keys": list(allowed_keys), "final_extra": final_extra or {}} + "allowed_keys": list(allowed_keys), "final_extra": final_extra or {}, + **{k: v for k, v in loop.items() if v is not None}} os.makedirs(workdir, exist_ok=True) with open(os.path.join(workdir, "spec.json"), "w") as f: json.dump(spec, f, indent=1) @@ -153,14 +165,15 @@ def collect(workdir, key_fields, allowed_keys): def run(*, model, system, prompt, workdir, cell_schema, key_fields, allowed_keys, tool_descriptions, min_research=0, research_tools=RESEARCH_TOOLS, final_extra=None, cells_tool=True, - thinking="high", pi=DEFAULT_PI, cwd=None, env=None, timeout=None): + thinking="high", pi=DEFAULT_PI, cwd=None, env=None, timeout=None, tool_parameters=None, **loop): """One pi call, in a fresh workdir/attempt-N (returned as "workdir"). Writes spec.json, prompt.txt, system.txt, events.jsonl, stderr.txt there - (and pi_tools.ts writes pieces.jsonl, final.json).""" + (and pi_tools.ts writes pieces.jsonl, final.json, debrief.json).""" workdir = new_attempt_dir(workdir) write_spec(workdir, cell_schema=cell_schema, key_fields=key_fields, allowed_keys=allowed_keys, tool_descriptions=tool_descriptions, min_research=min_research, - research_tools=research_tools, final_extra=final_extra, cells_tool=cells_tool) + research_tools=research_tools, final_extra=final_extra, cells_tool=cells_tool, + tool_parameters=tool_parameters, **loop) open(os.path.join(workdir, "prompt.txt"), "w").write(prompt) open(os.path.join(workdir, "system.txt"), "w").write(system) cmd = pi_command(model, system, prompt, thinking=thinking, pi=pi) @@ -177,5 +190,12 @@ def run(*, model, system, prompt, workdir, cell_schema, key_fields, allowed_keys usage["harness"] = "pi" usage["seconds"] = round(secs) cells, answer, delivery = collect(workdir, key_fields, allowed_keys) + final = read_json(os.path.join(workdir, "final.json")) or {} return {"rc": rc, "seconds": secs, "evidence": evidence, "usage": usage, "cells": cells, - "answer": answer, "delivery": delivery, "trace": events_path, "workdir": workdir} + "answer": answer, "delivery": delivery, "forced": bool(final.get("forced")), + "debrief": read_json(os.path.join(workdir, "debrief.json")), + "trace": events_path, "workdir": workdir} + + +def read_json(path): + return json.load(open(path)) if os.path.exists(path) else None diff --git a/fri_forecaster/pi_tools.ts b/fri_forecaster/pi_tools.ts index 12026ae..c077453 100644 --- a/fri_forecaster/pi_tools.ts +++ b/fri_forecaster/pi_tools.ts @@ -6,12 +6,34 @@ // tool_cli absolute path of tool_cli.py (the research tools) // research_tools which of web_search / read_page to offer ([] = no research) // min_research research calls that must return before submitting -// tools {name: {description}} for every tool offered +// tools {name: {description, parameters?}} for every tool offered; +// `parameters`, a JSON schema, replaces the built-in one // cell_schema JSON schema of ONE forecast cell // key_fields fields of a cell that identify it, e.g. ["question_id","horizon"] // allowed_keys every wanted key, fields joined by "|" // final_extra {name: JSON schema} extra properties of submit_forecast // +// Optional, each off by default (what the ForecastBench driver runs). AIRO's +// loop (redlines.llm.call_tools) sets them all: +// +// floor_counts "successful" (default): only research calls without an +// error count toward min_research; "returned": every call +// coverage question-group keys: submission waits for a web_search +// with `cause` = each key and `recent_days` set +// withhold_final the submit tools are not declared until the floor and +// coverage are met (default: declared, and refused) +// warn_missing_keys false: accept the final answer as sent (default true: +// one warning for unknown or missing keys) +// max_turns the turn limit: from the last turn, only +// submit_forecast is offered and the model is asked for it +// in words, for up to two more turns; then no tools. Also +// turns a reply in prose into a nudge to research or submit +// debrief {tool: {name, description, parameters}, prompt, retries, +// validator?: {file, function, kwargs}}: after the final +// answer, one more user turn offering only the debrief +// tool; deliveries merge, the validator's problems go back +// for up to `retries` more tries; written to debrief.json +// // Imports match the pi the harness is pinned to (harness.py // PINNED_PI_VERSION; package.json: @earendil-works/pi-coding-agent 0.99.2). // @@ -31,62 +53,107 @@ const spec = JSON.parse(readFileSync(`${OUT}/spec.json`, "utf8")); const keyFields: string[] = spec.key_fields; const allowed = new Set(spec.allowed_keys); const minResearch: number = spec.min_research; +const countsAll = spec.floor_counts === "returned"; +const coverage: string[] = spec.coverage || []; +const withholdFinal = !!spec.withhold_final; +const warnMissing = spec.warn_missing_keys !== false; +const maxTurns: number = spec.max_turns || 0; +const debrief = spec.debrief || null; let research = 0; // research calls that returned without an error -let researchFailed = 0; // ...and those that did: they do NOT count toward the floor +let researchFailed = 0; // ...and those that did: they count toward the floor only under floor_counts "returned" +const covered = new Set(); const received = new Map(); const cap = (s: string, n = 8000) => (s.length > n ? s.slice(0, n) : s); const keyOf = (f: any) => keyFields.map((k) => String(f?.[k])).join("|"); +const counted = () => (countsAll ? research + researchFailed : research); +const uncovered = () => coverage.filter((k) => !covered.has(k)); +let forcing = false; // past the turn limit: submit_forecast is offered regardless +const unlocked = () => forcing || (counted() >= minResearch && uncovered().length === 0); + +function runPy(args: string[], input?: string) { + return JSON.parse(execFileSync(PY, [spec.tool_cli, ...args], + { input, encoding: "utf8", maxBuffer: 64 * 1024 * 1024, env: process.env })); +} +const py = (tool: string, args: any) => runPy([tool, JSON.stringify(args)]); +// A debrief can outgrow one argv string (128 KB on Linux): it goes on stdin. +const validated = (value: any): string[] => + (runPy(["validate"], JSON.stringify({ ...debrief.validator, value })).problems || []).map(String); -function py(tool: string, args: any) { - const out = execFileSync(PY, [spec.tool_cli, tool, JSON.stringify(args)], - { encoding: "utf8", maxBuffer: 64 * 1024 * 1024, env: process.env }); - return JSON.parse(out); +// A refused call: the model reads `body` as the tool's result. +function refuse(body: any) { + return { content: [{ type: "text" as const, text: JSON.stringify(body) }], details: {}, isError: true }; } function notYet(name: string) { - return { content: [{ type: "text" as const, text: JSON.stringify({ error: `${name} is not offered yet: ${minResearch - research} more successful research call(s) must return first (${research} of ${minResearch} so far; ${researchFailed} failed)` }) }], details: {}, isError: true }; + const parts: string[] = []; + if (counted() < minResearch) { + parts.push(countsAll + ? `${minResearch - counted()} more research call(s) must return first (${counted()} of ${minResearch} so far)` + : `${minResearch - research} more successful research call(s) must return first (${research} of ${minResearch} so far; ${researchFailed} failed)`); + } + const missingCover = uncovered(); + if (missingCover.length) { + parts.push(`a web_search with recent_days set and cause= is still needed for: ${missingCover.join(", ")}`); + } + return refuse({ error: `${name} is not offered yet: ${parts.join("; ")}` }); } function missing() { return [...allowed].filter((k) => !received.has(k)); } +function say(text: string) { + return { type: "custom_message" as const, customType: "forecaster", content: text, display: true }; +} + let finalDone = false; // one final answer per call let finalWarned = false; // one chance to fix unknown/missing keys, then accepted as is const cells = Type.Array(Type.Unsafe(spec.cell_schema)); +const paramsOf = (name: string, builtin: any) => + spec.tools[name]?.parameters ? Type.Unsafe(spec.tools[name].parameters) : builtin; export default function (pi: ExtensionAPI) { const research_tools: string[] = spec.research_tools; + const submitTools = [...(spec.tools.submit_cells ? ["submit_cells"] : []), "submit_forecast"]; + const offer = () => pi.setActiveTools(unlocked() || !withholdFinal ? [...research_tools, ...submitTools] : research_tools); + + function researched(params: any, r: any, isSearch: boolean) { + const ok = r && !r.error; + if (ok) research += 1; else researchFailed += 1; + if (isSearch && params?.cause != null && params?.recent_days && (ok || countsAll)) covered.add(String(params.cause)); + if (withholdFinal) offer(); + } + if (research_tools.includes("web_search")) pi.registerTool({ name: "web_search", label: "web_search", description: spec.tools.web_search.description, - parameters: Type.Object({ + parameters: paramsOf("web_search", Type.Object({ query: Type.String(), max_results: Type.Optional(Type.Integer({ minimum: 1, maximum: 10 })), recent_days: Type.Optional(Type.Integer({ minimum: 1 })), - }), + })), async execute(_id, params) { const r = py("web_search", params); - if (r && !r.error) research += 1; else researchFailed += 1; + researched(params, r, true); return { content: [{ type: "text", text: cap(JSON.stringify(r)) }], details: r }; }, }); if (research_tools.includes("read_page")) pi.registerTool({ name: "read_page", label: "read_page", description: spec.tools.read_page.description, - parameters: Type.Object({ url: Type.String(), offset: Type.Optional(Type.Integer({ minimum: 0 })) }), + parameters: paramsOf("read_page", Type.Object({ url: Type.String(), offset: Type.Optional(Type.Integer({ minimum: 0 })) })), async execute(_id, params) { const r = py("read_page", params); - if (r && !r.error) research += 1; else researchFailed += 1; + researched(params, r, false); return { content: [{ type: "text", text: cap(JSON.stringify(r)) }], details: r }; }, }); if (spec.tools.submit_cells) pi.registerTool({ name: "submit_cells", label: "submit_cells", description: spec.tools.submit_cells.description, - parameters: Type.Object({ forecasts: cells }), + parameters: paramsOf("submit_cells", Type.Object({ forecasts: cells })), async execute(_id, params: any) { - if (research < minResearch) return notYet("submit_cells"); + if (!unlocked()) return notYet("submit_cells"); let n = 0; for (const f of params.forecasts || []) { if (f && typeof f === "object" && allowed.has(keyOf(f))) { @@ -109,21 +176,111 @@ export default function (pi: ExtensionAPI) { pi.registerTool({ name: "submit_forecast", label: "submit_forecast", description: spec.tools.submit_forecast.description, - parameters: Type.Object(finalProps), + parameters: paramsOf("submit_forecast", Type.Object(finalProps)), async execute(_id, params: any) { - if (research < minResearch) return notYet("submit_forecast"); - if (finalDone) return { content: [{ type: "text" as const, text: JSON.stringify({ error: "submit_forecast was already called; stop here." }) }], details: {}, isError: true }; + if (!unlocked()) return notYet("submit_forecast"); + if (finalDone) return refuse({ error: "submit_forecast was already called; stop here." }); const unknown = (params.forecasts || []).filter((f: any) => !allowed.has(keyOf(f))).map(keyOf); const have = new Set([...received.keys(), ...(params.forecasts || []).map(keyOf).filter((k: string) => allowed.has(k))]); const stillMissing = [...allowed].filter((k) => !have.has(k)); - if ((unknown.length || stillMissing.length) && !finalWarned) { + if (warnMissing && (unknown.length || stillMissing.length) && !finalWarned) { finalWarned = true; - return { content: [{ type: "text" as const, text: JSON.stringify({ error: "not submitted: fix and call submit_forecast again", unknown_keys: unknown, missing_keys: stillMissing }) }], details: {}, isError: true }; + return refuse({ error: "not submitted: fix and call submit_forecast again", unknown_keys: unknown, missing_keys: stillMissing }); } finalDone = true; - writeFileSync(`${OUT}/final.json`, JSON.stringify({ t: Date.now(), research, research_failed: researchFailed, answer: params }, null, 1)); + writeFileSync(`${OUT}/final.json`, JSON.stringify({ t: Date.now(), research, research_failed: researchFailed, forced: forcing, answer: params }, null, 1)); const r = { ok: true, received: received.size, of: allowed.size, in_final: (params.forecasts || []).length }; - return { content: [{ type: "text", text: JSON.stringify(r) + " Done: stop here." }], details: r }; + return { content: [{ type: "text", text: JSON.stringify(r) + (debrief ? "" : " Done: stop here.") }], details: r }; }, }); + + // The debrief: deliveries merge; the merged whole is validated. + const d = { result: null as any, complete: false, attempts: 0, problems: [] as string[], error: null as string | null }; + let phase: "forecast" | "debrief" | "done" = "forecast"; + let debriefCalled = false; + const saveDebrief = () => writeFileSync(`${OUT}/debrief.json`, JSON.stringify(d, null, 1)); + const parsed = (args: any) => { + if (!args || typeof args !== "object" || Array.isArray(args)) return args; + const out: Record = {}; + for (const [k, v] of Object.entries(args)) { + let x: any = v; + if (typeof v === "string" && v.trim().startsWith("{")) { try { x = JSON.parse(v); } catch { /* not JSON: kept */ } } + out[k] = x; + } + return out; + }; + if (debrief) pi.registerTool({ + name: debrief.tool.name, label: debrief.tool.name, + description: debrief.tool.description, + parameters: Type.Unsafe(debrief.tool.parameters), + async execute(_id, params: any) { + if (phase !== "debrief") return refuse({ error: `${debrief.tool.name} is offered after submit_forecast` }); + debriefCalled = true; + const args = parsed(params); + d.result = d.result && typeof d.result === "object" && args && typeof args === "object" ? { ...d.result, ...args } : args; + try { + d.problems = debrief.validator ? validated(d.result) : []; + } catch (e: any) { + d.problems = []; + d.error = `validator failed: ${String(e?.message || e).slice(0, 300)}`; + } + d.complete = d.problems.length === 0; + saveDebrief(); + const r = { accepted: d.complete, problems: d.problems }; + return { content: [{ type: "text", text: cap(JSON.stringify(r)) + (d.complete ? " Done: stop here." : "") }], details: r }; + }, + }); + + let turns = 0; + const limit = maxTurns ? maxTurns + 2 : 0; + const askFinal = "Now call submit_forecast with your final answer. Include any remaining forecast cells and required summary fields; previously accepted cells are retained."; + + pi.on("session_start", async () => { + if (debrief) pi.setActiveTools(pi.getActiveTools().filter((t) => t !== debrief.tool.name)); + if (withholdFinal) offer(); + }); + + pi.on("turn_end", async () => { + turns += 1; + if (phase === "forecast" && finalDone) { + if (!debrief) { phase = "done"; return; } + phase = "debrief"; + pi.setActiveTools([debrief.tool.name]); + saveDebrief(); + return { entries: [say(debrief.prompt)], continue: true }; + } + if (phase === "debrief") { + d.attempts += 1; + const called = debriefCalled; + debriefCalled = false; + if (!called) d.problems = [`${debrief.tool.name} was not called`]; + saveDebrief(); + if (d.complete || d.attempts >= 1 + (debrief.retries ?? 1)) { + phase = "done"; + pi.setActiveTools([]); + return; + } + return { entries: [say(called + ? `Call ${debrief.tool.name} again with what is still missing or wrong (entries you already sent are kept), fixing: ${d.problems.join("; ").slice(0, 4000)}` + : `Now call ${debrief.tool.name} with the debrief.`)], continue: true }; + } + if (phase === "forecast" && maxTurns) { + if (turns >= limit) { pi.setActiveTools([]); return; } + if (turns >= maxTurns - 1) { + forcing = true; + pi.setActiveTools(["submit_forecast"]); + return { entries: [say(askFinal)], continue: true }; + } + } + }); + + // A reply in prose (no tool call) ends pi's run; under max_turns, nudge instead. + pi.on("agent_before_settle", async () => { + if (phase !== "forecast" || finalDone || !maxTurns || turns >= limit || forcing) return; + const missingCover = uncovered(); + return { entries: [say(unlocked() + ? "Now call submit_forecast with your final answer." + : `Keep researching with the tools: submit_forecast is offered after ${minResearch} tool calls have returned (${counted()} so far)` + + (missingCover.length ? ` and after a web_search with recent_days on each of: ${missingCover.join(", ")}` : "") + ".")], continue: true }; + }); } diff --git a/fri_forecaster/tool_cli.py b/fri_forecaster/tool_cli.py index c8d0a21..f66e698 100644 --- a/fri_forecaster/tool_cli.py +++ b/fri_forecaster/tool_cli.py @@ -2,7 +2,12 @@ """The research tools as a one-shot CLI for the pi extension (pi_tools.ts): `tool_cli.py web_search '{"query": ...}'` prints the JSON fri_forecaster.tools returns. One implementation for every harness, so comparisons are of loops, -not search stacks. Keys come from the environment the harness passes down.""" +not search stacks. Keys come from the environment the harness passes down. + +`tool_cli.py validate` (arguments on stdin: {file, function, kwargs, value}) +runs a caller's debrief validator, `function(value, **kwargs)` from the +Python file `file`, and prints {"problems": [...]}.""" +import importlib.util import json import os import random @@ -35,7 +40,17 @@ def call(fn, args, sleep=time.sleep): return res +def validate(file, function, value, kwargs=None): + spec = importlib.util.spec_from_file_location("forecaster_validator", file) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return {"problems": [str(p) for p in (getattr(module, function)(value, **(kwargs or {})) or [])]} + + def main(): + if sys.argv[1] == "validate": + print(json.dumps(validate(**json.load(sys.stdin)))) + return load_keys() tool, args = sys.argv[1], json.loads(sys.argv[2] if len(sys.argv) > 2 else "{}") if tool not in FN: diff --git a/tests/test_pi_tools.py b/tests/test_pi_tools.py index 3c967b8..f76eeee 100644 --- a/tests/test_pi_tools.py +++ b/tests/test_pi_tools.py @@ -18,13 +18,20 @@ "properties": {"k": {"type": "string"}, "probability": {"type": "number"}}} DESC = {"web_search": "s", "read_page": "r", "submit_cells": "c", "submit_forecast": "f"} -# Research without the network: web_search answers, read_page fails. +# Research without the network: web_search answers, read_page fails. A +# debrief validator runs as the real tool_cli runs it. FAKE_TOOL_CLI = textwrap.dedent("""\ import json, sys + sys.path.insert(0, {repo!r}) tool = sys.argv[1] - print(json.dumps({"results": [{"url": "u", "title": "t"}]} if tool == "web_search" - else {"error": "page read failed: 404"})) - """) + if tool == "validate": + from fri_forecaster.tool_cli import validate + print(json.dumps(validate(**json.load(sys.stdin)))) + else: + print(json.dumps({{"results": [{{"url": "u", "title": "t"}}]}} if tool == "web_search" + else {{"error": "page read failed: 404"}})) + """).format(repo=os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +VALIDATOR = "def check(value, need):\n return [f'{k}: missing' for k in need if k not in (value or {})]\n" class ScriptedModel: @@ -69,6 +76,16 @@ def chunks(turn, n): yield dict(base, choices=[{"index": 0, "delta": {}, "finish_reason": finish}]) yield dict(base, choices=[], usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}) + def tools(self, n): + """The tools declared in request n (1-based).""" + return {t["function"]["name"] for t in self.requests[n - 1].get("tools") or []} + + def last_user(self, n): + """The text of the last user message in request n (1-based).""" + msg = [m for m in self.requests[n - 1]["messages"] if m.get("role") == "user"][-1] + content = msg["content"] + return content if isinstance(content, str) else "".join(p.get("text", "") for p in content) + def tool_result(self, n): """The tool result the model saw at the start of request n (1-based): the answer to the call it made in turn n-1.""" @@ -81,8 +98,14 @@ def close(self): # Skipped without node and pi, except in CI, where a skip would hide a broken install. -@unittest.skipUnless((harness.DEFAULT_PI and shutil.which("node")) or os.environ.get("CI"), "needs node and pi (npm ci)") -class PiToolsTest(unittest.TestCase): +NEEDS_PI = unittest.skipUnless((harness.DEFAULT_PI and shutil.which("node")) or os.environ.get("CI"), + "needs node and pi (npm ci)") + + +class PiCase(unittest.TestCase): + """A temp dir, the fake research tools, and run_pi: one harness.run + through the real pi against a ScriptedModel.""" + def setUp(self): self.tmp = tempfile.TemporaryDirectory() d = self.tmp.name @@ -109,6 +132,11 @@ def run_pi(self, script, **kw): args.update(kw) return harness.run(**args), model + +@NEEDS_PI +class PiToolsTest(PiCase): + """The default loop: the ForecastBench driver's.""" + def test_research_floor_pieces_and_final(self): res, model = self.run_pi([ ("submit_forecast", {"forecasts": [{"k": "a", "probability": 0.1}], "rationale": "early"}), # 1 @@ -158,5 +186,83 @@ def test_no_research_offers_only_submit_forecast(self): self.assertEqual({k: c["probability"] for k, c in res["cells"].items()}, {"a": 0.4, "b": 0.6}) +SEARCH_WITH_CAUSE = {"type": "object", "required": ["query", "cause"], + "properties": {"query": {"type": "string"}, "recent_days": {"type": "integer"}, + "cause": {"type": "string", "enum": ["ai", "bio"]}}} + + +@NEEDS_PI +class LoopOptionsTest(PiCase): + """The loop options AIRO's runner sets (pi_tools.ts header).""" + + def test_withheld_until_every_call_and_group_counts(self): + res, model = self.run_pi([ + ("read_page", {"url": "u"}), # fails, still counts + ("web_search", {"query": "q", "cause": "ai", "recent_days": 7}), + ("web_search", {"query": "q", "cause": "bio"}), # no recent_days: not coverage + ("web_search", {"query": "q", "cause": "bio", "recent_days": 30}), + ("submit_forecast", {"forecasts": [{"k": "a", "probability": 0.1}], "rationale": "r"}), + "done", + ], min_research=4, floor_counts="returned", coverage=["ai", "bio"], withhold_final=True, + warn_missing_keys=False, tool_parameters={"web_search": SEARCH_WITH_CAUSE}) + self.assertEqual(res["rc"], 0, open(os.path.join(res["workdir"], "stderr.txt")).read()) + for n in (1, 2, 3, 4): # 4 returned (one failed) and both groups covered after call 4 + self.assertNotIn("submit_forecast", model.tools(n)) + self.assertNotIn("submit_cells", model.tools(n)) + self.assertTrue({"submit_forecast", "submit_cells"} <= model.tools(5)) + sent = next(t for t in model.requests[0]["tools"] if t["function"]["name"] == "web_search") + self.assertEqual(sent["function"]["parameters"]["properties"]["cause"]["enum"], ["ai", "bio"]) + final = json.load(open(os.path.join(res["workdir"], "final.json"))) + self.assertEqual((final["research"], final["research_failed"], final["forced"]), (3, 1, False)) + self.assertEqual({k: c["probability"] for k, c in res["cells"].items()}, {"a": 0.1}) # no warning for b + + def test_turn_limit_offers_only_submission(self): + res, model = self.run_pi([ + ("web_search", {"query": "q1"}), + ("web_search", {"query": "q2"}), + ("submit_forecast", {"forecasts": [{"k": "a", "probability": 0.2}], "rationale": "forced"}), + "done", + ], min_research=5, withhold_final=True, warn_missing_keys=False, max_turns=3) + self.assertEqual(res["rc"], 0, open(os.path.join(res["workdir"], "stderr.txt")).read()) + self.assertEqual(model.tools(3), {"submit_forecast"}) + self.assertTrue(model.last_user(3).startswith("Now call submit_forecast with your final answer.")) + self.assertTrue(res["forced"]) + + def test_turn_limit_ends_a_call_that_never_submits(self): + res, model = self.run_pi(["thinking", "still thinking", "no", "no", "no", "no"], + min_research=1, withhold_final=True, max_turns=3) + self.assertEqual(res["rc"], 0, open(os.path.join(res["workdir"], "stderr.txt")).read()) + self.assertEqual(len(model.requests), 5) # the limit is max_turns + 2 + self.assertIn("Keep researching with the tools", model.last_user(2)) + self.assertEqual(model.tools(5), {"submit_forecast"}) + self.assertFalse(os.path.exists(os.path.join(res["workdir"], "final.json"))) + + def test_debrief_merges_and_retries(self): + validator = os.path.join(self.tmp.name, "validator.py") + open(validator, "w").write(VALIDATOR) + debrief = {"tool": {"name": "submit_debrief", "description": "d", + "parameters": {"type": "object", "properties": {"a": {"type": "string"}, + "b": {"type": "string"}}}}, + "prompt": "Debrief now.", "retries": 1, + "validator": {"file": validator, "function": "check", "kwargs": {"need": ["a", "b"]}}} + res, model = self.run_pi([ + ("web_search", {"query": "q1"}), + ("web_search", {"query": "q2"}), + ("submit_forecast", {"forecasts": [{"k": "a", "probability": 0.1}, {"k": "b", "probability": 0.2}], + "rationale": "r"}), + ("submit_debrief", {"a": "why a"}), + ("submit_debrief", {"b": "why b"}), + "done", + ], debrief=debrief) + self.assertEqual(res["rc"], 0, open(os.path.join(res["workdir"], "stderr.txt")).read()) + self.assertNotIn("submit_debrief", model.tools(1)) + self.assertEqual(model.tools(4), {"submit_debrief"}) + self.assertEqual(model.last_user(4), "Debrief now.") + self.assertIn("b: missing", model.tool_result(5)) + self.assertIn("fixing: b: missing", model.last_user(5)) + self.assertEqual(res["debrief"], {"result": {"a": "why a", "b": "why b"}, "complete": True, + "attempts": 2, "problems": [], "error": None}) + + if __name__ == "__main__": unittest.main()