longmemeval.py

"""longmemeval -- external-benchmark swing: does the unified controller help on a PUBLISHED
long-term-memory benchmark, on the slices our thesis targets?

LongMemEval (Wu et al., ICLR 2025). We run the KNOWLEDGE-UPDATE and ABSTENTION slices -- exactly the
abilities our write policy claims: recall the LATEST value of a fact that changed (knowledge-update),
and DECLINE when the answer was never stated (abstention). Oracle setting (evidence sessions only),
which isolates the WRITE/forget policy from retrieval-at-scale.

Ablation (same embedder + same top-k retrieval + same Qwen answerer; ONLY the write policy differs):
  store_all         -- keep every turn
  cosine_forget     -- dedup + cosine-only forgetting
  flashbulb         -- surprise admission + admission-via-contradiction + NLI-contradiction forgetting
                       + retention-tagged readout (its abstention signal)

Answerer = qwen-plus (D-011). Judge = qwen-max (D-011), disclosed (LongMemEval's own harness uses GPT-4o;
we substitute a Qwen judge to keep everything on Qwen and report that). ZERO LLM tokens on every write.
Checkpointed + resumable: re-running skips finished question_ids. Config via env: LME_SLICE
(knowledge-update,abstention -> default), LME_MAXQ (cap questions, default all), LME_K (top-k, default 8).
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
try:
    from dotenv import load_dotenv
    load_dotenv(Path(__file__).resolve().parents[1] / ".env")
except Exception:
    pass
from memory.sidecar import MemorySidecar, _get_embedder  # noqa: E402
from benchmarks.drift_suite import _Vec, _CosineSupersede  # noqa: E402
from agent import ANSWER_SYSTEM  # noqa: E402

DATA = Path(__file__).with_name("data") / "longmemeval_oracle.json"
OUT = Path(__file__).with_name("longmemeval_results.json")
EMB = "BAAI/bge-small-en-v1.5"


# ------------------------------------------------------------------ policies (uniform interface)
class _StoreAll:
    def __init__(self, emb): self.p = _Vec(emb)
    def ingest(self, text): self.p.write(text)
    def context(self, q, k): return "\n".join(f"- {t}" for t in self.p.read(q, k=k))
    def size(self): return self.p.size()

class _CosineForget:
    def __init__(self, emb): self.p = _CosineSupersede(emb)
    def ingest(self, text): self.p.write(text)
    def context(self, q, k): return "\n".join(f"- {t}" for t in self.p.read(q, k=k))
    def size(self): return self.p.size()

class _Flashbulb:
    def __init__(self, emb): self.sc = MemorySidecar("lme", tau=0.40, use_nli=True, seed=0, embedder=emb)
    def ingest(self, text): self.sc.write(text)
    def context(self, q, k): return self.sc.read(q, k=k)[1]  # retention-tagged, forgotten dropped
    def size(self): return len(self.sc._active_records())

POLICIES = {"store_all": _StoreAll, "cosine_forget": _CosineForget, "flashbulb": _Flashbulb}


# ------------------------------------------------------------------ Qwen answerer + judge
def _client(model):
    from llm_client import LLMClient
    return LLMClient(model=model)

def answer(context: str, question: str, ans_model) -> str:
    system = ANSWER_SYSTEM + (f"\n\nMEMORY:\n{context}" if context.strip() else "\n\nMEMORY:\n(empty)")
    return ans_model.chat([{"role": "system", "content": system},
                           {"role": "user", "content": question}]).text.strip()

def judge(question, gold, sys_answer, is_abstention, judge_model) -> bool:
    if is_abstention:
        prompt = (f"A user asked a memory assistant a question whose answer was NEVER provided, so the "
                  f"CORRECT behavior is to decline / say it doesn't know / that the info isn't available.\n"
                  f"Question: {question}\nAssistant answer: {sys_answer}\n\n"
                  f"Did the assistant correctly DECLINE / abstain (rather than fabricating a specific "
                  f"answer)? Reply with exactly YES or NO.")
    else:
        prompt = (f"You are grading a memory assistant.\nQuestion: {question}\nReference answer: {gold}\n"
                  f"Assistant answer: {sys_answer}\n\nDoes the assistant answer contain the key information "
                  f"in the reference answer (correct)? Ignore phrasing/extra detail. Reply exactly YES or NO.")
    r = judge_model.chat([{"role": "user", "content": prompt}]).text.strip().upper()
    return r.startswith("Y")


# ------------------------------------------------------------------ run
def _turns(entry):
    """All haystack turns, oldest->newest, as records (content only)."""
    out = []
    for sess in entry["haystack_sessions"]:
        for turn in sess:
            c = (turn.get("content") or "").strip()
            if c:
                out.append(c)
    return out

def main():
    slice_types = set(os.environ.get("LME_SLICE", "knowledge-update,abstention").split(","))
    maxq = int(os.environ.get("LME_MAXQ", "0")) or None
    k = int(os.environ.get("LME_K", "8"))
    policy_names = os.environ.get("LME_POLICIES", "store_all,cosine_forget,flashbulb").split(",")

    data = json.loads(DATA.read_text(encoding="utf-8"))
    def _slice(d):  # abstention is a _abs suffix, layered over the type
        is_abs = d["question_id"].endswith("_abs")
        return ("abstention" in slice_types and is_abs) or (d["question_type"] in slice_types)
    items = [d for d in data if _slice(d)]
    if maxq:
        items = items[:maxq]

    emb = _get_embedder(EMB, "cpu")
    ans_model, judge_model = _client(os.environ.get("LLM_MODEL", "qwen-plus")), _client(
        os.environ.get("LLM_JUDGE_MODEL", "qwen-max"))

    results = json.loads(OUT.read_text(encoding="utf-8")) if OUT.exists() else {"per_question": {}, "config": {}}
    done = results["per_question"]
    t0 = time.time()
    todo = [d for d in items if d["question_id"] not in done]
    print(f"LongMemEval slice={sorted(slice_types)} | {len(items)} questions ({len(todo)} to run) | "
          f"policies={policy_names} | k={k}", flush=True)

    for i, entry in enumerate(todo):
        qid = entry["question_id"]
        is_abs = qid.endswith("_abs")
        turns = _turns(entry)
        rec = {"type": entry["question_type"], "abstention": is_abs, "n_turns": len(turns), "policies": {}}
        for pname in policy_names:
            pol = POLICIES[pname](emb)
            for t in turns:
                pol.ingest(t)
            ctx = pol.context(entry["question"], k)
            ans = answer(ctx, entry["question"], ans_model)
            ok = judge(entry["question"], entry.get("answer", ""), ans, is_abs, judge_model)
            rec["policies"][pname] = {"correct": ok, "answer": ans[:300], "mem_size": pol.size()}
        done[qid] = rec
        # checkpoint + progress every question
        results["per_question"] = done
        results["config"] = {"slice": sorted(slice_types), "k": k, "policies": policy_names,
                             "answerer": os.environ.get("LLM_MODEL", "qwen-plus"),
                             "judge": os.environ.get("LLM_JUDGE_MODEL", "qwen-max"), "n_total": len(items)}
        results["summary"] = _summ(done, policy_names)
        OUT.write_text(json.dumps(results, indent=2), encoding="utf-8")
        n = i + 1
        el = time.time() - t0
        eta = (len(todo) - n) / (n / el) if el else 0
        acc = {p: results["summary"][p]["accuracy"] for p in policy_names}
        print(f"[{100*n/len(todo):5.1f}%] {n}/{len(todo)} {qid[:22]:22s} el {el/60:.1f}m eta {eta/60:.1f}m | "
              f"acc {acc}", flush=True)

    print("\nFINAL:", json.dumps(results["summary"], indent=2))
    print(f"wrote {OUT}")

def _summ(done, policy_names):
    out = {}
    for p in policy_names:
        rows = [r for r in done.values() if p in r["policies"]]
        n = len(rows) or 1
        correct = sum(r["policies"][p]["correct"] for r in rows)
        by_type = {}
        for r in rows:
            key = "abstention" if r["abstention"] else r["type"]
            by_type.setdefault(key, [0, 0])
            by_type[key][0] += int(r["policies"][p]["correct"]); by_type[key][1] += 1
        out[p] = {"accuracy": round(correct / n, 4), "correct": correct, "n": len(rows),
                  "by_type": {t: round(c / tot, 3) for t, (c, tot) in by_type.items()}}
    return out

if __name__ == "__main__":
    main()