"""End-to-end helpers used by the reproduction notebook.""" from __future__ import annotations import json import os import time from pathlib import Path from typing import Any from .clients import OpenAIJsonClient, ScriptedClient from .models import CanonicalItem from .offline import load_dataset from .paper_pipeline import PaperKernelPipeline, PaperPipelineConfig from .release import export_release from .store import RunStore from .surface import SurfacePipeline def _ensure_fresh(path: Path) -> None: if path.exists() and any(path.iterdir()): raise FileExistsError(f"refusing to reuse non-empty output directory {path}") def _review_accept() -> dict[str, str]: return { "verdict": "accept", "step_by_step_check": "n1 passes; n2 passes; n3 passes", "blocking_issues": "", "patch_suggestion": "", } def _offline_record() -> dict[str, Any]: return { "index": "demo-A-1", "question": r"Let \(a>0\). Prove that \(a+1/a\ge 2\).", "solution": ( r"Since \((a-1)^2\ge0\), expand and divide by \(a>0\) " r"to obtain \(a+1/a\ge2\)." ), "vars": ["a"], "params": [], "sci_consts": [], "problem_type": "proof", "variants": {}, } async def run_live_item( *, dataset_dir: Path, item_id: str, work_root: Path, model: str = "o3", api_key: str | None = None, ) -> dict[str, Any]: """Generate four surface variants and one verified kernel variant.""" if not (api_key or os.getenv("OPENAI_API_KEY")): raise RuntimeError( "OPENAI_API_KEY is missing. Set it in the environment or enter it " "with getpass in the notebook." ) _ensure_fresh(work_root) work_root.mkdir(parents=True, exist_ok=True) surface_root = work_root / "surface-runs" kernel_root = work_root / "kernel-runs" release_root = work_root / "release" records = load_dataset(dataset_dir) if item_id not in records: raise KeyError(f"{item_id!r} not found in {dataset_dir}") item = CanonicalItem.from_public_record(records[item_id]) started = time.monotonic() surface_store = RunStore(surface_root, item.item_id) surface_store.write_input(item) surface_store.write_config( {"protocol_name": "gap-surface-original-prompts", "model": model} ) surfaces = await SurfacePipeline( OpenAIJsonClient(model, api_key=api_key), surface_store, ).run_all(item) config = PaperPipelineConfig(proposer_model=model, judge_model=model) kernel = await PaperKernelPipeline( proposer=OpenAIJsonClient(model, api_key=api_key), judges=[OpenAIJsonClient(model, api_key=api_key) for _ in range(5)], store=RunStore(kernel_root, item.item_id), config=config, ).run(item) if kernel.status != "accepted": raise RuntimeError( f"kernel candidate was rejected after {len(kernel.iterations)} rounds" ) manifest = export_release( source_dataset=dataset_dir, surface_run_root=surface_root, kernel_run_root=kernel_root, output_root=release_root, item_ids={item_id}, ) return { "mode": "live", "item_id": item_id, "model": model, "elapsed_seconds": round(time.monotonic() - started, 2), "surface_families": sorted(surfaces), "kernel_status": kernel.status, "verification_rounds": len(kernel.iterations), "export_status": manifest["status"], "exported_item_count": manifest["exported_item_count"], "work_root": str(work_root.resolve()), "release_record": manifest["exported"][0]["path"], "manifest": str((release_root / "manifest.json").resolve()), } async def run_offline_smoke(work_root: Path) -> dict[str, Any]: """Exercise the end-to-end wiring with deterministic model responses.""" _ensure_fresh(work_root) work_root.mkdir(parents=True, exist_ok=True) source_root = work_root / "source" source_root.mkdir() surface_root = work_root / "surface-runs" kernel_root = work_root / "kernel-runs" release_root = work_root / "release" record = _offline_record() item_id = str(record["index"]) (source_root / f"{item_id}.json").write_text( json.dumps(record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) item = CanonicalItem.from_public_record(record) names = { "descriptive_long": "positivequantity", "descriptive_long_confusing": "walnutvioletterrace", "descriptive_long_misleading": "primefieldorder", "garbled_string": "qzxwvtnphjgrksla", } surface_responses = { f"{item_id}.surface.{family}": { "map": {"a": name}, "question": record["question"].replace("a", name), "solution": record["solution"].replace("a", name), } for family, name in names.items() } surface_store = RunStore(surface_root, item_id) surface_store.write_input(item) surface_store.write_config( {"protocol_name": "gap-surface-smoke", "model": "scripted"} ) surfaces = await SurfacePipeline( ScriptedClient(surface_responses), surface_store, ).run_all(item) proposer = ScriptedClient( { f"{item_id}.stage1.dag": { "nodes": [ { "node_id": "n1", "claim": "(a-1)^2 >= 0", "dependencies": [], }, { "node_id": "n2", "claim": "a^2-2a+1 >= 0", "dependencies": ["n1"], }, { "node_id": "n3", "claim": "a+1/a >= 2", "dependencies": ["n2"], }, ], "terminal_node_id": "n3", }, f"{item_id}.stage2.methods": { "nodes": [ { "node_id": "n1", "method_label": "use nonnegativity of a square", }, { "node_id": "n2", "method_label": "expand the square", }, { "node_id": "n3", "method_label": "divide by a positive quantity", }, ] }, f"{item_id}.stage3.replacement": { "changes": [ { "slot_id": "slot1", "source_node_id": "n1", "description": "positive square root at equality", "original_value": "1", "replacement_value": "2", "guard_condition": "replacement is positive", "guard_justification": "2 is positive", } ], "closure_statement": "The equality value is the only change.", }, f"{item_id}.stage4.diffusion": { "nodes": [ { "node_id": "n1", "dependencies": [], "method_label": "use nonnegativity of a square", "instantiated_claim": "(x-2)^2 >= 0", "justification": "squares are nonnegative", }, { "node_id": "n2", "dependencies": ["n1"], "method_label": "expand the square", "instantiated_claim": "x^2-4x+4 >= 0", "justification": "expand n1", }, { "node_id": "n3", "dependencies": ["n2"], "method_label": "divide by a positive quantity", "instantiated_claim": "x+4/x >= 4", "justification": "divide n2 by x>0", }, ], "terminal_node_id": "n3", "terminal_answer": "4", }, f"{item_id}.stage5.render": { "question": "Let x>0. Prove that x+4/x >= 4.", "solution": ( "[n1] Since (x-2)^2 >= 0. [n2] Expanding gives " "x^2-4x+4 >= 0. [n3] Divide by x>0 to get x+4/x >= 4." ), "node_order": ["n1", "n2", "n3"], "terminal_answer": "4", }, } ) judges = [ ScriptedClient( {f"{item_id}.verify": [_review_accept(), _review_accept()]} ) for _ in range(5) ] kernel = await PaperKernelPipeline( proposer=proposer, judges=judges, store=RunStore(kernel_root, item_id), config=PaperPipelineConfig( proposer_model="scripted", judge_model="scripted", ), ).run(item) manifest = export_release( source_dataset=source_root, surface_run_root=surface_root, kernel_run_root=kernel_root, output_root=release_root, item_ids={item_id}, ) return { "mode": "offline-smoke", "item_id": item_id, "surface_families": sorted(surfaces), "kernel_status": kernel.status, "verification_rounds": len(kernel.iterations), "export_status": manifest["status"], "exported_item_count": manifest["exported_item_count"], "work_root": str(work_root.resolve()), "release_record": manifest["exported"][0]["path"], "manifest": str((release_root / "manifest.json").resolve()), }