"""End-to-end helpers used by the reproducibility 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 .pipeline import KernelPipeline, PipelineConfig from .release import export_release from .store import RunStore from .surface import SURFACE_FAMILIES, 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 _accept_verdict() -> dict[str, Any]: return { "verdict": "accept", "step_by_step_check": "n1 and n2 instantiate the fixed method plan", "blocking_issues": "", "patch_suggestion": "", } def _offline_fixture() -> dict[str, Any]: plan = { "steps": [ { "node_id": "n1", "method_label": "use nonnegativity of a square", "input_roles": ["positive scalar"], "output_role": "nonnegative expression", "invariants": ["scalar is real"], }, { "node_id": "n2", "method_label": "expand and divide by a positive quantity", "input_roles": ["nonnegative expression", "positive scalar"], "output_role": "target inequality", "invariants": ["divisor is positive"], }, ], "terminal_node_id": "n2", } replacement = { "target_node_id": "n1", "original_object": "(a-1)^2", "replacement_object": "(x-2)^2", "guard_conditions": ["x>0"], "guard_evidence": ["positive real x satisfies the division guard"], "expected_downstream_changes": ["expand around 2", "divide by x"], "rationale": "the square-nonnegativity plan is unchanged", } diffused = { "node_instantiations": [ { "node_id": "n1", "method_label": plan["steps"][0]["method_label"], "instantiated_claim": "(x-2)^2 >= 0", "instantiated_derivation": "a square is nonnegative", "dependencies": [], }, { "node_id": "n2", "method_label": plan["steps"][1]["method_label"], "instantiated_claim": "x+4/x >= 4", "instantiated_derivation": "expand and divide by positive x", "dependencies": ["n1"], }, ], "regenerated_proof": ( "Since (x-2)^2 >= 0, expansion and division by x>0 give " "x+4/x >= 4." ), "terminal_answer": "x+4/x >= 4", } candidate = { "problem": "Let x>0. Prove that x+4/x >= 4.", "proof": diffused["regenerated_proof"], "terminal_answer": diffused["terminal_answer"], "node_instantiations": diffused["node_instantiations"], "replacement": replacement, } return { "record": { "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": {}, }, "dag": { "nodes": [ { "node_id": "n1", "claim": "(a-1)^2 >= 0", "derivation": "a square is nonnegative", "dependencies": [], "source_span": "Since (a-1)^2 >= 0", "mathematical_objects": ["a", "(a-1)^2"], }, { "node_id": "n2", "claim": "a+1/a >= 2", "derivation": "expand and divide by positive a", "dependencies": ["n1"], "source_span": "divide by a>0", "mathematical_objects": ["a"], }, ], "terminal_node_id": "n2", }, "plan": plan, "replacement": replacement, "diffused": diffused, "candidate": candidate, } 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 all five GAP variants for one Putnam item and export them.""" 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-v1", "proposer_model": model, } ) surface_pipeline = SurfacePipeline( OpenAIJsonClient(model, api_key=api_key), surface_store, ) surfaces = await surface_pipeline.run_all(item) config = PipelineConfig(proposer_model=model, judge_model=model) kernel_pipeline = KernelPipeline( 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, ) kernel = await kernel_pipeline.run(item) if kernel.status != "accepted": raise RuntimeError( f"kernel candidate was rejected after {len(kernel.iterations)} rounds; " f"inspect {kernel_root / 'items' / item.item_id}" ) 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 same end-to-end code path with deterministic 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" fixture = _offline_fixture() record = fixture["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) surface_responses = { f"{item_id}.surface.{family}.a": { "replacement": { "descriptive_long": "positiveQuantity", "descriptive_long_confusing": "walnutVioletTerrace", "descriptive_long_misleading": "primeFieldOrder", "garbled_string": "xcQ7h2ZfRw9v", }[family] } for family in SURFACE_FAMILIES } surface_store = RunStore(surface_root, item_id) surface_store.write_input(item) surface_store.write_config( {"protocol_name": "gap-surface-smoke", "proposer_model": "scripted"} ) surfaces = await SurfacePipeline( ScriptedClient(surface_responses), surface_store, ).run_all(item) proposer = ScriptedClient( { f"{item_id}.stage1": fixture["dag"], f"{item_id}.stage2": fixture["plan"], f"{item_id}.stage3": fixture["replacement"], f"{item_id}.stage4": fixture["diffused"], f"{item_id}.stage5": fixture["candidate"], } ) judges = [ ScriptedClient( { f"{item_id}.verify": [ _accept_verdict(), _accept_verdict(), ] } ) for _ in range(5) ] kernel = await KernelPipeline( proposer=proposer, judges=judges, store=RunStore(kernel_root, item_id), config=PipelineConfig( 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()), }