SGK86/procedural-graphs-repro

1"""Offline self-evolution loop (paper Sec 3.3, Algorithm 1).2 3Each round: diagnostic rollout on a train batch -> feedback-driven mutation by an4offline refiner LLM (topology + attribute edits) -> validation gate (commit only if5held-out performance does not decrease) -> rejection memory for rejected candidates.6 7The current graph's validation metrics are cached between rounds so the gate only runs8the candidate evaluation (the paper's gate compares against the last committed val score).9"""10import copy11import json12import time13 14from hotpot import em_score, evaluate_batch15 16REFINER_SYS = (17    "You are an offline refiner for a Procedural Graph used by a QA agent. You are given "18    "the current graph, traces from successful and failed tasks, and previously REJECTED "19    "candidate edits (do not re-propose them). Identify repeated error loops in failures "20    "and multi-step shortcuts in successes, then emit a compact edit set. Reply with ONLY "21    "a JSON object with keys: add_nodes (list of {name, description}), add_edges (list of "22    "{source, target, condition, guidance, pitfalls}), prune_edges (list of {source, target}), "23    "edit_edges (list of {source, target, condition, guidance, pitfalls}). All edge endpoints "24    "must reference nodes that exist after your add_nodes. Keep edits minimal (<=4 operations)."25)26 27 28def run_batch(agent, records):29    preds, traces = [], []30    t0 = time.time()31    for j, r in enumerate(records):32        out = agent.answer(r["question"], r["context"])33        preds.append(out["answer"])34        traces.append({"question": r["question"], "gold": r["answer"], "pred": out["answer"],35                       "node": out["node"], "guidance": out["guidance"]})36        if (j + 1) % 16 == 0:37            print(f"  [batch] {j+1}/{len(records)} records, {time.time()-t0:.0f}s", flush=True)38    return preds, traces, evaluate_batch(records, preds)39 40 41def diagnose(traces, records, preds, max_succ=6, max_fail=6):42    succ, fail = [], []43    for tr, rec, pred in zip(traces, records, preds):44        (succ if em_score(pred, rec["answer"]) else fail).append(tr)45    return succ[:max_succ], fail[:max_fail]46 47 48def make_edits(llm, graph, succ_traces, fail_traces, rejection_memory):49    user = (f"CURRENT GRAPH:\n{graph.serialize()}\n\n"50            f"SUCCESSFUL TRACES (high score):\n{json.dumps(succ_traces, indent=1, ensure_ascii=False)}\n\n"51            f"FAILED TRACES (low score):\n{json.dumps(fail_traces, indent=1, ensure_ascii=False)}\n\n"52            f"REJECTED CANDIDATES (do not re-propose):\n{json.dumps(rejection_memory[-4:], indent=1, ensure_ascii=False)}\n\n"53            "Emit the edit set JSON now.")54    edits = llm.chat_json([{"role": "system", "content": REFINER_SYS},55                           {"role": "user", "content": user}], max_tokens=1200)56    for key in ("add_nodes", "add_edges", "prune_edges", "edit_edges"):57        edits.setdefault(key, [])58    return edits59 60 61def evolve_round(llm, graph, agent, train_records, val_records, rejection_memory,62                 val_prev_metrics=None):63    """One evolution round; returns (new_graph, diagnostics, accepted, val_metrics)."""64    preds, traces, train_metrics = run_batch(agent, train_records)65    succ, fail = diagnose(traces, train_records, preds)66    edits = make_edits(llm, graph, succ, fail, rejection_memory)67    cand = copy.deepcopy(graph)68    try:69        applied = cand.apply_edits(edits)70        if not applied:71            raise ValueError("refiner proposed no applicable edits")72        if not cand.reachable() >= set(cand.nodes):73            raise ValueError("edit produced unreachable nodes")74    except Exception as e:75        print(f"[evolve] invalid candidate ({e}); rejected", flush=True)76        rejection_memory.append(edits)77        return graph, {"train_metrics": train_metrics, "applied": [], "error": str(e)}, False, None78    agent_cand = agent.__class__(llm, cand, mode=agent.mode)79    _, _, val_metrics = run_batch(agent_cand, val_records)80    if val_prev_metrics is None:81        val_prev_metrics = run_batch(agent.__class__(llm, graph, mode=agent.mode), val_records)[2]82    accepted = val_metrics["f1"] >= val_prev_metrics["f1"]83    print(f"[evolve] val f1: prev {val_prev_metrics['f1']:.2f} vs cand {val_metrics['f1']:.2f} -> "84          f"{'ACCEPT' if accepted else 'REJECT'} (applied: {applied})", flush=True)85    if not accepted:86        rejection_memory.append(edits)87        return graph, {"train_metrics": train_metrics, "val_prev": val_prev_metrics,88                       "val_cand": val_metrics, "applied": applied}, False, val_prev_metrics89    return cand, {"train_metrics": train_metrics, "val_prev": val_prev_metrics,90                  "val_cand": val_metrics, "applied": applied}, True, val_metrics