diff options
Diffstat (limited to 'tests/test_paper_pipeline.py')
| -rw-r--r-- | tests/test_paper_pipeline.py | 192 |
1 files changed, 192 insertions, 0 deletions
diff --git a/tests/test_paper_pipeline.py b/tests/test_paper_pipeline.py new file mode 100644 index 0000000..7f1de46 --- /dev/null +++ b/tests/test_paper_pipeline.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import asyncio +import json + +from gap_pipeline.clients import ScriptedClient +from gap_pipeline.paper_pipeline import PaperKernelPipeline, PaperPipelineConfig +from gap_pipeline.prompts import JUDGE_SYSTEM_PROMPT +from gap_pipeline.store import RunStore + + +def _dag() -> dict: + return { + "nodes": [ + { + "node_id": "n1", + "claim": "(a-1)^2 >= 0", + "dependencies": [], + }, + { + "node_id": "n2", + "claim": "a+1/a >= 2", + "dependencies": ["n1"], + }, + ], + "terminal_node_id": "n2", + } + + +def _methods() -> dict: + return { + "nodes": [ + { + "node_id": "n1", + "method_label": "use nonnegativity of a square", + }, + { + "node_id": "n2", + "method_label": "expand and divide by a positive quantity", + }, + ] + } + + +def _replacement() -> dict: + return { + "changes": [ + { + "slot_id": "slot1", + "source_node_id": "n1", + "description": "positive equality value", + "original_value": "1", + "replacement_value": "2", + "guard_condition": "replacement is positive", + "guard_justification": "2 is positive", + } + ], + "closure_statement": "No undeclared changes.", + } + + +def _diffused() -> dict: + return { + "nodes": [ + { + "node_id": "n1", + "dependencies": [], + "method_label": "use nonnegativity of a square", + "instantiated_claim": "(x-2)^2 >= 0", + "justification": "squares are nonnegative", + }, + { + "node_id": "n2", + "dependencies": ["n1"], + "method_label": "expand and divide by a positive quantity", + "instantiated_claim": "x+4/x >= 4", + "justification": "expand and divide by x>0", + }, + ], + "terminal_node_id": "n2", + "terminal_answer": "4", + } + + +def _variant(suffix: str = "") -> dict: + return { + "question": f"Let x>0. Prove that x+4/x >= 4.{suffix}", + "solution": ( + "[n1] Since (x-2)^2 >= 0. " + "[n2] Expand and divide by x>0 to obtain x+4/x >= 4." + ), + "node_order": ["n1", "n2"], + "terminal_answer": "4", + } + + +def _accept() -> dict: + return { + "verdict": "accept", + "step_by_step_check": "n1 is valid; n2 is valid", + "blocking_issues": "", + "patch_suggestion": "", + } + + +def _reject() -> dict: + return { + "verdict": "reject", + "step_by_step_check": "n1 is valid; n2 needs a wording correction", + "blocking_issues": "the terminal wording is ambiguous", + "patch_suggestion": "clarify the terminal wording", + } + + +def test_five_stage_repair_uses_prior_bundle_and_appendix_judge( + tmp_path, + item, +) -> None: + proposer = ScriptedClient( + { + f"{item.item_id}.stage1.dag": _dag(), + f"{item.item_id}.stage2.methods": _methods(), + f"{item.item_id}.stage3.replacement": [ + _replacement(), + _replacement(), + ], + f"{item.item_id}.stage4.diffusion": [ + _diffused(), + _diffused(), + ], + f"{item.item_id}.stage5.render": [ + _variant(), + _variant(" The requested bound is explicit."), + ], + } + ) + judges = [] + for judge_id in range(1, 6): + first = _reject() if judge_id == 1 else _accept() + judges.append( + ScriptedClient( + {f"{item.item_id}.verify": [first, _accept(), _accept()]} + ) + ) + + run_root = tmp_path / "run" + result = asyncio.run( + PaperKernelPipeline( + proposer=proposer, + judges=judges, + store=RunStore(run_root, item.item_id), + config=PaperPipelineConfig( + proposer_model="scripted", + judge_model="scripted", + ), + ).run(item) + ) + + assert result.status == "accepted" + assert [row.pass_streak_after for row in result.iterations] == [0, 1, 2] + assert result.iterations[0].bundle_sha256 != result.iterations[1].bundle_sha256 + assert result.iterations[1].bundle_sha256 == result.iterations[2].bundle_sha256 + + calls_dir = run_root / "items" / item.item_id / "calls" + repair_stage_calls = [ + json.loads( + ( + calls_dir + / f"{item.item_id}.stage{stage}.{name}.v02.json" + ).read_text() + ) + for stage, name in [ + (3, "replacement"), + (4, "diffusion"), + (5, "render"), + ] + ] + assert "PREVIOUS REPLACEMENT PLAN" in repair_stage_calls[0]["user_prompt"] + assert "PREVIOUS DIFFUSED PROOF" in repair_stage_calls[1]["user_prompt"] + assert "PREVIOUS RENDERED VARIANT" in repair_stage_calls[2]["user_prompt"] + assert all( + "the terminal wording is ambiguous" in call["user_prompt"] + for call in repair_stage_calls + ) + + judge_calls = sorted(calls_dir.glob(f"{item.item_id}.verify.*.json")) + assert len(judge_calls) == 15 + assert not any(".a1." in path.name for path in judge_calls) + first_judge_call = json.loads(judge_calls[0].read_text()) + assert first_judge_call["system_prompt"] == JUDGE_SYSTEM_PROMPT + assert "METHOD-LABEL SEQUENCE (abstract plan):" in first_judge_call["user_prompt"] + assert "SOURCE PROOF DAG:" not in first_judge_call["user_prompt"] |