path: root/src/gap_pipeline/e2e.py

"""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()),
    }