1"""Track A: text, structured messages and evidence (A0-A8). Each encoder returns a Message with realized budgets.23Question access: `question=None` -> question-blind; otherwise question-aware. Budgets are receiver tokens (recorded as4realized counts under the receiver tokenizer) and UTF-8 bytes.5"""6from __future__ import annotations7import json, math, re, time, gzip8from dataclasses import dataclass, field, asdict9from typing import Any, Optional10from ..prompts import SUMMARY_SYSTEM, summary_user, EVIDENCE_SYSTEM, evidence_user11from ..cache.serialize import codec_compress, codec_decompress1213@dataclass14class Message:15 method_id: str16 text: str # material the receiver reads (after decoding)17 material_label: str18 utf8_bytes: int19 payload_bytes: int # bytes actually transported (after codec, if any)20 side_information_bytes: int # e.g. reference index, codec headers21 encode_ms: float22 decode_ms: float23 sender_usage: dict = field(default_factory=dict) # sender tokens in/out24 meta: dict = field(default_factory=dict)25 artifacts: dict = field(default_factory=dict) # e.g., evidence index for A6 retrieval2627def _sentences(text: str) -> list[str]:28 parts = re.split(r"(?<=[.!?])\s+|\n+", text)29 return [p.strip() for p in parts if p.strip()]3031def _tok(text: str) -> list[str]:32 return re.findall(r"[a-z0-9]+(?:[.,'-][a-z0-9]+)*", text.lower())3334# ---------------------------------------------------------------- A0 / A135def a0_no_message() -> Message:36 return Message("A0", "", "Material", 0, 0, 0, 0.0, 0.0, meta={"note": "receiver gets no source context"})3738def a1_full(doc: str) -> Message:39 b = len(doc.encode("utf-8"))40 return Message("A1", doc, "Document", b, b, 0, 0.0, 0.0)4142# ---------------------------------------------------------------- A2 head/tail truncation (token budget under receiver tokenizer)43def a2_truncate(doc: str, budget_tokens: int, tokenizer, head_frac: float = 0.5) -> Message:44 t0 = time.perf_counter()45 ids = tokenizer.encode(doc, add_special_tokens=False)46 if len(ids) <= budget_tokens:47 text = doc48 else:49 h = int(budget_tokens * head_frac); t = budget_tokens - h50 text = tokenizer.decode(ids[:h]) + "\n[...]\n" + tokenizer.decode(ids[len(ids) - t:])51 b = len(text.encode("utf-8"))52 return Message("A2", text, "Document excerpt (head and tail)", b, b, 0, (time.perf_counter() - t0) * 1000, 0.0, meta={"budget_tokens": budget_tokens, "source_tokens": len(ids)})5354# ---------------------------------------------------------------- A3 extractive selection (deterministic)55def _bm25_scores(sents: list[str], query_toks: list[str], k1=1.5, b=0.75) -> list[float]:56 docs = [_tok(s) for s in sents]57 N = len(docs); avgdl = sum(len(d) for d in docs) / max(1, N)58 df = {}59 for d in docs:60 for w in set(d): df[w] = df.get(w, 0) + 161 out = []62 for d in docs:63 sc = 0.0; L = len(d)64 tf = {}65 for w in d: tf[w] = tf.get(w, 0) + 166 for w in query_toks:67 if w in tf:68 idf = math.log(1 + (N - df[w] + 0.5) / (df[w] + 0.5))69 sc += idf * tf[w] * (k1 + 1) / (tf[w] + k1 * (1 - b + b * L / avgdl))70 out.append(sc)71 return out7273def a3_extractive(doc: str, budget_tokens: int, tokenizer, question: Optional[str]) -> Message:74 """Question-aware: BM25 of the question over sentences. Question-blind: centrality (sum of BM25 similarity to all other75 sentences, favouring sentences dense in numbers/identifiers), i.e. a generic 'keep concrete facts' heuristic."""76 t0 = time.perf_counter()77 sents = _sentences(doc)78 if question:79 scores = _bm25_scores(sents, _tok(question))80 else:81 toks = [_tok(s) for s in sents]82 allq = [w for t in toks for w in set(t)]83 base = _bm25_scores(sents, list(set(allq)))84 conc = [len(re.findall(r"\d", s)) + 2 * len(re.findall(r"\b[A-Z]{2}-\d{2}[A-Z]-\d{4}\b", s)) for s in sents]85 scores = [b_ * (1 + 0.1 * c) for b_, c in zip(base, conc)]86 order = sorted(range(len(sents)), key=lambda i: -scores[i])87 chosen, used = [], 088 for i in order:89 n = len(tokenizer.encode(sents[i], add_special_tokens=False))90 if used + n > budget_tokens:91 continue92 chosen.append(i); used += n93 if used >= budget_tokens * 0.97:94 break95 chosen.sort()96 text = " ".join(sents[i] for i in chosen)97 b = len(text.encode("utf-8"))98 return Message("A3", text, "Selected passages", b, b, 0, (time.perf_counter() - t0) * 1000, 0.0,99 meta={"budget_tokens": budget_tokens, "realized_tokens": used, "n_sentences": len(chosen), "question_aware": question is not None})100101# ---------------------------------------------------------------- A4 / A5 sender-generated102def _sender_generate(sender, system: str, user: str, max_tokens: int) -> tuple[str, dict, float]:103 ids = sender.chat_prompt_tokens([{"role": "system", "content": system}, {"role": "user", "content": user}])104 t0 = time.perf_counter()105 g = sender.generate_text(ids, max_tokens=max_tokens)106 ms = (time.perf_counter() - t0) * 1000107 usage = {"sender_input_tokens": len(ids), "sender_output_tokens": g["output_tokens"], "sender_prefill_ms": g["prefill_ms"], "sender_ttft_ms": g["ttft_ms"], "sender_decode_ms": g["decode_ms"]}108 return g["text"], usage, ms109110def a4_summary(doc: str, budget_tokens: int, sender, question: Optional[str]) -> Message:111 text, usage, ms = _sender_generate(sender, SUMMARY_SYSTEM, summary_user(doc, budget_tokens, question), max_tokens=budget_tokens)112 b = len(text.encode("utf-8"))113 return Message("A4", text, "Summary written by another agent", b, b, 0, ms, 0.0, usage, meta={"budget_tokens": budget_tokens, "question_aware": question is not None})114115def _parse_json_loose(s: str):116 s = s.strip()117 if s.startswith("```"):118 s = re.sub(r"^```(?:json)?", "", s).rstrip("`").strip()119 i, j = s.find("{"), s.rfind("}")120 if i >= 0 and j > i:121 s2 = s[i:j + 1]122 try:123 return json.loads(s2)124 except Exception:125 # try to repair truncated JSON: cut at last complete fact object126 k = s2.rfind("}")127 while k > 0:128 try:129 return json.loads(s2[:k + 1] + "]}")130 except Exception:131 k = s2.rfind("}", 0, k)132 return None133134def a5_structured(doc: str, budget_tokens: int, sender, question: Optional[str]) -> Message:135 """Typed facts with exact quotes resolved to source spans; unknowns kept; optional notes field. Quotes that do not136 resolve to a substring of the document are flagged (unresolved) and counted."""137 raw, usage, ms = _sender_generate(sender, EVIDENCE_SYSTEM, evidence_user(doc, budget_tokens, question), max_tokens=budget_tokens)138 obj = _parse_json_loose(raw)139 facts = []140 unresolved = 0141 if isinstance(obj, dict) and isinstance(obj.get("facts"), list):142 for f in obj["facts"]:143 if not isinstance(f, dict):144 continue145 qt = str(f.get("quote", ""))146 pos = doc.find(qt) if qt else -1147 if pos < 0:148 unresolved += 1149 facts.append({"entity": str(f.get("entity", "")), "attribute": str(f.get("attribute", "")), "value": str(f.get("value", "")),150 "span": [pos, pos + len(qt)] if pos >= 0 else None})151 notes = obj.get("notes"); unknowns = obj.get("unknowns")152 else:153 notes, unknowns = None, None154 msg_obj = {"facts": facts}155 if notes: msg_obj["notes"] = notes156 if unknowns: msg_obj["unknowns"] = unknowns157 text = json.dumps(msg_obj, ensure_ascii=False)158 b = len(text.encode("utf-8"))159 return Message("A5", text, "Structured evidence extracted by another agent (JSON; span = character offsets in the source)", b, b, 0, ms, 0.0, usage,160 meta={"budget_tokens": budget_tokens, "n_facts": len(facts), "unresolved_quotes": unresolved, "parse_ok": obj is not None, "question_aware": question is not None},161 artifacts={"raw": raw})162163# ---------------------------------------------------------------- A6 references + metered retrieval164def a6_references(doc: str, budget_tokens: int, sender, question: Optional[str], evidence_msg: Optional[Message] = None) -> Message:165 """Message = compact reference index (entity/attribute + span) WITHOUT values; the receiver must retrieve spans."""166 ev = evidence_msg or a5_structured(doc, budget_tokens, sender, question)167 facts = json.loads(ev.text).get("facts", [])168 refs = [{"ref": i, "entity": f["entity"], "attribute": f["attribute"], "span": f["span"]} for i, f in enumerate(facts) if f.get("span")]169 text = json.dumps({"references": refs}, ensure_ascii=False)170 b = len(text.encode("utf-8"))171 return Message("A6", text, "Reference index", b, b, 0, ev.encode_ms, 0.0, ev.sender_usage,172 meta={"n_refs": len(refs), "question_aware": question is not None}, artifacts={"refs": refs, "doc": doc})173174def a6_retrieve(msg: Message, question: str, max_fetch_bytes: int, max_lookups: int) -> tuple[str, dict]:175 """Receiver-side retrieval: rank references by BM25 of the question against entity/attribute labels; fetch spans176 (plus a small context window) from the source under the byte and lookup budget. Every fetch is metered."""177 refs, doc = msg.artifacts["refs"], msg.artifacts["doc"]178 labels = [f"{r['entity']} {r['attribute']}" for r in refs]179 scores = _bm25_scores(labels, _tok(question)) if refs else []180 order = sorted(range(len(refs)), key=lambda i: -scores[i])181 fetched, nbytes, lookups, pieces = [], 0, 0, []182 for i in order:183 if lookups >= max_lookups:184 break185 s, e = refs[i]["span"]186 s2, e2 = max(0, s - 40), min(len(doc), e + 160)187 snippet = doc[s2:e2]188 nb = len(snippet.encode("utf-8"))189 if nbytes + nb > max_fetch_bytes:190 continue191 lookups += 1; nbytes += nb192 pieces.append(f"[ref {refs[i]['ref']}: {refs[i]['entity']} / {refs[i]['attribute']}] ...{snippet}...")193 material = msg.text + "\n\nRetrieved source passages:\n" + "\n".join(pieces)194 return material, {"lookups": lookups, "fetched_bytes": nbytes, "retrieved_refs": lookups}195196# ---------------------------------------------------------------- A7 lossless transport codec197def a7_codec(doc: str, codec: str = "zstd") -> Message:198 raw = doc.encode("utf-8")199 comp, enc_ms = codec_compress(raw, codec)200 back, dec_ms = codec_decompress(comp, codec)201 assert back == raw202 return Message("A7", doc, "Document", len(raw), len(comp), 0, enc_ms, dec_ms, meta={"codec": codec, "ratio": len(raw) / max(1, len(comp))})203204# ---------------------------------------------------------------- A8 LLMLingua-2 (learned prompt compression)205_LL = None206def a8_llmlingua(doc: str, budget_tokens: int, tokenizer, question: Optional[str], device: str = "mps") -> Message:207 global _LL208 t0 = time.perf_counter()209 if _LL is None:210 from llmlingua import PromptCompressor211 _LL = PromptCompressor(model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank", use_llmlingua2=True, device_map=device)212 src_tokens = len(tokenizer.encode(doc, add_special_tokens=False))213 rate = min(1.0, max(0.02, budget_tokens / max(1, src_tokens)))214 res = _LL.compress_prompt(doc, rate=rate, force_tokens=["\n", ".", ",", ":", "-"], drop_consecutive=True)215 text = res["compressed_prompt"]216 # LLMLingua-2 is question-agnostic by design (token classification); question is recorded but not used217 b = len(text.encode("utf-8"))218 return Message("A8", text, "Compressed document", b, b, 0, (time.perf_counter() - t0) * 1000, 0.0,219 meta={"budget_tokens": budget_tokens, "rate": rate, "origin_tokens": res.get("origin_tokens"), "compressed_tokens": res.get("compressed_tokens"), "question_aware": False,220 "compressor": "microsoft/llmlingua-2-xlm-roberta-large-meetingbank", "device": device})221