diff --git a/autoresearch/prefill/program.md b/autoresearch/prefill/program.md index bb03008..a41e990 100644 --- a/autoresearch/prefill/program.md +++ b/autoresearch/prefill/program.md @@ -57,6 +57,14 @@ benchmark here. Model/tokenizer/quantization/rope/window/cache-format changes must be deployed outside this supervisor. Cold benchmarks are explicit, separate invocations of `scripts/benchmark_prefill_architecture.py`. +Prefill budgets are hard admission limits, never truncation instructions. +Strategy input must fit 4096 tokens by carrying the complete active leaf +ancestry and its exact experiment records. Generator and Critic inputs must fit +6144 tokens; the Critic always receives the complete current Generator output. +If any complete semantic unit exceeds its budget, reject it before remote +Prefill and preserve the checkpoint. Never slice, sample, summarize, or drop +the tail of an over-budget input. + Do not optimize output wording, scores, prizes, or other proof-irrelevant content. Prefill performance is a tertiary objective after mathematical decomposition progress, while preserving the complete semantic contract. diff --git a/autoresearch/prefill/supervisor.py b/autoresearch/prefill/supervisor.py index 1ea0b8a..34770b6 100644 --- a/autoresearch/prefill/supervisor.py +++ b/autoresearch/prefill/supervisor.py @@ -6,6 +6,7 @@ import ast import csv import hashlib +import io import json import os import re @@ -328,7 +329,7 @@ def _extract_json(text: str) -> dict: if line.strip() and not line.strip().startswith(("#", "```")) ] - objective = " ".join(prose[:3])[:1200] + objective = " ".join(prose) if not objective: raise ValueError("strategy agent returned no usable candidate") digest = hashlib.sha256(stripped.encode()).hexdigest()[:12] @@ -338,7 +339,7 @@ def _extract_json(text: str) -> dict: target_matches[-1] if target_matches else "" ), "hypothesis": objective, - "plan": {"steps": steps[:8]}, + "plan": {"steps": steps}, "strategy_parse_mode": "prose", } @@ -417,6 +418,106 @@ def _pending_leaf_ids(ledger: dict) -> list[str]: return sorted(unresolved - unresolved_parents) +def build_strategy_research_state( + *, + current: dict, + ledger: dict, + results_text: str, +) -> dict: + target_id = _select_repair_target(current, ledger) + obligations = { + str(item.get("obligation_id", "")): item + for item in ledger.get("obligations", []) + } + ancestry = [] + cursor = target_id + visited = set() + while cursor and cursor not in visited: + visited.add(cursor) + item = obligations[cursor] + ancestry.append({ + "obligation_id": cursor, + "statement": item.get("statement", ""), + "status": item.get("status", ""), + "parent_id": item.get("parent_id", ""), + "last_run_id": item.get("last_run_id", ""), + "last_evidence": item.get("last_evidence", ""), + }) + cursor = str(item.get("parent_id", "")) + ancestry.reverse() + ancestry_ids = { + item["obligation_id"] for item in ancestry + } + relevant_results = [] + if results_text.strip(): + for row in csv.DictReader( + io.StringIO(results_text), + delimiter="\t", + ): + if row.get("target_obligation_id") not in ancestry_ids: + continue + relevant_results.append({ + "candidate_id": row.get("candidate_id", ""), + "target_obligation_id": row.get( + "target_obligation_id", + "", + ), + "hypothesis_sha256": row.get("hypothesis_sha256", ""), + "research_outcome": row.get("research_outcome", ""), + "research_evidence": row.get("research_evidence", ""), + "new_frontier": row.get("new_frontier", ""), + "kept": row.get("kept", ""), + "error": row.get("error", ""), + }) + return { + "target_leaf_id": target_id, + "target_ancestry": ancestry, + "relevant_experiments": relevant_results, + "current_candidate": { + "candidate_id": current.get("candidate_id", ""), + "target_obligation_id": current.get( + "target_obligation_id", + "", + ), + "hypothesis": current.get("hypothesis", ""), + "prefill_compute_chunk_tokens": current.get( + "prefill_compute_chunk_tokens", + ), + }, + } + + +def build_strategy_prompt( + *, + program: str, + current: dict, + results_text: str, + ledger: dict, +) -> str: + research_state = build_strategy_research_state( + current=current, + ledger=ledger, + results_text=results_text, + ) + return ( + "You are the AutoResearch strategy agent. Follow the human-owned " + "program exactly. Attack TARGET_LEAF_ID and propose " + "one falsifiable GAN strategy experiment. Return JSON only with keys: " + + ", ".join(REQUIRED_CANDIDATE_FIELDS) + + ". prefill_compute_chunk_tokens is immutable and must equal " + f"{current['prefill_compute_chunk_tokens']}. " + "Do not weaken full context, final-only snapshots, or no-fallback rules." + " The hypothesis must not repeat any hypothesis hash in RESEARCH_STATE. " + "It must either construct a concrete object or attempt a concrete " + "counterexample for the target leaf. target_obligation_id must equal " + "TARGET_LEAF_ID. Every statement and evidence item below is complete; " + "do not infer omitted text from unrelated branches." + f"\n\nPROGRAM:\n{program}" + "\n\nRESEARCH_STATE:\n" + f"{json.dumps(research_state, ensure_ascii=False)}" + ) + + class StrategyPrefillHeartbeat: def __init__( self, @@ -495,28 +596,18 @@ def propose_candidate( current: dict, results_text: str, ledger: dict, + max_prefill_tokens: int = 4096, ) -> dict: from kakeya import Client from transformers import AutoTokenizer from scripts.chat_grpc import _resolve_eos_token_ids tokenizer = AutoTokenizer.from_pretrained(tokenizer_id) - prompt = ( - "You are the AutoResearch strategy agent. Follow the human-owned " - "program exactly. Select one unresolved proof obligation and propose " - "one falsifiable GAN strategy experiment. Return JSON only with keys: " - + ", ".join(REQUIRED_CANDIDATE_FIELDS) - + ". prefill_compute_chunk_tokens is immutable and must equal " - f"{current['prefill_compute_chunk_tokens']}. " - "Do not weaken full context, final-only snapshots, or no-fallback rules." - " The hypothesis must not repeat any hypothesis or hash in RESULTS. " - "It must either construct a concrete object or attempt a concrete " - "counterexample for one smaller decomposition leaf. The " - "target_obligation_id must be one of PENDING_LEAF_IDS." - f"\n\nPROGRAM:\n{program}\n\nCURRENT:\n{json.dumps(current)}" - f"\n\nRESULTS:\n{results_text[-12000:]}" - f"\n\nLEDGER:\n{json.dumps(ledger)}" - f"\n\nPENDING_LEAF_IDS:\n{json.dumps(_pending_leaf_ids(ledger))}" + prompt = build_strategy_prompt( + program=program, + current=current, + results_text=results_text, + ledger=ledger, ) ids = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], @@ -525,6 +616,11 @@ def propose_candidate( return_dict=False, enable_thinking=False, ) + if len(ids) > max_prefill_tokens: + raise ValueError( + "Strategy Prefill token budget exceeded without truncation: " + f"{len(ids)} > {max_prefill_tokens}", + ) generated: list[int] = [] print( f"[autoresearch] Strategy Prefill: 0/{len(ids)} tokens (0.0%)", @@ -820,6 +916,7 @@ def run_iteration(args, iteration: int) -> dict: results_path.read_text() if results_path.exists() else "" ), ledger=ledger_data, + max_prefill_tokens=args.strategy_max_prefill_tokens, ) if proposed["target_obligation_id"] not in _pending_leaf_ids( ledger_data, @@ -987,6 +1084,11 @@ def main() -> int: ) parser.add_argument("--address", default="127.0.0.1:51051") parser.add_argument("--dashboard", default="http://127.0.0.1:8090") + parser.add_argument( + "--strategy-max-prefill-tokens", + type=int, + default=4096, + ) parser.add_argument( "--tokenizer-id", default=str( @@ -1014,6 +1116,8 @@ def main() -> int: args = parser.parse_args() if args.iterations <= 0: raise SystemExit("iterations must be > 0") + if args.strategy_max_prefill_tokens <= 0: + raise SystemExit("strategy-max-prefill-tokens must be > 0") for iteration in range(args.iterations): row = run_iteration(args, iteration) print(json.dumps(row, indent=2, sort_keys=True)) diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index cc7524e..e41d99f 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -670,6 +670,19 @@ def recover_checkpoint_from_log( ) +def enforce_prefill_token_budget( + stage: str, + token_ids, + max_tokens: int, +) -> None: + token_count = len(token_ids) + if token_count > max_tokens: + raise ValueError( + f"{stage} Prefill token budget exceeded without truncation: " + f"{token_count} > {max_tokens}", + ) + + def build_generator_messages( goal: str, *, @@ -914,6 +927,12 @@ def main() -> int: parser.add_argument("--api-key-file", default="~/.kakeya/network_api_key") parser.add_argument("--tokenizer-id", required=True) parser.add_argument("--output-tokens", type=int, default=64) + parser.add_argument( + "--max-prefill-tokens", + type=int, + default=6144, + help="Hard per-stage Prefill budget; over-budget input is rejected.", + ) parser.add_argument( "--max-response-tokens", type=int, @@ -968,6 +987,8 @@ def main() -> int: args = parser.parse_args() if args.output_tokens <= 0: raise SystemExit("output-tokens must be > 0") + if args.max_prefill_tokens <= 0: + raise SystemExit("max-prefill-tokens must be > 0") if args.auto_loop_boundary_wait_s < 0: raise SystemExit("auto-loop-boundary-wait-s must be >= 0") research_candidate = None @@ -1245,6 +1266,11 @@ def get_stats(): return_dict=False, enable_thinking=False, ) + enforce_prefill_token_budget( + "Generator", + generator_ids, + args.max_prefill_tokens, + ) print( f"[allens] Generator Prefill: {len(generator_ids)} tokens...", flush=True, @@ -1334,6 +1360,11 @@ def get_stats(): return_dict=False, enable_thinking=False, ) + enforce_prefill_token_budget( + "Critic", + critic_ids, + args.max_prefill_tokens, + ) print( f"[allens] Critic Prefill: {len(critic_ids)} tokens...", flush=True, diff --git a/tests/inference_engine/bench/test_autoresearch_supervisor.py b/tests/inference_engine/bench/test_autoresearch_supervisor.py index a5620be..1a7df9b 100644 --- a/tests/inference_engine/bench/test_autoresearch_supervisor.py +++ b/tests/inference_engine/bench/test_autoresearch_supervisor.py @@ -1,6 +1,7 @@ from autoresearch.prefill.supervisor import ( append_result, best_kept, + build_strategy_research_state, check_runtime_health, parse_research_verdict, read_results, @@ -281,6 +282,51 @@ def test_pending_leaf_ids_excludes_unresolved_parents(): assert _pending_leaf_ids(ledger) == ["RH-C1-child", "RH-C2"] +def test_strategy_state_keeps_complete_active_ancestry_only(): + ledger = {"obligations": [ + { + "obligation_id": "RH-C1", + "statement": "unrelated root", + "status": "UNRESOLVED", + "parent_id": "", + "last_evidence": "unrelated evidence", + }, + { + "obligation_id": "RH-C2", + "statement": "exact root statement", + "status": "UNRESOLVED", + "parent_id": "", + "last_evidence": "exact root evidence", + }, + { + "obligation_id": "RH-C2-child", + "statement": "exact child statement", + "status": "UNRESOLVED", + "parent_id": "RH-C2", + "last_evidence": "exact child evidence", + }, + ]} + results = ( + "candidate_id\ttarget_obligation_id\tresearch_outcome\t" + "research_evidence\tnew_frontier\tkept\terror\t" + "hypothesis_sha256\n" + "c1\tRH-C1\tINCONCLUSIVE\tunrelated result\tx\tFalse\t\th1\n" + "c2\tRH-C2-child\tDECOMPOSED\texact result\tfrontier\tTrue\t\th2\n" + ) + state = build_strategy_research_state( + current={**_candidate(), "target_obligation_id": "RH-C2"}, + ledger=ledger, + results_text=results, + ) + assert state["target_leaf_id"] == "RH-C2-child" + serialized = str(state) + assert "exact root statement" in serialized + assert "exact child evidence" in serialized + assert "exact result" in serialized + assert "unrelated root" not in serialized + assert "unrelated result" not in serialized + + def test_results_are_append_only_and_best_is_selected(tmp_path): path = tmp_path / "results.tsv" common = { @@ -390,6 +436,7 @@ def test_supervisor_preserves_runtime_and_cache_across_iterations(): assert "clear_primary_cache" not in source assert "launchctl" not in source assert "bootout" not in source + assert "results_text[-" not in source def test_gan_subprocess_output_is_streamed_not_captured(): diff --git a/tests/inference_engine/bridge/test_agent_gan_repl.py b/tests/inference_engine/bridge/test_agent_gan_repl.py index 1c36a1a..dee7321 100644 --- a/tests/inference_engine/bridge/test_agent_gan_repl.py +++ b/tests/inference_engine/bridge/test_agent_gan_repl.py @@ -20,6 +20,7 @@ build_critic_messages, build_generator_messages, extract_obligation_history, + enforce_prefill_token_budget, install_signal_protection, is_runtime_artifact_prompt, consume_critic_issue_batch, @@ -324,6 +325,19 @@ def test_generator_history_is_scoped_to_target_leaf(): assert "unrelated operator branch" not in prompt +def test_prefill_budget_rejects_whole_input_without_truncation(): + token_ids = list(range(7)) + enforce_prefill_token_budget("Generator", token_ids, 7) + try: + enforce_prefill_token_budget("Critic", token_ids, 6) + except ValueError as exc: + assert "without truncation" in str(exc) + assert "7 > 6" in str(exc) + else: + raise AssertionError("over-budget Prefill must be rejected") + assert token_ids == list(range(7)) + + def test_runtime_output_cannot_replace_research_goal(): for text in ( "critic> ### Central Claim",