Align five-stage pipeline with manuscript

1 files changed, 76 insertions, 77 deletions

diff --git a/src/gap_pipeline/kernel_prompts.py b/src/gap_pipeline/kernel_prompts.py
index c4a6128..f932d0a 100644
--- a/src/gap_pipeline/kernel_prompts.py
+++ b/src/gap_pipeline/kernel_prompts.py
@@ -1,8 +1,9 @@
-"""Prompts for an explicit implementation of the paper's five kernel stages.
+"""Prompts for the paper's explicit five-stage kernel implementation.
 
-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.
+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
@@ -16,6 +17,7 @@ from .kernel_models import (
     ReplacementPlan,
 )
 from .models import CanonicalItem
+from .prompts import JUDGE_SYSTEM_PROMPT, JUDGE_USER_TEMPLATE
 
 
 DAG_SYSTEM = "You are a rigorous competition-math proof analyst."
@@ -55,18 +57,19 @@ replacements that create a genuinely new problem while preserving the supplied
 method plan.
 
 For every change:
-- state the exact source DAG node;
+- 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.
-{feedback}
+{repair_context}
 
 Return JSON only:
 {{"changes":[
-  {{"slot_id":"s1","source_node_id":"n1","description":"...",
+  {{"slot_id":"slot1","source_node_id":"n1","description":"...",
    "original_value":"...","replacement_value":"...",
    "guard_condition":"...","guard_justification":"..."}}
  ],
@@ -81,6 +84,9 @@ OFFICIAL SOLUTION:
 CONCRETE PROOF DAG:
 {dag}
 
+SOURCE LEAF NODE IDS:
+{leaf_node_ids}
+
 CONTENT-FREE METHOD PLAN:
 {methods}"""
 
@@ -90,6 +96,7 @@ 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":[
@@ -115,6 +122,7 @@ DECLARED 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;
@@ -135,51 +143,7 @@ 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}"""
+JUDGE_SYSTEM = JUDGE_SYSTEM_PROMPT
 
 
 def _dump(value: object) -> str:
@@ -201,18 +165,27 @@ def replacement_user(
     dag: ProofDAG,
     methods: MethodPlan,
     *,
+    previous_replacements: ReplacementPlan | None = None,
     feedback: str = "",
 ) -> str:
-    feedback_block = (
-        f"Previous verification feedback to address:\n{feedback}"
-        if feedback
-        else "This is the initial replacement proposal."
+    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(
-        feedback=feedback_block,
+        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),
     )
 
@@ -222,8 +195,23 @@ def diffusion_user(
     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),
@@ -234,8 +222,23 @@ def diffusion_user(
 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,
@@ -247,30 +250,26 @@ def render_user(
 
 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],
+    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,
         ),
-        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),
+        slot_replacement=_dump(replacements),
+        candidate_problem=str(variant_payload["question"]),
+        candidate_proof=str(variant_payload["solution"]),
     )