summaryrefslogtreecommitdiff
path: root/src/gap_pipeline/e2e.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/gap_pipeline/e2e.py')
-rw-r--r--src/gap_pipeline/e2e.py301
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()),
+ }