path: root/src/gap_pipeline/kernel_prompts.py

"""Prompts for the paper's explicit five-stage kernel implementation.

The consolidated Prompt-A/Prompt-B interface remains byte-pinned in
``prompts.py``. The default pipeline exposes the manuscript's five operations
as separately validated artifacts and uses its Appendix F.3 judge prompt
verbatim.
"""

from __future__ import annotations

import json

from .kernel_models import (
    DiffusedProof,
    MethodPlan,
    ProofDAG,
    ReplacementPlan,
)
from .models import CanonicalItem
from .prompts import JUDGE_SYSTEM_PROMPT, JUDGE_USER_TEMPLATE


DAG_SYSTEM = "You are a rigorous competition-math proof analyst."
DAG_USER = """Parse the official solution into a concrete proof DAG.

Each node must be an intermediate mathematical claim, not a method name.
Each dependency must be a local entailment used to establish that node.
List nodes in topological order with IDs n1, n2, ... . Every node must
contribute to one terminal node.

Return JSON only:
{{"nodes":[{{"node_id":"n1","claim":"...","dependencies":[]}}],
 "terminal_node_id":"nN"}}

ORIGINAL PROBLEM:
<<<{question}>>>

OFFICIAL SOLUTION:
<<<{solution}>>>"""


METHOD_SYSTEM = "You abstract concrete proofs into content-free method plans."
METHOD_USER = """For every concrete proof-DAG node, provide one content-free
method label. Preserve node IDs and order exactly. A method label describes the
operation, not the source constants, variable names, or final numerical answer.

Return JSON only:
{{"nodes":[{{"node_id":"n1","method_label":"..."}}]}}

PROOF DAG:
{dag}"""


REPLACEMENT_SYSTEM = "You design guarded, proof-plan-preserving math replacements."
REPLACEMENT_USER = """Choose one or more substantive numerical or structural
replacements that create a genuinely new problem while preserving the supplied
method plan.

For every change:
- target one of the disclosed source leaf DAG nodes;
- state that exact source DAG node;
- record the original and replacement values explicitly;
- state a mathematical guard condition derived from the original problem;
- explain why the replacement satisfies the guard.

The list must be closed: stages 4 and 5 may introduce no mathematical change
that is not declared here. Variable renaming alone is not a kernel change.
{repair_context}

Return JSON only:
{{"changes":[
  {{"slot_id":"slot1","source_node_id":"n1","description":"...",
   "original_value":"...","replacement_value":"...",
   "guard_condition":"...","guard_justification":"..."}}
 ],
 "closure_statement":"All intended mathematical changes are listed above."}}

ORIGINAL PROBLEM:
<<<{question}>>>

OFFICIAL SOLUTION:
<<<{solution}>>>

CONCRETE PROOF DAG:
{dag}

SOURCE LEAF NODE IDS:
{leaf_node_ids}

CONTENT-FREE METHOD PLAN:
{methods}"""


DIFFUSION_SYSTEM = "You re-instantiate a proof DAG under declared guarded replacements."
DIFFUSION_USER = """Propagate only the declared replacements through the proof
DAG, node by node. Return exactly one row for every source node, preserving its
ID, dependencies, and method label. Do not introduce undeclared constants,
objects, assumptions, reductions, or proof methods.
{repair_context}

Return JSON only:
{{"nodes":[
  {{"node_id":"n1","dependencies":[],"method_label":"...",
   "instantiated_claim":"...","justification":"..."}}
 ],
 "terminal_node_id":"nN",
 "terminal_answer":"..."}}

ORIGINAL PROBLEM:
<<<{question}>>>

SOURCE PROOF DAG:
{dag}

METHOD PLAN:
{methods}

DECLARED REPLACEMENTS:
{replacements}"""


RENDER_SYSTEM = "You render a verified proof DAG into a self-contained Putnam problem."
RENDER_USER = """Render the diffused proof into one complete problem statement
and one complete solution.
{repair_context}

Requirements:
- the question must be fully determined by the diffused terminal claim;
- the solution must follow the diffused nodes in order and visibly mark each
  paragraph with [n1], [n2], ...;
- node_order must equal exactly {node_order};
- do not add any mathematical change beyond the declared replacements;
- copy the terminal answer exactly.

Return JSON only:
{{"question":"...","solution":"...","node_order":{node_order},
 "terminal_answer":"..."}}

DECLARED REPLACEMENTS:
{replacements}

DIFFUSED PROOF:
{diffused}"""


JUDGE_SYSTEM = JUDGE_SYSTEM_PROMPT


def _dump(value: object) -> str:
    if hasattr(value, "model_dump"):
        value = value.model_dump(mode="json")
    return json.dumps(value, ensure_ascii=False, indent=2)


def dag_user(item: CanonicalItem) -> str:
    return DAG_USER.format(question=item.problem, solution=item.solution)


def method_user(dag: ProofDAG) -> str:
    return METHOD_USER.format(dag=_dump(dag))


def replacement_user(
    item: CanonicalItem,
    dag: ProofDAG,
    methods: MethodPlan,
    *,
    previous_replacements: ReplacementPlan | None = None,
    feedback: str = "",
) -> str:
    repair_context = (
        "This is the initial replacement proposal."
        if previous_replacements is None
        else (
            "This is a repair pass. Preserve the previous slot IDs, source leaf "
            "nodes, and accepted changes unless the verification feedback "
            "identifies them as the blocking issue. Make only the smallest "
            "necessary correction.\n\nPREVIOUS REPLACEMENT PLAN:\n"
            f"{_dump(previous_replacements)}\n\nVERIFICATION FEEDBACK:\n"
            f"{feedback or 'No additional textual feedback was supplied.'}"
        )
    )
    return REPLACEMENT_USER.format(
        repair_context=repair_context,
        question=item.problem,
        solution=item.solution,
        dag=_dump(dag),
        leaf_node_ids=json.dumps(dag.leaf_node_ids(), ensure_ascii=False),
        methods=_dump(methods),
    )


def diffusion_user(
    item: CanonicalItem,
    dag: ProofDAG,
    methods: MethodPlan,
    replacements: ReplacementPlan,
    *,
    previous_diffused: DiffusedProof | None = None,
    feedback: str = "",
) -> str:
    repair_context = (
        "This is the initial DAG diffusion."
        if previous_diffused is None
        else (
            "This is a repair pass. Apply only corrections required by the "
            "verification feedback; preserve every unaffected node.\n\n"
            f"PREVIOUS DIFFUSED PROOF:\n{_dump(previous_diffused)}\n\n"
            f"VERIFICATION FEEDBACK:\n"
            f"{feedback or 'No additional textual feedback was supplied.'}"
        )
    )
    return DIFFUSION_USER.format(
        repair_context=repair_context,
        question=item.problem,
        dag=_dump(dag),
        methods=_dump(methods),
        replacements=_dump(replacements),
    )


def render_user(
    replacements: ReplacementPlan,
    diffused: DiffusedProof,
    *,
    previous_variant: object | None = None,
    feedback: str = "",
) -> str:
    repair_context = (
        "This is the initial rendering."
        if previous_variant is None
        else (
            "This is a repair pass. Preserve all unaffected wording and apply "
            "only corrections required by the verification feedback.\n\n"
            f"PREVIOUS RENDERED VARIANT:\n{_dump(previous_variant)}\n\n"
            f"VERIFICATION FEEDBACK:\n"
            f"{feedback or 'No additional textual feedback was supplied.'}"
        )
    )
    return RENDER_USER.format(
        repair_context=repair_context,
        node_order=json.dumps(
            [node.node_id for node in diffused.nodes],
            ensure_ascii=False,
        ),
        replacements=_dump(replacements),
        diffused=_dump(diffused),
    )


def judge_user(
    item: CanonicalItem,
    methods: MethodPlan,
    replacements: ReplacementPlan,
    variant: object,
) -> str:
    variant_payload = (
        variant.model_dump(mode="json")
        if hasattr(variant, "model_dump")
        else variant
    )
    return JUDGE_USER_TEMPLATE.format(
        original_problem=item.problem,
        original_solution=item.solution,
        method_labels=json.dumps(
            [
                f"{node.node_id}: {node.method_label}"
                for node in methods.nodes
            ],
            ensure_ascii=False,
        ),
        slot_replacement=_dump(replacements),
        candidate_problem=str(variant_payload["question"]),
        candidate_proof=str(variant_payload["solution"]),
    )