Add minimal GAP reproduction package

1 files changed, 301 insertions, 0 deletions

diff --git a/src/gap_pipeline/e2e.py b/src/gap_pipeline/e2e.py
new file mode 100644
index 0000000..093068f
--- /dev/null
+++ b/src/gap_pipeline/e2e.py
@@ -0,0 +1,301 @@
+"""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()),
+    }