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 TM1819HARDWARE_ID = "laptop-m5max-48g"2021RELAY_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 "22 "{budget} tokens for the next colleague, preserving every concrete fact (names, identifiers, numbers, units, dates, statuses, negations) "23 "exactly as written. Do not add facts that are not in the notes.{repair}\n\nNotes:\n{msg}\n\nSummary:")24RELAY_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 "25 "schema (keys facts, optional notes, optional unknowns), under {budget} tokens, keeping every fact and value exactly; drop nothing unless the "26 "budget forces you to, and never invent values.{repair}\n\nEvidence:\n{msg}\n\nJSON:")2728def _ledger_metrics(message: str, facts: list[dict], doc_text: str) -> dict:29 """Retention: a fact counts as retained if its value (normalized) appears in the message. Unsupported additions:30 numbers with >=4 significant digits, identifiers (XX-00X-0000) and 4-digit-year dates in the message that do not31 occur in the document."""32 # structured messages carry character-offset spans ("span": [a, b]); strip them so offsets are not counted as numbers33 message = re.sub(r'"span":\s*\[[^\]]*\]', '"span": []', message)34 msg_n = normalize_answer(message); msg_nc = message.replace(",", "")35 retained = 0; by_kind = {}36 for f in facts:37 v = f["value"]38 ok = normalize_answer(v) in msg_n or v.replace(",", "") in msg_nc39 by_kind.setdefault(f["kind"], [0, 0]); by_kind[f["kind"]][1] += 140 if ok:41 retained += 1; by_kind[f["kind"]][0] += 142 doc_nums = {str(int(x)) for x in extract_numbers(doc_text) if abs(x) >= 1000}43 msg_nums = {str(int(x)) for x in extract_numbers(message) if abs(x) >= 1000}44 unsupported_nums = msg_nums - doc_nums45 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))46 unsupported_ids = msg_ids - doc_ids47 # distorted amounts: a project code co-occurring (same sentence) with an unsupported number48 distorted = 049 for sent in re.split(r"(?<=[.!?])\s+|\n+", message):50 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):51 distorted += 152 return {"facts_total": len(facts), "facts_retained": retained, "retention_rate": retained / max(1, len(facts)),53 "retention_by_kind": {k: v[0] / v[1] for k, v in by_kind.items()},54 "unsupported_numbers": len(unsupported_nums), "unsupported_identifiers": len(unsupported_ids), "distorted_sentences": distorted,55 "message_bytes": len(message.encode("utf-8"))}5657def _repair_lookup(prev_msg: str, doc_text: str, max_bytes: int, max_lookups: int) -> tuple[str, dict]:58 """Repairable chain: the agent re-reads source sentences for entity codes mentioned in the notes (metered)."""59 codes = []60 for c in re.findall(r"\b[A-Z][a-z]+-\d{1,2}\b", prev_msg):61 if c not in codes: codes.append(c)62 sents = re.split(r"(?<=[.!?])\s+", doc_text)63 fetched, nb, n = [], 0, 064 for c in codes:65 if n >= max_lookups: break66 hits = [s for s in sents if c in s]67 if not hits: continue68 chunk = " ".join(hits)69 b = len(chunk.encode("utf-8"))70 if nb + b > max_bytes: continue71 fetched.append(f"[source passages for {c}] {chunk}"); nb += b; n += 172 return "\n".join(fetched), {"lookups": n, "fetched_bytes": nb, "codes_mentioned": len(codes)}7374def run_chains(config_path: str, resume: bool = False) -> int:75 cfg = yaml.safe_load(open(config_path))76 mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]77 models = {}78 def get(a):79 if a not in models: models[a] = MLXModel(mcfg[a]["repo"])80 return models[a]81 man = read_json(ROOT / "data/manifests/splits.json")82 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"]]83 if cfg.get("limit_docs"): items = items[: cfg["limit_docs"]]84 out_path = ROOT / "runs/raw" / f"{cfg['experiment_id']}.jsonl"85 log_path = ROOT / "runs/logs" / f"{cfg['experiment_id']}.log"86 done = {r["trial_id"] for r in read_jsonl(out_path) if r["status"] == "success"} if (resume and out_path.exists()) else set()87 protocol_hash = cfg.get("protocol_hash", "unfrozen"); config_hash = sha256_json(cfg)88 def log(m):89 line = f"{utc_now()} {m}"; print(line, flush=True)90 with open(log_path, "a") as f: f.write(line + "\n")91 log(f"chains start {cfg['experiment_id']} items={len(items)} done={len(done)}")92 HOPS = cfg["hops"]; budget = cfg["budget_tokens"]93 for it in items:94 doc = load_doc(it["document_id"], "synthetic"); key = load_key(it["document_id"], "synthetic")95 qkeys = {q["question_id"]: q for q in key["questions"]}96 for chain in cfg["chains"]:97 rep, mode, agents = chain["representation"], chain["mode"], chain["agents"] # agents: list of aliases cycled98 cid = f"{rep}|{mode}|{'-'.join(agents)}"99 # skip if all hop trials done100 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"]]101 if all(k in done for k in keys_needed):102 continue103 msg = None; hop_msgs = {}; hop_meta = {}104 t_chain = 0.0; fetched_total = 0; lookups_total = 0105 for h in range(1, max(HOPS) + 1):106 agent = get(agents[(h - 1) % len(agents)])107 t0 = time.perf_counter()108 repair_txt, rinfo = "", {"lookups": 0, "fetched_bytes": 0}109 if h == 1:110 m = TM.a4_summary(doc["text"], budget, agent, None) if rep == "A4" else TM.a5_structured(doc["text"], budget, agent, None)111 text = m.text; usage = m.sender_usage112 else:113 if mode == "repairable":114 repair_txt, rinfo = _repair_lookup(msg, doc["text"], cfg.get("max_fetch_bytes", 1500), cfg.get("max_lookups", 4))115 rep_clause = ("\n\nYou may also use these passages re-read from the original source:\n" + repair_txt) if repair_txt else ""116 if rep == "A4":117 user = RELAY_SUMMARY_USER.format(budget=budget, msg=msg, repair=rep_clause)118 text, usage, _ = TM._sender_generate(agent, SUMMARY_SYSTEM, user, budget)119 else:120 user = RELAY_EVIDENCE_USER.format(budget=budget, msg=msg, repair=rep_clause)121 raw, usage, _ = TM._sender_generate(agent, EVIDENCE_SYSTEM, user, budget)122 obj = TM._parse_json_loose(raw)123 text = json.dumps(obj, ensure_ascii=False) if isinstance(obj, dict) else raw124 hop_ms = (time.perf_counter() - t0) * 1000; t_chain += hop_ms125 fetched_total += rinfo["fetched_bytes"]; lookups_total += rinfo["lookups"]126 msg = text; hop_msgs[h] = text127 hop_meta[h] = {"hop_ms": hop_ms, "agent": agent.info.repo_id, "usage": usage, "repair": rinfo, "ledger": _ledger_metrics(text, key["facts"], doc["text"])}128 if h in HOPS:129 receiver = get(agents[(h - 1) % len(agents)]) # the last agent in the chain answers130 for q in doc["questions"]:131 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)132 if tk in done: continue133 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,134 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"],135 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}",136 tokenizer_hashes={"receiver": receiver.info.tokenizer_hash}, backend_versions={"mlx": mx.__version__, "mlx_lm": receiver.info.mlx_lm_version}, hardware_id=HARDWARE_ID,137 device_state={}, weight_precision="bf16", method_id=rep, implementation_fidelity="custom", compression_config={"mode": mode, "agents": agents, "budget_tokens": budget},138 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,139 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")),140 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()}},141 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,142 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()}},143 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)144 try:145 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")}])146 t1 = time.perf_counter(); g = receiver.generate_text(ids, max_tokens=48)147 rec.update({"prediction_artifact": g["text"], "output_hash": sha256_text(g["text"]), "query_tokens": len(ids), "message_tokens": len(receiver.encode(text)),148 "ttft_ms": g["ttft_ms"], "total_wall_ms": (time.perf_counter() - t1) * 1000 + t_chain, "output_usage": {"output_tokens": g["output_tokens"]},149 "task_scores": score_synthetic(g["text"], qkeys[q["question_id"]]), "status": "success" if g["text"].strip() else "invalid_output"})150 rec["timing_spans"].update({"receiver_prefill_ms": g["prefill_ms"], "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"]})151 except Exception as e:152 rec["error_code"] = f"{type(e).__name__}: {str(e)[:200]}"153 write_record(rec, out_path)154 L = hop_meta[h]["ledger"]155 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}")156 # store hop messages for audit157 Path(ROOT / "runs/raw/chains").mkdir(parents=True, exist_ok=True)158 (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))159 mx.clear_cache()160 log("chains end")161 return 0162