path: root/src/gap_pipeline/kernel_prompts.py

"""Prompts for an explicit implementation of the paper's five kernel stages.

The byte-pinned historical Prompt-A/Prompt-B remain in ``prompts.py``.  These
prompts make the richer five-stage manuscript description executable and keep
every transformation decision in a typed artifact.
"""

from __future__ import annotations

import json

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


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:
- state the 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.
{feedback}

Return JSON only:
{{"changes":[
  {{"slot_id":"s1","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}

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.

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.

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 = """You are a verification judge for a literal five-stage GAP
kernel transformation. Verification, not de novo solving, is your task."""
JUDGE_USER = """Check the candidate against every disclosed artifact.

You must verify:
1. every replacement satisfies its guard and all mathematical changes are
   declared in the replacement plan;
2. each diffused node instantiates the matching source node and method label;
3. dependencies and proof order are preserved;
4. the rendered problem is well-posed and the rendered solution proves it;
5. the terminal answer agrees across diffusion, rendering, and solution.

In step_by_step_check, mention each of these node IDs explicitly:
{required_node_ids}
In replacement_check, mention each of these slot IDs explicitly:
{required_slot_ids}
{format_feedback}

Return JSON only:
{{"verdict":"accept" or "reject",
 "step_by_step_check":"...",
 "replacement_check":"...",
 "blocking_issues":"...",
 "patch_suggestion":"..."}}

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

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

SOURCE PROOF DAG:
{dag}

METHOD PLAN:
{methods}

REPLACEMENT PLAN:
{replacements}

DIFFUSED PROOF:
{diffused}

RENDERED VARIANT:
{variant}"""


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,
    *,
    feedback: str = "",
) -> str:
    feedback_block = (
        f"Previous verification feedback to address:\n{feedback}"
        if feedback
        else "This is the initial replacement proposal."
    )
    return REPLACEMENT_USER.format(
        feedback=feedback_block,
        question=item.problem,
        solution=item.solution,
        dag=_dump(dag),
        methods=_dump(methods),
    )


def diffusion_user(
    item: CanonicalItem,
    dag: ProofDAG,
    methods: MethodPlan,
    replacements: ReplacementPlan,
) -> str:
    return DIFFUSION_USER.format(
        question=item.problem,
        dag=_dump(dag),
        methods=_dump(methods),
        replacements=_dump(replacements),
    )


def render_user(
    replacements: ReplacementPlan,
    diffused: DiffusedProof,
) -> str:
    return RENDER_USER.format(
        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,
    dag: ProofDAG,
    methods: MethodPlan,
    replacements: ReplacementPlan,
    diffused: DiffusedProof,
    variant: object,
    *,
    format_feedback: str = "",
) -> str:
    return JUDGE_USER.format(
        required_node_ids=json.dumps(dag.node_ids(), ensure_ascii=False),
        required_slot_ids=json.dumps(
            [change.slot_id for change in replacements.changes],
            ensure_ascii=False,
        ),
        format_feedback=(
            f"Previous report-format error: {format_feedback}"
            if format_feedback
            else ""
        ),
        question=item.problem,
        solution=item.solution,
        dag=_dump(dag),
        methods=_dump(methods),
        replacements=_dump(replacements),
        diffused=_dump(diffused),
        variant=_dump(variant),
    )