diff options
Diffstat (limited to 'src/gap_pipeline/pipeline.py')
| -rw-r--r-- | src/gap_pipeline/pipeline.py | 385 |
1 files changed, 116 insertions, 269 deletions
diff --git a/src/gap_pipeline/pipeline.py b/src/gap_pipeline/pipeline.py index 2c0deb4..307ba81 100644 --- a/src/gap_pipeline/pipeline.py +++ b/src/gap_pipeline/pipeline.py @@ -1,42 +1,30 @@ -"""Five-stage kernel generation and the J=5, K=2, T=15 verifier.""" +"""Prompt-faithful kernel generation and the J=5, K=2, T=15 loop.""" from __future__ import annotations import asyncio -import hashlib -from pathlib import Path -from typing import Any from pydantic import BaseModel, ConfigDict, model_validator from .clients import JsonLLM from .models import ( CanonicalItem, - DiffusedProof, IterationRecord, JudgeVerdict, KernelCandidate, + KernelPlan, KernelRunResult, - MethodPlan, ModelCallRecord, - ProofDAG, - ReplacementSpec, ) from .prompts import ( - DAG_SYSTEM, - DIFFUSION_SYSTEM, - JUDGE_SYSTEM, - METHOD_SYSTEM, - QUESTION_SYSTEM, - REPAIR_SYSTEM, - REPLACEMENT_SYSTEM, - dag_user, - diffusion_user, + FIX_SYSTEM_PROMPT, + JUDGE_SYSTEM_PROMPT, + KERNEL_GENERATE_SYSTEM, + KERNEL_PLAN_SYSTEM, + fix_user, judge_user, - method_user, - question_user, - repair_user, - replacement_user, + kernel_generate_user, + kernel_plan_user, ) from .store import RunStore, sha256_payload @@ -50,22 +38,11 @@ class PipelineConfig(BaseModel): judge_count: int = 5 streak_length: int = 2 max_iterations: int = 15 - human_audit_fraction: float = 0.10 - audit_seed: int = 0 - difficulty_instruction: str = ( - "Preserve the proof plan and avoid trivializing the source problem. " - "Routine algebraic adaptation is allowed, but do not introduce a new " - "pivotal lemma or an independent proof strategy." - ) @model_validator(mode="after") - def enforce_submitted_protocol(self) -> "PipelineConfig": + def enforce_protocol(self) -> "PipelineConfig": if (self.judge_count, self.streak_length, self.max_iterations) != (5, 2, 15): - raise ValueError( - "the GAP protocol is fixed at J=5, K=2, T=15" - ) - if self.human_audit_fraction != 0.10: - raise ValueError("the GAP post-hoc audit fraction is 10%") + raise ValueError("the GAP protocol is fixed at J=5, K=2, T=15") return self @@ -80,7 +57,7 @@ class KernelPipeline: ) -> None: if len(judges) != config.judge_count: raise ValueError( - f"expected {config.judge_count} judge clients, received {len(judges)}" + f"expected {config.judge_count} judges, received {len(judges)}" ) self.proposer = proposer self.judges = judges @@ -94,8 +71,7 @@ class KernelPipeline: request_id: str, system_prompt: str, user_prompt: str, - response_type: type[BaseModel], - ) -> BaseModel: + ) -> dict: response = await client.generate_json( system_prompt=system_prompt, user_prompt=user_prompt, @@ -113,193 +89,125 @@ class KernelPipeline: usage=response.usage, ) ) - return response_type.model_validate(response.data) + return response.data - async def construct_dag(self, item: CanonicalItem) -> ProofDAG: - request_id = f"{item.item_id}.stage1.dag" - dag = await self._call( + async def extract_plan(self, item: CanonicalItem) -> KernelPlan: + request_id = f"{item.item_id}.plan" + payload = await self._call( self.proposer, request_id=request_id, - system_prompt=DAG_SYSTEM, - user_prompt=dag_user(item), - response_type=ProofDAG, + system_prompt=KERNEL_PLAN_SYSTEM, + user_prompt=kernel_plan_user(item), ) - assert isinstance(dag, ProofDAG) - self.store.write_stage("01_proof_dag", dag, request_id=request_id) - return dag - - async def summarize_methods(self, item: CanonicalItem, dag: ProofDAG) -> MethodPlan: - request_id = f"{item.item_id}.stage2.methods" - plan = await self._call( - self.proposer, + plan = KernelPlan.model_validate(payload) + nodes = [ + { + "node_id": f"n{index}", + "method_label": step, + "dependencies": [] if index == 1 else [f"n{index - 1}"], + } + for index, step in enumerate(plan.core_steps, start=1) + ] + self.store.write_stage( + "01_proof_dag", + { + "nodes": nodes, + "terminal_node_id": nodes[-1]["node_id"], + "construction": "ordered core_steps returned verbatim by Prompt-A", + }, request_id=request_id, - system_prompt=METHOD_SYSTEM, - user_prompt=method_user(dag), - response_type=MethodPlan, ) - assert isinstance(plan, MethodPlan) - plan.validate_against(dag) - self.store.write_stage("02_method_plan", plan, request_id=request_id) - return plan - - async def generate_replacement( - self, - item: CanonicalItem, - dag: ProofDAG, - plan: MethodPlan, - ) -> ReplacementSpec: - request_id = f"{item.item_id}.stage3.replacement" - replacement = await self._call( - self.proposer, + self.store.write_stage( + "02_method_plan", + {"method_labels": plan.core_steps}, request_id=request_id, - system_prompt=REPLACEMENT_SYSTEM, - user_prompt=replacement_user( - item, - dag, - plan, - difficulty_instruction=self.config.difficulty_instruction, - ), - response_type=ReplacementSpec, ) - assert isinstance(replacement, ReplacementSpec) - replacement.validate_against(dag) self.store.write_stage( - "03_replacement", - replacement, + "03_mutable_slots", + { + key: value.model_dump(mode="json") + for key, value in plan.mutable_slots.items() + }, request_id=request_id, ) - return replacement + return plan - async def diffuse_dag( + async def generate_candidate( self, item: CanonicalItem, - dag: ProofDAG, - plan: MethodPlan, - replacement: ReplacementSpec, - ) -> DiffusedProof: - request_id = f"{item.item_id}.stage4.diffusion" - diffused = await self._call( + plan: KernelPlan, + ) -> KernelCandidate: + request_id = f"{item.item_id}.candidate" + payload = await self._call( self.proposer, request_id=request_id, - system_prompt=DIFFUSION_SYSTEM, - user_prompt=diffusion_user(item, dag, plan, replacement), - response_type=DiffusedProof, + system_prompt=KERNEL_GENERATE_SYSTEM, + user_prompt=kernel_generate_user(item, plan), ) - assert isinstance(diffused, DiffusedProof) - diffused.validate_against(dag, plan) + candidate = KernelCandidate.model_validate(payload) self.store.write_stage( - "04_diffused_proof", - diffused, + "04_regenerated_proof", + {"solution": candidate.solution}, request_id=request_id, ) - return diffused - - async def render_question( - self, - item: CanonicalItem, - dag: ProofDAG, - plan: MethodPlan, - replacement: ReplacementSpec, - diffused: DiffusedProof, - ) -> KernelCandidate: - request_id = f"{item.item_id}.stage5.question" - candidate = await self._call( - self.proposer, - request_id=request_id, - system_prompt=QUESTION_SYSTEM, - user_prompt=question_user( - item, - dag, - plan, - replacement, - diffused.model_dump(mode="json"), - ), - response_type=KernelCandidate, - ) - assert isinstance(candidate, KernelCandidate) - candidate.validate_against(dag, plan) self.store.write_stage( - "05_draft_candidate", - candidate, + "05_variant_question", + {"question": candidate.question}, request_id=request_id, ) return candidate - async def build_candidate( - self, - item: CanonicalItem, - ) -> tuple[ProofDAG, MethodPlan, KernelCandidate]: - dag = await self.construct_dag(item) - plan = await self.summarize_methods(item, dag) - replacement = await self.generate_replacement(item, dag, plan) - diffused = await self.diffuse_dag(item, dag, plan, replacement) - candidate = await self.render_question( - item, - dag, - plan, - replacement, - diffused, - ) - return dag, plan, candidate - async def _judge_once( self, judge: JsonLLM, *, item: CanonicalItem, - plan: MethodPlan, + plan: KernelPlan, candidate: KernelCandidate, iteration: int, judge_id: int, ) -> JudgeVerdict: request_id = f"{item.item_id}.verify.t{iteration:02d}.j{judge_id}" - verdict = await self._call( + payload = await self._call( judge, request_id=request_id, - system_prompt=JUDGE_SYSTEM, - user_prompt=judge_user( - item, - plan, - candidate, - judge_id=judge_id, - iteration=iteration, - ), - response_type=JudgeVerdict, + system_prompt=JUDGE_SYSTEM_PROMPT, + user_prompt=judge_user(item, plan, candidate), ) - assert isinstance(verdict, JudgeVerdict) - return verdict + return JudgeVerdict.model_validate(payload) async def _repair( self, *, item: CanonicalItem, - dag: ProofDAG, - plan: MethodPlan, candidate: KernelCandidate, verdicts: list[JudgeVerdict], iteration: int, ) -> KernelCandidate: + problem_issues: list[str] = [] + solution_issues: list[str] = [] + for verdict in verdicts: + if verdict.verdict == "reject": + problem_issues.append(verdict.blocking_issues) + solution_issues.append( + verdict.patch_suggestion or verdict.blocking_issues + ) request_id = f"{item.item_id}.repair.after_t{iteration:02d}" - issues = [ - { - "judge_id": judge_id, - "blocking_issues": verdict.blocking_issues, - "patch_suggestion": verdict.patch_suggestion, - } - for judge_id, verdict in enumerate(verdicts, start=1) - if verdict.verdict == "reject" - ] - repaired = await self._call( + payload = await self._call( self.proposer, request_id=request_id, - system_prompt=REPAIR_SYSTEM, - user_prompt=repair_user(item, dag, plan, candidate, issues), - response_type=KernelCandidate, + system_prompt=FIX_SYSTEM_PROMPT, + user_prompt=fix_user( + item, + candidate, + problem_issues="; ".join(problem_issues), + solution_issues="; ".join(solution_issues), + ), + ) + repaired = KernelCandidate( + question=str(payload["corrected_question"]), + solution=str(payload["corrected_solution"]), ) - assert isinstance(repaired, KernelCandidate) - repaired.validate_against(dag, plan) - if sha256_payload(repaired) == sha256_payload(candidate): - raise ValueError("repair call returned a byte-equivalent candidate") self.store.write_stage( f"repair_after_{iteration:02d}", repaired, @@ -307,25 +215,17 @@ class KernelPipeline: ) return repaired - def _audit_selected(self, item_id: str) -> bool: - digest = hashlib.sha256( - f"{self.config.audit_seed}:{item_id}".encode("utf-8") - ).digest() - draw = int.from_bytes(digest[:8], "big") / float(2**64) - return draw < self.config.human_audit_fraction - async def verify_candidate( self, *, item: CanonicalItem, - dag: ProofDAG, - plan: MethodPlan, + plan: KernelPlan, candidate: KernelCandidate, ) -> KernelRunResult: + current = candidate iterations: list[IterationRecord] = [] pass_streak = 0 - current = candidate - streak_candidate_sha: str | None = None + streak_sha: str | None = None for iteration in range(1, self.config.max_iterations + 1): current_sha = sha256_payload(current) @@ -345,115 +245,62 @@ class KernelPipeline: ) ) unanimous = all(verdict.verdict == "accept" for verdict in verdicts) - if unanimous: - if streak_candidate_sha not in {None, current_sha}: + if streak_sha not in {None, current_sha}: raise AssertionError("pass streak crossed candidate versions") - streak_candidate_sha = current_sha + streak_sha = current_sha pass_streak += 1 - record = IterationRecord( - iteration=iteration, - candidate_sha256=current_sha, - verdicts=verdicts, - unanimous=True, - pass_streak_after=pass_streak, - ) - iterations.append(record) - self.store.write_iteration(iteration, record) - if pass_streak >= self.config.streak_length: - result = KernelRunResult( - item_id=item.item_id, - status="accepted", - accepted_candidate=current, - iterations=iterations, - ) - self.store.write_final( - { - **result.model_dump(mode="json"), - "accepted_candidate_sha256": current_sha, - "human_audit_selected": self._audit_selected(item.item_id), - } - ) - return result - continue + else: + pass_streak = 0 + streak_sha = None - pass_streak = 0 - streak_candidate_sha = None - if iteration == self.config.max_iterations: - record = IterationRecord( - iteration=iteration, - candidate_sha256=current_sha, - verdicts=verdicts, - unanimous=False, - pass_streak_after=0, - ) - iterations.append(record) - self.store.write_iteration(iteration, record) - break - - repaired = await self._repair( - item=item, - dag=dag, - plan=plan, - candidate=current, - verdicts=verdicts, - iteration=iteration, - ) - repaired_sha = sha256_payload(repaired) record = IterationRecord( iteration=iteration, candidate_sha256=current_sha, verdicts=verdicts, - unanimous=False, - pass_streak_after=0, - repaired=True, - repaired_candidate_sha256=repaired_sha, + unanimous=unanimous, + pass_streak_after=pass_streak, ) iterations.append(record) self.store.write_iteration(iteration, record) - current = repaired + + if pass_streak == self.config.streak_length: + result = KernelRunResult( + item_id=item.item_id, + status="accepted", + plan=plan, + accepted_candidate=current, + iterations=iterations, + accepted_candidate_sha256=current_sha, + ) + self.store.write_final(result) + return result + + if not unanimous and iteration < self.config.max_iterations: + current = await self._repair( + item=item, + candidate=current, + verdicts=verdicts, + iteration=iteration, + ) result = KernelRunResult( item_id=item.item_id, status="rejected", + plan=plan, iterations=iterations, - rejection_reason=( - f"no two consecutive unanimous passes within " - f"{self.config.max_iterations} iterations" - ), - ) - self.store.write_final( - { - **result.model_dump(mode="json"), - "human_audit_selected": False, - } + rejection_reason="no two consecutive unanimous rounds within T=15", ) + self.store.write_final(result) return result async def run(self, item: CanonicalItem) -> KernelRunResult: self.store.write_input(item) self.store.write_config(self.config) - dag, plan, candidate = await self.build_candidate(item) + plan = await self.extract_plan(item) + candidate = await self.generate_candidate(item, plan) return await self.verify_candidate( item=item, - dag=dag, plan=plan, candidate=candidate, ) - - -def make_pipeline( - *, - proposer: JsonLLM, - judge_factory: Any, - run_dir: Path, - item_id: str, - config: PipelineConfig, -) -> KernelPipeline: - judges = [judge_factory(judge_id) for judge_id in range(1, 6)] - return KernelPipeline( - proposer=proposer, - judges=judges, - store=RunStore(run_dir, item_id), - config=config, - ) |