Align five-stage pipeline with manuscript

1 files changed, 36 insertions, 3 deletions

diff --git a/tests/test_prompts.py b/tests/test_prompts.py
index 4f6568b..9125d26 100644
--- a/tests/test_prompts.py
+++ b/tests/test_prompts.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import hashlib
 import inspect
 
-from gap_pipeline import prompts
+from gap_pipeline import kernel_prompts, prompts
 from gap_pipeline.clients import OpenAIJsonClient
 from gap_pipeline.kernel_models import (
     DiffusedProof,
@@ -12,6 +12,7 @@ from gap_pipeline.kernel_models import (
     ReplacementPlan,
 )
 from gap_pipeline.kernel_prompts import (
+    JUDGE_SYSTEM,
     dag_user,
     diffusion_user,
     judge_user,
@@ -40,6 +41,19 @@ EXPECTED = {
     "SURFACE_USER_TEMPLATE": "5d9043d7d1c1033db1aea0696812a2c4cbb6f3e28ca9d7b6436c036b40831fe7",
 }
 
+EXPECTED_FIVE_STAGE = {
+    "DAG_SYSTEM": "86e596a0c3aea07100ba642bb07d22e25385c2a51007b8f7578b662502ae4650",
+    "DAG_USER": "28757954499cee438cb3676fbb8aff39b2bbb4644a1f9dac2604b71d24b03547",
+    "METHOD_SYSTEM": "1091d4db002673abf67f4407418d0d7e34f648878f907d21e1ee910318d447ce",
+    "METHOD_USER": "79901153a7e2ef16ca4ade9763e4fd05a6d9d248a76085998723389f9934dca6",
+    "REPLACEMENT_SYSTEM": "2bf314a8cdcf986b344d45af081f7a1f675da6be5687ae4cd594db2ed7c0b25b",
+    "REPLACEMENT_USER": "6e88657245d951e47893eb9bcd6dbc5cd7a826fc67b8b2ce849da5b559f19fd9",
+    "DIFFUSION_SYSTEM": "f6ce12280b592830f5bacca3c8bc8131da26c9c0f94ffb9bfb7e9f53347f63ac",
+    "DIFFUSION_USER": "e92b39f15d1af1acc12daf34fabc0dfaf63bcebb4c947821fadb121aec44c9c1",
+    "RENDER_SYSTEM": "ead495a3f7df7d9655dc24ac2a437d4358ec080cd5e4477b3e97641828c3233e",
+    "RENDER_USER": "9e3eabd0fecf6946552db3217efa2eea7938099bcb093c9ca02fabf5b7d8f70c",
+}
+
 
 def test_prompt_values_are_byte_locked() -> None:
     actual = {
@@ -49,6 +63,16 @@ def test_prompt_values_are_byte_locked() -> None:
     assert actual == EXPECTED
 
 
+def test_five_stage_prompt_values_are_byte_locked() -> None:
+    actual = {
+        name: hashlib.sha256(
+            getattr(kernel_prompts, name).encode("utf-8")
+        ).hexdigest()
+        for name in EXPECTED_FIVE_STAGE
+    }
+    assert actual == EXPECTED_FIVE_STAGE
+
+
 def test_o3_adapter_does_not_send_temperature() -> None:
     source = inspect.getsource(OpenAIJsonClient.generate_json)
     assert '"temperature"' not in source
@@ -68,7 +92,7 @@ def test_literal_five_stage_prompts_render(item) -> None:
         {
             "changes": [
                 {
-                    "slot_id": "s1",
+                    "slot_id": "slot1",
                     "source_node_id": "n1",
                     "description": "constant",
                     "original_value": "1",
@@ -108,6 +132,15 @@ def test_literal_five_stage_prompts_render(item) -> None:
         replacement_user(item, dag, methods),
         diffusion_user(item, dag, methods, replacements),
         render_user(replacements, diffused),
-        judge_user(item, dag, methods, replacements, diffused, variant),
+        judge_user(item, methods, replacements, variant),
     ]
     assert all("{" in value and "}" in value for value in rendered)
+    assert JUDGE_SYSTEM == prompts.JUDGE_SYSTEM_PROMPT
+    assert rendered[-1] == prompts.JUDGE_USER_TEMPLATE.format(
+        original_problem=item.problem,
+        original_solution=item.solution,
+        method_labels='["n1: method"]',
+        slot_replacement=replacements.model_dump_json(indent=2),
+        candidate_problem="question",
+        candidate_proof="[n1] solution",
+    )