1"""E4: repeated handoffs (relay chains) with an evaluator-only fact ledger.23Chain: agent_1 reads the document and writes m_1 (A4 prose or A5 structured, question-blind). agent_i (i>1) reads only4m_{i-1} (closed) or m_{i-1} plus metered source lookups (repairable) and writes m_i under the same budget. After hops5h in HOPS the receiver answers all questions from m_h. Retention/distortion is scored from the ledger (never shown to agents).6"""7from __future__ import annotations8import json, re, time, yaml9from pathlib import Path10import mlx.core as mx11from .util import ROOT, read_json, sha256_json, sha256_text, utc_now, read_jsonl12from .records import new_record, write_record, trial_key13from .data.prepare import load_doc, load_key14from .scoring import score_synthetic, normalize_answer, extract_numbers15from .prompts import SYSTEM_RECEIVER, receiver_user_text, SUMMARY_SYSTEM, EVIDENCE_SYSTEM16from .backends.mlx_backend import MLXModel, reset_peak17from .methods import text as TM1819import platform as _pf20HARDWARE_ID = {"MacBookPro-de-simon-pierre-2.local": "laptop-m5max-48g", "MacStudnpierre2": "m3u96a-m3ultra-96g", "m3u96b": "m3u96b-m3ultra-96g"}.get(_pf.node().split(".")[0] if _pf.node().startswith("MacStud") or _pf.node().startswith("m3u") else _pf.node(), _pf.node())2122RELAY_SUMMARY_USER = ("You received the following notes from a colleague about a document you cannot see. Rewrite them as a faithful summary of at most "23 "{budget} tokens for the next colleague, preserving every concrete fact (names, identifiers, numbers, units, dates, statuses, negations) "24 "exactly as written. Do not add facts that are not in the notes.{repair}\n\nNotes:\n{msg}\n\nSummary:")25RELAY_EVIDENCE_USER = ("You received the following structured evidence (JSON) from a colleague about a document you cannot see. Re-emit it as a JSON object with the same "26 "schema (keys facts, optional notes, optional unknowns), under {budget} tokens, keeping every fact and value exactly; drop nothing unless the "27 "budget forces you to, and never invent values.{repair}\n\nEvidence:\n{msg}\n\nJSON:")2829def _ledger_metrics(message: str, facts: list[dict], doc_text: str) -> dict:30 """Retention: a fact counts as retained if its value (normalized) appears in the message. Unsupported additions:31 numbers with >=4 significant digits, identifiers (XX-00X-0000) and 4-digit-year dates in the message that do not32 occur in the document."""33 # structured messages carry character-offset spans ("span": [a, b]); strip them so offsets are not counted as numbers34 message = re.sub(r'"span":\s*\[[^\]]*\]', '"span": []', message)35 msg_n = normalize_answer(message); msg_nc = message.replace(",", "")36 retained = 0; by_kind = {}37 for f in facts:38 v = f["value"]39 ok = normalize_answer(v) in msg_n or v.replace(",", "") in msg_nc40 by_kind.setdefault(f["kind"], [0, 0]); by_kind[f["kind"]][1] += 141 if ok:42 retained += 1; by_kind[f["kind"]][0] += 143 doc_nums = {str(int(x)) for x in extract_numbers(doc_text) if abs(x) >= 1000}44 msg_nums = {str(int(x)) for x in extract_numbers(message) if abs(x) >= 1000}45 unsupported_nums = msg_nums - doc_nums46 doc_ids = set(re.findall(r"\b[A-Z]{2}-\d{2}[A-Z]-\d{4}\b", doc_text)); msg_ids = set(re.findall(r"\b[A-Z]{2}-\d{2}[A-Z]-\d{4}\b", message))47 unsupported_ids = msg_ids - doc_ids48 # distorted amounts: a project code co-occurring (same sentence) with an unsupported number49 distorted = 050 for sent in re.split(r"(?<=[.!?])\s+|\n+", message):51 if re.search(r"\b[A-Z][a-z]+-\d{1,2}\b", sent) and any(n in sent.replace(",", "") for n in unsupported_nums):52 distorted += 153 return {"facts_total": len(facts), "facts_retained": retained, "retention_rate": retained / max(1, len(facts)),54 "retention_by_kind": {k: v[0] / v[1] for k, v in by_kind.items()},55 "unsupported_numbers": len(unsupported_nums), "unsupported_identifiers": len(unsupported_ids), "distorted_sentences": distorted,56 "message_bytes": len(message.encode("utf-8"))}5758def _repair_lookup(prev_msg: str, doc_text: str, max_bytes: int, max_lookups: int) -> tuple[str, dict]:59 """Repairable chain: the agent re-reads source sentences for entity codes mentioned in the notes (metered)."""60 codes = []61 for c in re.findall(r"\b[A-Z][a-z]+-\d{1,2}\b", prev_msg):62 if c not in codes: codes.append(c)63 sents = re.split(r"(?<=[.!?])\s+", doc_text)64 fetched, nb, n = [], 0, 065 for c in codes:66 if n >= max_lookups: break67 hits = [s for s in sents if c in s]68 if not hits: continue69 chunk = " ".join(hits)70 b = len(chunk.encode("utf-8"))71 if nb + b > max_bytes: continue72 fetched.append(f"[source passages for {c}] {chunk}"); nb += b; n += 173 return "\n".join(fetched), {"lookups": n, "fetched_bytes": nb, "codes_mentioned": len(codes)}7475def run_chains(config_path: str, resume: bool = False) -> int:76 cfg = yaml.safe_load(open(config_path))77 mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]78 models = {}79 def get(a):80 if a not in models: models[a] = MLXModel(mcfg[a]["repo"])81 return models[a]82 man = read_json(ROOT / "data/manifests/splits.json")83 items = [it for it in man["items"] if it["family"] == "synthetic" and it["split"] in cfg["splits"] and it.get("target_tokens") in cfg["lengths"]]84 if cfg.get("limit_docs"): items = items[: cfg["limit_docs"]]85 items = items[cfg.get("doc_offset", 0):: cfg.get("doc_stride", 1)] # sharding across machines (same experiment id)86 out_path = ROOT / "runs/raw" / f"{cfg['experiment_id']}.jsonl"87 log_path = ROOT / "runs/logs" / f"{cfg['experiment_id']}.log"88 done = {r["trial_id"] for r in read_jsonl(out_path) if r["status"] == "success"} if (resume and out_path.exists()) else set()89 protocol_hash = cfg.get("protocol_hash", "unfrozen"); config_hash = sha256_json(cfg)90 def log(m):91 line = f"{utc_now()} {m}"; print(line, flush=True)92 with open(log_path, "a") as f: f.write(line + "\n")93 log(f"chains start {cfg['experiment_id']} items={len(items)} done={len(done)}")94 HOPS = cfg["hops"]; budget = cfg["budget_tokens"]95 for it in items:96 doc = load_doc(it["document_id"], "synthetic"); key = load_key(it["document_id"], "synthetic")97 qkeys = {q["question_id"]: q for q in key["questions"]}98 for chain in cfg["chains"]:99 rep, mode, agents = chain["representation"], chain["mode"], chain["agents"] # agents: list of aliases cycled100 cid = f"{rep}|{mode}|{'-'.join(agents)}"101 # skip if all hop trials done102 keys_needed = [trial_key(cfg["experiment_id"], it["document_id"], q["question_id"], rep, {"mode": mode, "agents": agents, "budget": budget}, agents[0], agents[-1], "blind", protocol_hash, 0, "chain", h) for h in HOPS for q in doc["questions"]]103 if all(k in done for k in keys_needed):104 continue105 msg = None; hop_msgs = {}; hop_meta = {}106 t_chain = 0.0; fetched_total = 0; lookups_total = 0107 for h in range(1, max(HOPS) + 1):108 agent = get(agents[(h - 1) % len(agents)])109 t0 = time.perf_counter()110 repair_txt, rinfo = "", {"lookups": 0, "fetched_bytes": 0}111 if h == 1:112 m = TM.a4_summary(doc["text"], budget, agent, None) if rep == "A4" else TM.a5_structured(doc["text"], budget, agent, None)113 text = m.text; usage = m.sender_usage114 else:115 if mode == "repairable":116 repair_txt, rinfo = _repair_lookup(msg, doc["text"], cfg.get("max_fetch_bytes", 1500), cfg.get("max_lookups", 4))117 rep_clause = ("\n\nYou may also use these passages re-read from the original source:\n" + repair_txt) if repair_txt else ""118 if rep == "A4":119 user = RELAY_SUMMARY_USER.format(budget=budget, msg=msg, repair=rep_clause)120 text, usage, _ = TM._sender_generate(agent, SUMMARY_SYSTEM, user, budget)121 else:122 user = RELAY_EVIDENCE_USER.format(budget=budget, msg=msg, repair=rep_clause)123 raw, usage, _ = TM._sender_generate(agent, EVIDENCE_SYSTEM, user, budget)124 obj = TM._parse_json_loose(raw)125 text = json.dumps(obj, ensure_ascii=False) if isinstance(obj, dict) else raw126 hop_ms = (time.perf_counter() - t0) * 1000; t_chain += hop_ms127 fetched_total += rinfo["fetched_bytes"]; lookups_total += rinfo["lookups"]128 msg = text; hop_msgs[h] = text129 hop_meta[h] = {"hop_ms": hop_ms, "agent": agent.info.repo_id, "usage": usage, "repair": rinfo, "ledger": _ledger_metrics(text, key["facts"], doc["text"])}130 if h in HOPS:131 receiver = get(agents[(h - 1) % len(agents)]) # the last agent in the chain answers132 for q in doc["questions"]:133 tk = trial_key(cfg["experiment_id"], it["document_id"], q["question_id"], rep, {"mode": mode, "agents": agents, "budget": budget}, agents[0], agents[-1], "blind", protocol_hash, 0, "chain", h)134 if tk in done: continue135 rec = new_record(experiment_id=cfg["experiment_id"], trial_id=tk, run_group_id=f"{cfg['experiment_id']}:{cid}:h{h}", protocol_hash=protocol_hash, config_hash=config_hash,136 task_family="synthetic", document_id=it["document_id"], document_hash=it["document_hash"], question_id=q["question_id"], split=it["split"], cluster_id=it["cluster_id"],137 sender_model_revision="|".join(f"{get(a).info.repo_id}@{get(a).info.revision}" for a in agents), receiver_model_revision=f"{receiver.info.repo_id}@{receiver.info.revision}",138 tokenizer_hashes={"receiver": receiver.info.tokenizer_hash}, backend_versions={"mlx": mx.__version__, "mlx_lm": receiver.info.mlx_lm_version}, hardware_id=HARDWARE_ID,139 device_state={}, weight_precision="bf16", method_id=rep, implementation_fidelity="custom", compression_config={"mode": mode, "agents": agents, "budget_tokens": budget},140 question_access="blind", topology=f"chain_{mode}", handoff_index=h, reuse_count=len(doc["questions"]), cache_hit=False, cache_compatibility_hash=None, random_seed=0,141 source_tokens_sender=None, source_tokens_receiver=len(receiver.encode(doc["text"])), message_utf8_bytes=len(text.encode("utf-8")), serialized_payload_bytes=len(text.encode("utf-8")),142 cache_tensor_bytes=None, side_information_bytes=0, fetched_bytes=fetched_total, timing_spans={"chain_ms": t_chain, "hops": {str(k): v["hop_ms"] for k, v in hop_meta.items()}},143 input_usage={"lookups": lookups_total}, output_usage={}, provider_cached_usage=None, process_memory=None, backend_memory=None, peak_prefill_memory=None, peak_decode_memory=None,144 mapper_memory=None, cost_with_source=None, evidence_artifact={"ledger": hop_meta[h]["ledger"], "ledger_by_hop": {str(k): v["ledger"] for k, v in hop_meta.items()}},145 reference_key_id=f"keys/synthetic/{it['document_id']}.json", status="implementation_error", error_code=None, log_artifact=str(log_path.relative_to(ROOT)), payload_hash=sha256_text(text), output_hash=None)146 try:147 ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text("Notes received after %d handoff(s)" % h, text, q["question"], "synthetic")}])148 t1 = time.perf_counter(); g = receiver.generate_text(ids, max_tokens=48)149 rec.update({"prediction_artifact": g["text"], "output_hash": sha256_text(g["text"]), "query_tokens": len(ids), "message_tokens": len(receiver.encode(text)),150 "ttft_ms": g["ttft_ms"], "total_wall_ms": (time.perf_counter() - t1) * 1000 + t_chain, "output_usage": {"output_tokens": g["output_tokens"]},151 "task_scores": score_synthetic(g["text"], qkeys[q["question_id"]]), "status": "success" if g["text"].strip() else "invalid_output"})152 rec["timing_spans"].update({"receiver_prefill_ms": g["prefill_ms"], "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"]})153 except Exception as e:154 rec["error_code"] = f"{type(e).__name__}: {str(e)[:200]}"155 write_record(rec, out_path)156 L = hop_meta[h]["ledger"]157 log(f"{cid:>28s} {it['document_id']} hop={h} retention={L['retention_rate']:.2f} unsupported={L['unsupported_numbers']}+{L['unsupported_identifiers']} bytes={L['message_bytes']} fetched={fetched_total}")158 # store hop messages for audit159 Path(ROOT / "runs/raw/chains").mkdir(parents=True, exist_ok=True)160 (ROOT / "runs/raw/chains" / f"{cfg['experiment_id']}_{it['document_id']}_{cid.replace('|','_').replace('-','_')}.json").write_text(json.dumps({"hops": hop_msgs, "meta": hop_meta}, ensure_ascii=False, indent=1, default=str))161 mx.clear_cache()162 log("chains end")163 return 0164