From 473d829c9e7b8fd390d19096566cc3f16ee3c494 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Fri, 17 Jul 2026 12:14:30 +0800 Subject: [PATCH] fix(agents): reuse deterministic interactive Prefill KV Remove per-run prompt entropy and compact Agent GAN templates so repeated tasks can hit allens KV while new Critic contexts require substantially fewer Prefill tokens. Co-authored-by: Cursor --- docs/ops/distributed-prefill-kv-network.md | 4 +- scripts/agent_gan_repl.py | 92 ++++++++++--------- .../bridge/test_agent_gan_repl.py | 25 +++++ 3 files changed, 78 insertions(+), 43 deletions(-) diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 5d62ab7..f3ccf7b 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -344,9 +344,11 @@ The REPL ignores external `SIGTERM`; its shell supervisor restarts signal-based exits. Only `/quit`, `/exit`, or EOF is treated as approval to stop. Generator output is always streamed in full to Terminal. To keep the 16GB allens Critic Prefill interactive, Critic receives a labeled extractive evidence -window (default 128 Generator tokens: beginning + conclusion) with the omitted +window (default 64 Generator tokens: beginning + conclusion) with the omitted token count and EOS status. It must not interpret evidence-window omission as Generator truncation. Long Prefill operations emit a heartbeat every 30 seconds. +Interactive prompt templates are deterministic and contain no per-run nonce, so +repeating the same task can reuse allens cold-tier and Primary hot-tier KV. ## Rollback diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index ace3c82..44281cc 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -46,6 +46,46 @@ def _telemetry_request(url: str, **kwargs): return None +def build_generator_messages(prompt: str) -> list[dict[str, str]]: + return [ + { + "role": "system", + "content": ( + "Answer rigorously. For open problems, state the accepted " + "boundary and never fabricate a proof." + ), + }, + {"role": "user", "content": prompt}, + ] + + +def build_critic_messages( + prompt: str, + evidence: str, + *, + stop_reason: str, + complete: bool, +) -> list[dict[str, str]]: + return [ + { + "role": "system", + "content": ( + "Score the answer 0-10, identify false claims, and give " + "specific corrections. Do not penalize an honest statement " + "that an open problem is unsolved. Evidence is intentionally " + "bounded; omitted tokens do not imply truncation." + ), + }, + { + "role": "user", + "content": ( + f"Task:\n{prompt}\n\nEvidence:\n{evidence}\n\n" + f"Completion: {stop_reason}; complete={complete}" + ), + }, + ] + + class TokenPrinter: def __init__(self, tokenizer, label: str) -> None: self.tokenizer = tokenizer @@ -132,9 +172,11 @@ def main() -> int: default=0, help="Optional client response cap; 0 means generate until model EOS.", ) - parser.add_argument("--critic-evidence-tokens", type=int, default=128) + parser.add_argument("--critic-evidence-tokens", type=int, default=64) parser.add_argument("--skip-ensure", action="store_true") args = parser.parse_args() + if min(args.output_tokens, args.critic_evidence_tokens) <= 0: + raise SystemExit("output-tokens and critic-evidence-tokens must be > 0") from kakeya import Client from transformers import AutoTokenizer @@ -195,20 +237,7 @@ def get_stats(): remote_run = run is not None run_id = run["id"] if remote_run else f"local_{run_nonce[:16]}" try: - generator_messages = [ - { - "role": "system", - "content": ( - "You are the Generator agent. Produce a concrete, " - "technically rigorous answer. For open or unsolved " - "problems, state the accepted boundary honestly, " - "provide rigorous context, and never fabricate a " - "proof. Internal run " - f"{run_nonce}." - ), - }, - {"role": "user", "content": prompt}, - ] + generator_messages = build_generator_messages(prompt) generator_ids = tokenizer.apply_chat_template( generator_messages, add_generation_prompt=True, @@ -260,33 +289,12 @@ def get_stats(): generator_text, args.critic_evidence_tokens, ) - critic_messages = [ - { - "role": "system", - "content": ( - "You are the Critic/Discriminator. Score the answer " - "0-10, identify false assumptions, and propose " - "specific corrections. Do not penalize a correct " - "statement that an open problem has no accepted " - "proof. Call an answer incomplete only when its " - "completion status is not EOS or its syntax is " - "visibly cut off. " - "You receive a bounded evidence window; omitted " - "tokens are not evidence of Generator truncation. " - f"Internal run {run_nonce}." - ), - }, - { - "role": "user", - "content": ( - f"Original task:\n{prompt}\n\n" - f"Generator evidence window:\n{evidence}\n\n" - "Generator completion status: " - f"{generator_actual['stop_reason']}; " - f"complete={generator_actual['complete']}" - ), - }, - ] + critic_messages = build_critic_messages( + prompt, + evidence, + stop_reason=generator_actual["stop_reason"], + complete=generator_actual["complete"], + ) critic_ids = tokenizer.apply_chat_template( critic_messages, add_generation_prompt=True, diff --git a/tests/inference_engine/bridge/test_agent_gan_repl.py b/tests/inference_engine/bridge/test_agent_gan_repl.py index 4847e8e..aff7bb7 100644 --- a/tests/inference_engine/bridge/test_agent_gan_repl.py +++ b/tests/inference_engine/bridge/test_agent_gan_repl.py @@ -7,6 +7,8 @@ TokenPrinter, _stage, _telemetry_request, + build_critic_messages, + build_generator_messages, install_signal_protection, ) @@ -139,3 +141,26 @@ def timeout(*_args, **_kwargs): output = capsys.readouterr().out assert "telemetry-warning" in output assert "inference will continue" in output + + +def test_interactive_prompts_are_deterministic_for_kv_reuse(): + generator_a = build_generator_messages("prove RH") + generator_b = build_generator_messages("prove RH") + critic_a = build_critic_messages( + "prove RH", + "bounded evidence", + stop_reason="eos", + complete=True, + ) + critic_b = build_critic_messages( + "prove RH", + "bounded evidence", + stop_reason="eos", + complete=True, + ) + assert generator_a == generator_b + assert critic_a == critic_b + combined = repr(generator_a + critic_a) + assert "Internal run" not in combined + assert "open problem" in combined + assert "omitted tokens do not imply truncation" in combined