path: root/src/gap_pipeline/e2e.py

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