Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions docs/ops/distributed-prefill-kv-network.md
Original file line number Diff line number Diff line change
Expand Up @@ -357,6 +357,13 @@ not wait for Terminal automation. A short boundary window still accepts an
explicit whitelisted command such as `/quit` or `/steer`. Any inference
exception pauses auto-loop and preserves the last complete checkpoint.
`--no-auto-loop` restores manual one-turn-at-a-time operation.
Rigorous mathematical problems discovered outside the running turn can be
queued in `~/.kakeya/agent_gan_critic_inbox.json`. The next turn appends every
pending issue to the previous Critic correction for Generator remediation and
to the new Critic steering for verification. The full issue list is timestamped
in the Critic log and recorded by ID/count in benchmark config. Issues remain
pending across failures and are marked consumed only after a complete,
successful Generator/Critic turn.
Generator output is always streamed in full to Terminal and passed verbatim to
the Gemma Critic. Sampling, truncation, summarization, independent chunk scores,
and semantic fallback are forbidden. A global Critic score is valid only when
Expand Down
113 changes: 111 additions & 2 deletions scripts/agent_gan_repl.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,65 @@ class ReplCheckpoint:
schema_version: int = 1


@dataclass
class CriticIssueBatch:
issue_id: str
issues: list[str]
status: str = "pending"
consumed_by_run: str = ""
schema_version: int = 1


def save_critic_issue_batch(path: Path, batch: CriticIssueBatch) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
temporary.write_text(
json.dumps(asdict(batch), ensure_ascii=False, indent=2),
encoding="utf-8",
)
os.chmod(temporary, 0o600)
temporary.replace(path)


def load_pending_critic_issues(path: Path) -> CriticIssueBatch | None:
if not path.exists():
return None
batch = CriticIssueBatch(
**json.loads(path.read_text(encoding="utf-8")),
)
if (
batch.schema_version != 1
or not batch.issue_id
or not batch.issues
or any(not str(issue).strip() for issue in batch.issues)
):
raise ValueError("invalid Critic issue inbox")
return batch if batch.status == "pending" else None


def format_critic_issue_injection(batch: CriticIssueBatch) -> str:
items = "\n".join(
f"{index}. {issue.strip()}"
for index, issue in enumerate(batch.issues, start=1)
)
return (
"\n\nEXTERNAL RIGOROUS MATHEMATICAL ISSUES "
f"(id={batch.issue_id}):\n{items}\n"
"The Generator must address every issue explicitly. The Critic must "
"verify every correction and keep unresolved items in the frontier."
)


def consume_critic_issue_batch(
path: Path,
batch: CriticIssueBatch,
run_id: str,
) -> None:
batch.status = "consumed"
batch.consumed_by_run = run_id
save_critic_issue_batch(path, batch)


def parse_repl_command(raw: str, phase: ReplPhase) -> ReplCommand:
text = raw.strip()
lower = text.lower()
Expand Down Expand Up @@ -476,6 +535,11 @@ def main() -> int:
default="~/.kakeya/agent_gan_state.json",
help="Private resumable Generator/Critic checkpoint.",
)
parser.add_argument(
"--critic-inbox",
default="~/.kakeya/agent_gan_critic_inbox.json",
help="Private pending rigorous-math issues for the next turn.",
)
parser.add_argument("--recover-run", default="")
parser.add_argument(
"--recover-log",
Expand Down Expand Up @@ -535,6 +599,7 @@ def get_stats():
return stats

state_path = Path(args.state_file).expanduser()
critic_inbox_path = Path(args.critic_inbox).expanduser()
if args.recover_run:
recovered = recover_checkpoint_from_log(
Path(args.recover_log).expanduser(),
Expand Down Expand Up @@ -571,6 +636,7 @@ def get_stats():
)
print(f"[log] {transcript.log_path}", flush=True)
print(f"[state] {state_path} phase={phase.value}", flush=True)
print(f"[critic-inbox] {critic_inbox_path}", flush=True)
if checkpoint:
print(
f"[state-restored] run={checkpoint.last_run_id or '(none)'} "
Expand Down Expand Up @@ -626,6 +692,28 @@ def get_stats():
if command.action in {"continue", "steer"} and args.auto_loop:
auto_loop_active = True
phase = ReplPhase.RUNNING
critic_issue_batch = load_pending_critic_issues(
critic_inbox_path,
)
critic_issue_injection = (
format_critic_issue_injection(critic_issue_batch)
if critic_issue_batch is not None else ""
)
if critic_issue_batch is not None:
print(
f"[critic-issue-injection] id="
f"{critic_issue_batch.issue_id} "
f"count={len(critic_issue_batch.issues)}",
flush=True,
)
for index, issue in enumerate(
critic_issue_batch.issues,
start=1,
):
print(
f"[critic-issue-{index}] {issue}",
flush=True,
)
run_nonce = uuid.uuid4().hex
telemetry_state["degraded"] = False
run = _telemetry_request(
Expand All @@ -644,6 +732,14 @@ def get_stats():
research_goal.encode(),
).hexdigest(),
"feedback_applied": bool(previous_critic),
"critic_issue_id": (
critic_issue_batch.issue_id
if critic_issue_batch is not None else ""
),
"critic_issue_count": (
len(critic_issue_batch.issues)
if critic_issue_batch is not None else 0
),
},
},
)
Expand All @@ -662,7 +758,9 @@ def get_stats():
research_goal,
steering=steering,
previous_generator=previous_generator,
previous_critic=previous_critic,
previous_critic=(
previous_critic + critic_issue_injection
),
)
generator_ids = tokenizer.apply_chat_template(
generator_messages,
Expand Down Expand Up @@ -728,7 +826,7 @@ def get_stats():
critic_messages = build_critic_messages(
research_goal,
critic_context,
steering=steering,
steering=steering + critic_issue_injection,
stop_reason=generator_actual["stop_reason"],
complete=generator_actual["complete"],
)
Expand Down Expand Up @@ -814,6 +912,17 @@ def get_stats():
last_run_id=run_id,
),
)
if critic_issue_batch is not None:
consume_critic_issue_batch(
critic_inbox_path,
critic_issue_batch,
run_id,
)
print(
f"[critic-issues-consumed] id="
f"{critic_issue_batch.issue_id} run={run_id}",
flush=True,
)
phase = ReplPhase.READY
if auto_loop_active:
print(
Expand Down
32 changes: 32 additions & 0 deletions tests/inference_engine/bridge/test_agent_gan_repl.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import io
import json
import signal
import time
from pathlib import Path

from scripts.agent_gan_repl import (
PrefillHeartbeat,
CriticIssueBatch,
ReplCheckpoint,
ReplPhase,
TimestampedTee,
Expand All @@ -16,9 +18,13 @@
build_generator_messages,
install_signal_protection,
is_runtime_artifact_prompt,
consume_critic_issue_batch,
format_critic_issue_injection,
load_checkpoint,
load_pending_critic_issues,
parse_repl_command,
recover_checkpoint_from_log,
save_critic_issue_batch,
save_checkpoint,
)

Expand Down Expand Up @@ -328,6 +334,32 @@ def test_checkpoint_round_trip_is_private(tmp_path):
assert load_checkpoint(tmp_path / "missing.json") is None


def test_critic_issue_inbox_retries_until_successful_consumption(tmp_path):
path = tmp_path / "critic-inbox.json"
batch = CriticIssueBatch(
issue_id="math-review-1",
issues=[
"Do not identify -zeta'/zeta with xi.",
"Prove zero convergence before invoking Hurwitz.",
],
)
save_critic_issue_batch(path, batch)
loaded = load_pending_critic_issues(path)
assert loaded == batch
assert path.stat().st_mode & 0o777 == 0o600
injection = format_critic_issue_injection(loaded)
assert "math-review-1" in injection
assert "1. Do not identify" in injection
assert "2. Prove zero convergence" in injection
# Merely loading/injecting does not consume the issues.
assert load_pending_critic_issues(path).status == "pending"
consume_critic_issue_batch(path, loaded, "br_success")
assert load_pending_critic_issues(path) is None
persisted = json.loads(path.read_text())
assert persisted["status"] == "consumed"
assert persisted["consumed_by_run"] == "br_success"


def test_recover_complete_checkpoint_from_timestamped_log(tmp_path):
path = tmp_path / "agent.log"
path.write_text(
Expand Down
Loading