path: root/src/gap_pipeline/prompts.py

"""Prompts and strict JSON contracts for the GAP pipeline."""

from __future__ import annotations

import json

from .models import (
    CanonicalItem,
    KernelCandidate,
    MethodPlan,
    ProofDAG,
    ReplacementSpec,
)


DAG_SYSTEM = """You are a mathematical proof-structure parser.
Convert a supplied reference solution into a directed acyclic graph. Each node
must be one locally checkable intermediate claim. Edges point from prerequisites
to claims that use them. Preserve every essential proof step and do not invent
new mathematics. Return JSON only."""


def dag_user(item: CanonicalItem) -> str:
    return f"""ITEM ID: {item.item_id}

PROBLEM:
{item.problem}

REFERENCE SOLUTION:
{item.solution}

Return:
{{
  "nodes": [
    {{
      "node_id": "n1",
      "claim": "concrete intermediate mathematical claim",
      "derivation": "why this claim follows",
      "dependencies": [],
      "source_span": "corresponding reference-solution text",
      "mathematical_objects": ["objects used at this node"]
    }}
  ],
  "terminal_node_id": "node containing the requested conclusion"
}}

Requirements:
- every dependency must refer to an earlier logical prerequisite;
- the graph must be acyclic and connected to the terminal conclusion;
- leaf nodes must expose the initial mathematical objects that could be
  replaced to produce a genuinely new setting;
- retain all essential cases, guards, and endpoint checks."""


METHOD_SYSTEM = """You abstract concrete proof steps into a reusable proof plan.
For every DAG node, produce a content-light method label that says what operation
or lemma is used without copying the node's particular constants or object names.
Keep exactly the DAG's topological order and node IDs. Return JSON only."""


def method_user(dag: ProofDAG) -> str:
    return f"""PROOF DAG:
{dag.model_dump_json(indent=2)}

Return:
{{
  "steps": [
    {{
      "node_id": "n1",
      "method_label": "content-free operation or lemma",
      "input_roles": ["abstract role consumed by this step"],
      "output_role": "abstract role produced by this step",
      "invariants": ["conditions that must remain true"]
    }}
  ],
  "terminal_node_id": "{dag.terminal_node_id}"
}}

Do not merge, omit, reorder, or add nodes."""


REPLACEMENT_SYSTEM = """You design one guarded replacement at a leaf of a proof
DAG. The replacement must create a genuinely different mathematical setting
while allowing the exact method-label plan to remain executable. Prefer a
replacement that does not trivialize the problem. State every guard and cite
concrete evidence from the source inequalities or assumptions. Return JSON only."""


def replacement_user(
    item: CanonicalItem,
    dag: ProofDAG,
    plan: MethodPlan,
    *,
    difficulty_instruction: str,
) -> str:
    return f"""ORIGINAL PROBLEM:
{item.problem}

ORIGINAL SOLUTION:
{item.solution}

PROOF DAG:
{dag.model_dump_json(indent=2)}

METHOD PLAN:
{plan.model_dump_json(indent=2)}

DIFFICULTY INSTRUCTION:
{difficulty_instruction}

Return:
{{
  "target_node_id": "a leaf node ID",
  "original_object": "leaf object being replaced",
  "replacement_object": "new object or setting",
  "guard_conditions": ["condition required for every downstream step"],
  "guard_evidence": ["source-derived evidence that the guard is satisfiable"],
  "expected_downstream_changes": ["concrete changes expected during diffusion"],
  "rationale": "why the same plan remains valid and the task is not trivialized"
}}"""


DIFFUSION_SYSTEM = """You re-instantiate a fixed method plan after one guarded
leaf replacement. Propagate the replacement through every dependent DAG node.
You must keep exactly the original node order, dependency structure, and method
labels. All algebra, cases, domains, and endpoint checks must be valid in the
new setting. Return JSON only."""


def diffusion_user(
    item: CanonicalItem,
    dag: ProofDAG,
    plan: MethodPlan,
    replacement: ReplacementSpec,
) -> str:
    return f"""ORIGINAL PROBLEM:
{item.problem}

ORIGINAL REFERENCE SOLUTION:
{item.solution}

PROOF DAG:
{dag.model_dump_json(indent=2)}

FIXED METHOD PLAN:
{plan.model_dump_json(indent=2)}

GUARDED REPLACEMENT:
{replacement.model_dump_json(indent=2)}

Return:
{{
  "node_instantiations": [
    {{
      "node_id": "same node ID",
      "method_label": "exact corresponding method label",
      "instantiated_claim": "new concrete claim",
      "instantiated_derivation": "valid derivation in the new setting",
      "dependencies": ["same dependency IDs"]
    }}
  ],
  "regenerated_proof": "complete rigorous proof in the new setting",
  "terminal_answer": "terminal conclusion established by that proof"
}}

The sequence and labels must match exactly. Do not write a problem statement
yet; work forward from the replacement to a correct terminal answer."""


QUESTION_SYSTEM = """You reverse-engineer a self-contained competition
mathematics problem from a fully regenerated proof and terminal answer. State
exactly the assumptions used by the proof, ask exactly for its terminal
conclusion, and do not expose the solution strategy. Return JSON only."""


def question_user(
    item: CanonicalItem,
    dag: ProofDAG,
    plan: MethodPlan,
    replacement: ReplacementSpec,
    diffused_payload: dict,
) -> str:
    return f"""SOURCE PROBLEM (style reference only):
{item.problem}

FIXED METHOD PLAN:
{plan.model_dump_json(indent=2)}

REPLACEMENT:
{replacement.model_dump_json(indent=2)}

REGENERATED PROOF:
{json.dumps(diffused_payload, ensure_ascii=False, indent=2)}

Return:
{{
  "problem": "self-contained candidate problem",
  "proof": "the complete regenerated proof",
  "terminal_answer": "answer requested by the candidate problem",
  "node_instantiations": {json.dumps(diffused_payload.get("node_instantiations", []), ensure_ascii=False)},
  "replacement": {replacement.model_dump_json()}
}}

The problem must be genuinely different from the source and must neither omit
an assumption used by the proof nor add an assumption that trivializes it."""


# Four-field verifier JSON contract used by the GAP pipeline.
JUDGE_SYSTEM = """You are a verification judge for a kernel-variant generation
pipeline. You must decide whether a CANDIDATE variant problem and its CANDIDATE
proof are mathematically equivalent to the ORIGINAL problem under the given
METHOD-LABEL sequence (the abstract proof plan).

Your job is verification, not solving. You receive:
- the ORIGINAL problem statement and reference solution,
- the abstract METHOD-LABEL sequence,
- the SLOT replacement that was applied,
- the CANDIDATE variant statement and regenerated proof.

Check step by step that:
1. Every candidate step instantiates the corresponding method label.
2. The candidate proof is valid, with no unjustified leap.
3. The candidate problem is well posed and its requested terminal answer
   matches the regenerated proof.
4. The candidate is genuinely different but uses the same plan.

Return one JSON object only."""


def judge_user(
    item: CanonicalItem,
    plan: MethodPlan,
    candidate: KernelCandidate,
    *,
    judge_id: int,
    iteration: int,
) -> str:
    return f"""JUDGE REPLICATE: {judge_id}
ITERATION: {iteration}

ORIGINAL PROBLEM:
{item.problem}

ORIGINAL REFERENCE SOLUTION:
{item.solution}

METHOD-LABEL SEQUENCE:
{plan.model_dump_json(indent=2)}

SLOT REPLACEMENT:
{candidate.replacement.model_dump_json(indent=2)}

CANDIDATE VARIANT PROBLEM:
{candidate.problem}

CANDIDATE REGENERATED PROOF:
{candidate.proof}

CANDIDATE NODE INSTANTIATIONS:
{json.dumps([node.model_dump() for node in candidate.node_instantiations], ensure_ascii=False, indent=2)}

Return:
{{
  "verdict": "accept or reject",
  "step_by_step_check": "for every METHOD label, state whether the candidate step instantiates it correctly",
  "blocking_issues": "logical gap, computation error, ill-posed statement, or plan deviation; empty if accepted",
  "patch_suggestion": "minimal correction if rejected; empty if accepted"
}}

The step-by-step check must explicitly cover every method-plan node."""


REPAIR_SYSTEM = """You repair a rejected kernel candidate. Apply only the
minimal changes needed to address all blocking issues while preserving the
guarded replacement, DAG node order, dependencies, and exact method-label
sequence. Return the complete corrected candidate as JSON only."""


def repair_user(
    item: CanonicalItem,
    dag: ProofDAG,
    plan: MethodPlan,
    candidate: KernelCandidate,
    issues: list[dict],
) -> str:
    return f"""ORIGINAL PROBLEM:
{item.problem}

PROOF DAG:
{dag.model_dump_json(indent=2)}

FIXED METHOD PLAN:
{plan.model_dump_json(indent=2)}

CURRENT CANDIDATE:
{candidate.model_dump_json(indent=2)}

JUDGE FEEDBACK:
{json.dumps(issues, ensure_ascii=False, indent=2)}

Return the complete corrected object with fields:
problem, proof, terminal_answer, node_instantiations, replacement.
Do not change the replacement unless a judge proves it violates a guard; if it
must change, retain the same target leaf and explain the corrected guards in
the replacement rationale."""


SURFACE_NAME_SYSTEM = """You propose identifier replacements for a mathematical
problem. A replacement must be a single ASCII identifier beginning with a
letter, must not collide with supplied identifiers, and must fit the requested
family. Return JSON only."""


def surface_name_user(
    *,
    symbol: str,
    role: str,
    family: str,
    context: str,
    forbidden: list[str],
) -> str:
    descriptions = {
        "descriptive_long": "a semantically helpful descriptive identifier",
        "descriptive_long_confusing": "2-5 unrelated common words concatenated",
        "descriptive_long_misleading": (
            "mathematical jargon suggesting a different role"
        ),
        "garbled_string": "a 4-16 character semantically meaningless string",
    }
    return f"""SYMBOL: {symbol}
ROLE: {role}
FAMILY: {family} ({descriptions[family]})
CONTEXT:
{context}
FORBIDDEN IDENTIFIERS:
{json.dumps(forbidden)}

Return {{"replacement": "one valid identifier", "rationale": "brief reason"}}."""