meme-bench / code_release /code /memebench /memory_runner.py
clinical-agent-memory-bench's picture
Upload 99 files
750b88d verified
Raw History Blame Contribute Delete
24.6 kB
"""Streaming evaluation: per-visit memory construction and questions at their checkpoints.
Visits are fed to a method one at a time in chronological order. Each question is answered
when the stream reaches its checkpoint visit, so the method has seen exactly the visits up to
and including that checkpoint. Answers are graded afterwards by the judge in a thread pool.
A context overflow is recorded as '[CONTEXT_EXCEEDED]' and any other API failure as
'[err ...]'; both are graded as failed.
"""
from __future__ import annotations
import json
import os
import re
import shutil
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from .dataset import Patient, Question, render_visit
from .llm_client import USAGE, ContextOverflow, make_answerer, make_judge
from .methods import INTERNAL_LLM_METHODS, get_method
from .prompts import MEMORIZE_TEMPLATE, ANSWER_SYS, ANSWER_USER, qtext
from . import judges
from .report import MEMORY_LOG_MAX_CHARS, build_report, evidence_diagnostic, print_report
ROWS_MISMATCH = "stored rows do not match the questions file; start a fresh run"
def parse_qids(spec: str, n: int) -> list[int]:
"""Parse a question selector: '0-5,7' -> [0, 1, 2, 3, 4, 5, 7]; '' selects all."""
sel: set[int] = set()
for part in filter(None, (p.strip() for p in spec.split(","))):
if "-" in part:
a, b = part.split("-", 1)
sel |= set(range(int(a), int(b) + 1))
else:
sel.add(int(part))
if sel and not sel <= set(range(n)):
raise SystemExit(f"--qids out of range 0..{n - 1}: {sorted(sel - set(range(n)))}")
return sorted(sel) if sel else list(range(n))
def _memorize_text(agent, session: dict) -> str:
"""The rendered visit, wrapped in the memorize template when the method builds memory with an LLM."""
text = render_visit(session)
if agent.MEMORIZE_WITH_LLM:
return MEMORIZE_TEMPLATE.format(context=text)
return text
RUN_DIR_RE = re.compile(r"^\d{8}_\d{6}(-\d+)?$")
def _latest_run(run_root: Path, complete_only: bool = False) -> Path | None:
"""Return the newest timestamped run under `run_root`, or None; complete_only skips runs without report.json."""
runs = sorted((d for d in run_root.glob("*") if d.is_dir() and RUN_DIR_RE.match(d.name)
and (not complete_only or (d / "report.json").exists())),
key=lambda d: d.name)
return runs[-1] if runs else None
def _mark_latest(run_root: Path, out_dir: Path) -> None:
"""Point <run_root>/latest at this run (atomic symlink replace)."""
link = run_root / "latest"
tmp = run_root / f".latest.{os.getpid()}"
if tmp.is_symlink() or tmp.exists():
tmp.unlink()
tmp.symlink_to(out_dir.name)
os.replace(tmp, link)
def _write_rows(path: Path, rows: list[dict]) -> None:
"""Write rows as JSONL with an atomic replace."""
tmp = path.with_suffix(path.suffix + ".tmp")
with open(tmp, "w", encoding="utf-8") as fh:
for r in rows:
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
os.replace(tmp, path)
def _answer_one(answerer, agent, q: Question) -> dict:
"""Build the MEMORY block for one question, call the answering model, and return the ungraded row."""
USAGE.set_phase("query")
ctx = agent.build_memory_context(q.question, q.ckpt_date,
display_question=qtext(q.raw))
api_error = ctx.error
if ctx.answer is not None:
ans = str(ctx.answer).strip()
if not ans:
ans, api_error = "[err no_answer]", True
elif ctx.error:
ans = "[err agent_error]"
else:
prompt = ANSWER_USER.format(ctx=ctx.mem, as_of=q.ckpt_date, qtext=qtext(q.raw))
try:
ans = str(answerer.complete_json(ANSWER_SYS, prompt).get("answer") or "").strip()
except ContextOverflow:
ans, api_error = "[CONTEXT_EXCEEDED]", True
except Exception as e:
ans, api_error = f"[err {str(e)[:60]}]", True
if not ans and not api_error:
ans, api_error = "[err no_answer]", True
error_cause = None
if api_error:
if ctx.error:
error_cause = "agent_error"
elif ans == "[CONTEXT_EXCEEDED]":
error_cause = "context_overflow"
elif ans == "[err no_answer]":
error_cause = ("no_answer_from_model"
if getattr(agent, "ANSWERS_INTERNALLY", False) or ctx.n_tok
else "empty_memory")
else:
error_cause = "api_error"
return {
"qid": q.qid, "qid_str": q.raw.get("qid"), "api_error": api_error,
"error_cause": error_cause,
"axis": q.axis, "answer_type": q.answer_type, "grounding": q.grounding,
"ckpt_visit": q.ckpt_visit, "ckpt_date": q.ckpt_date,
"n_tok": ctx.n_tok, "context_truncated": ctx.context_truncated,
**evidence_diagnostic(
q,
ctx.retrieved_visits if (ctx.retrieved_visits or not getattr(agent, "ANSWERS_INTERNALLY", False)
or ctx.retrieved_count) else None),
"retrieved_count": ctx.retrieved_count,
"retrieved_visits": ctx.retrieved_visits,
"memory_preview": ctx.mem[:MEMORY_LOG_MAX_CHARS],
"agent_extra": ctx.extra,
"answer": ans,
}
def _grade_rows(rows: list[dict], questions: dict[int, Question], judge_client, workers: int) -> None:
"""Fill verdict, pass and judge_detail of each row in place, grading rows in parallel."""
USAGE.set_phase("judge")
def grade(row):
if row["api_error"]:
row.update(verdict="error", **{"pass": False}, judge_detail={})
return
ok, verdict, det = judges.score(judge_client, questions[row["qid"]], row["answer"])
row.update(verdict=str(verdict), **{"pass": bool(ok)}, judge_detail=det)
with ThreadPoolExecutor(max_workers=max(1, workers)) as ex:
list(ex.map(grade, rows))
def _mismatched_rows(rows: list[dict], qmap: dict[int, Question]) -> list[dict]:
"""Return the stored rows that do not belong to the current questions: an index that is not in the
questions file, or another question id at the same index."""
bad = []
for r in rows:
q = qmap.get(r.get("qid"))
if (q is None
or (r.get("qid_str") and q.raw.get("qid") not in (None, r["qid_str"]))):
bad.append(r)
return bad
def rejudge_patient(method_name: str, patient: Patient, out_root: Path,
answer_model: str | None = None, judge_model: str | None = None,
judge_workers: int = 8) -> dict:
"""Re-grade every stored answer of the newest completed run with the current judge."""
t_start = time.time()
run_root = Path(out_root) / f"{method_name}_{answer_model}" / patient.name
out_dir = _latest_run(run_root, complete_only=True)
if out_dir is None:
raise SystemExit(f"no completed run to re-judge under {run_root}")
answers_path = out_dir / "answers.jsonl"
if not answers_path.exists():
raise SystemExit(f"no answers to re-judge at {answers_path}")
USAGE.bind(out_dir / "usage.jsonl")
_, questions = patient.load()
qmap = {q.qid: q for q in questions}
rows = [json.loads(l) for l in open(answers_path, encoding="utf-8") if l.strip()]
if _mismatched_rows(rows, qmap) or {r["qid"] for r in rows} - set(qmap):
raise SystemExit(ROWS_MISMATCH)
before = {r["qid"]: r["pass"] for r in rows}
judge_client = make_judge(model=judge_model)
_grade_rows(rows, qmap, judge_client, judge_workers)
flips = [(r["qid"], before[r["qid"]], r["pass"]) for r in rows if before[r["qid"]] != r["pass"]]
stamp = time.strftime("%Y%m%d_%H%M%S", time.gmtime())
shutil.copy2(answers_path, out_dir / f"answers.before_rejudge_{stamp}.jsonl")
if (out_dir / "report.json").exists():
shutil.copy2(out_dir / "report.json", out_dir / f"report.before_rejudge_{stamp}.json")
_write_rows(answers_path, rows)
rep = json.load(open(out_dir / "report.json", encoding="utf-8"))
original_usage = rep.get("llm_usage", {})
rep.update(build_report(rows, {k: v for k, v in rep.items() if k not in (
"total", "pass", "accuracy", "by_axis", "by_answer_type", "by_grounding",
"evidence_truncated_count", "evidence_recall", "context_truncated_count",
"errors", "judge_failures", "llm_usage")}, original_usage))
rep["judge_model"] = getattr(judge_client, "model", judge_model)
rep["rejudged"] = {"regraded": len(rows), "flipped": len(flips),
"judge_model": rep["judge_model"], "usage": USAGE.totals,
"wall_time_s": round(time.time() - t_start, 1)}
(out_dir / "report.json").write_text(json.dumps(rep, ensure_ascii=False, indent=2))
print(f" {method_name} x {patient.name}: re-graded {len(rows)}, "
f"{len(flips)} verdict(s) flipped {[f'q{q}:{a}->{b}' for q, a, b in flips]} "
f"-> {rep['pass']}/{rep['total']} ({rep['accuracy']:.2f})")
return rep
def _endpoint_key(name: str) -> bool:
"""Parameters that name an endpoint or a key; they are not written to run_settings.json or report.json."""
return "base_url" in name or "api_key" in name
def _run_settings(method_name, patient, method_params, answer_model, judge_model, use_internal, answer_max_tokens,
answer_temperature) -> dict:
"""The settings a resumed run must match: method, patient, models, method parameters and answering settings."""
params = {k: v for k, v in method_params.items() if k != "working_dir" and not _endpoint_key(k)}
return {"method": method_name, "patient": patient.name,
"answer_model": answer_model, "judge_model": judge_model,
"internal_model": method_params.get("model") if use_internal else None,
"method_params": json.loads(json.dumps(params, default=str)),
"answer_max_tokens": answer_max_tokens, "answer_temperature": answer_temperature}
def run_patient(method_name: str, patient: Patient, method_params: dict, out_root: Path, *,
answer_max_tokens: int,
answer_model: str | None = None, judge_model: str | None = None,
resume: bool = False,
qids: str = "",
judge_workers: int = 8,
answer_temperature: float | None = None,
internal_model: str | None = None,
internal_base_url: str | None = None,
answer_base_url: str | None = None) -> dict:
"""Stream one patient's visits through one memory method; answer + grade its questions."""
t_start = time.time()
run_root = Path(out_root) / f"{method_name}_{answer_model}" / patient.name
if qids and not resume:
raise SystemExit("--qids without --resume would start a new partial run; "
"add --resume to merge into the newest existing run")
out_dir = _latest_run(run_root) if resume else None
if method_name == "memos" and out_dir is not None:
raise SystemExit("memos cannot resume an existing run; start a fresh run.")
if out_dir is None:
stamp = time.strftime("%Y%m%d_%H%M%S", time.gmtime())
run_root.mkdir(parents=True, exist_ok=True)
out_dir, n = run_root / stamp, 1
while True:
try:
out_dir.mkdir(exist_ok=False)
break
except FileExistsError:
n += 1
out_dir = run_root / f"{stamp}-{n}"
else:
out_dir.mkdir(parents=True, exist_ok=True)
USAGE.bind(out_dir / "usage.jsonl")
sessions, questions = patient.load()
qmap = {q.qid: q for q in questions}
torun = set(parse_qids(qids, len(questions)))
answers_path = out_dir / "answers.jsonl"
prev_rows: list[dict] = []
carried: list[dict] = []
if answers_path.exists() and resume:
prev_rows = [json.loads(l) for l in open(answers_path, encoding="utf-8") if l.strip()]
if resume:
done = {r["qid"] for r in prev_rows if r.get("verdict") not in (None, "error", "?")}
if qids:
print(f"[resume] answering the {len(torun)} selected qids again "
f"({len(torun & done)} already graded); their stored rows are replaced")
else:
torun -= done
print(f"[resume] {len(done)} qids already graded, {len(torun)} to run")
if resume:
partial = out_dir / "answers_partial.jsonl"
partial_rows = []
if partial.exists():
plines = [l for l in open(partial, encoding="utf-8") if l.strip()]
for i, line in enumerate(plines):
try:
partial_rows.append(json.loads(line))
except json.JSONDecodeError:
if i == len(plines) - 1:
print(f"[resume] dropping torn final line of {partial.name}")
continue
raise SystemExit(f"cannot resume: line {i + 1} of {partial} is corrupt")
if _mismatched_rows(prev_rows + partial_rows, qmap):
raise SystemExit(ROWS_MISMATCH)
if partial_rows:
seen = {r["qid"] for r in prev_rows}
for r in partial_rows:
if r["qid"] in torun and r["qid"] not in seen:
carried.append(r)
seen.add(r["qid"])
if carried:
torun -= {r["qid"] for r in carried}
print(f"[resume] rescued {len(carried)} answered-but-ungraded qids from "
f"answers_partial.jsonl; {len(torun)} left to answer")
by_ckpt: dict[int, list[Question]] = defaultdict(list)
for q in questions:
if q.qid in torun:
by_ckpt[q.ckpt_visit].append(q)
consolidate_at = {q.ckpt_visit for q in questions if resume or q.qid in torun}
method_params = dict(method_params)
if "answer_max_tokens" in method_params:
raise SystemExit("answer_max_tokens is set once, at the top of the methods config")
answer_max_tokens = int(answer_max_tokens)
if method_name == "lightmem":
method_params.setdefault("answer_max_tokens", answer_max_tokens)
use_internal = bool(internal_model) and method_name in INTERNAL_LLM_METHODS
if use_internal:
method_params.setdefault("model", internal_model)
method_params.setdefault("internal_base_url", internal_base_url)
elif method_name in INTERNAL_LLM_METHODS and answer_base_url:
method_params.setdefault("internal_base_url", answer_base_url)
method_params.setdefault("model", answer_model)
if method_params.pop("reuse_store", False):
raise SystemExit("reuse_store is not supported: every run replays the visits from the first one, so a "
"reused store would hold each visit twice and early questions would see later visits")
method_params.setdefault("working_dir", str(out_dir / "store"))
settings = _run_settings(method_name, patient, method_params, answer_model, judge_model, use_internal,
answer_max_tokens, answer_temperature)
settings_path = out_dir / "run_settings.json"
if resume:
if not settings_path.exists():
raise SystemExit(f"cannot resume {out_dir}: it has no run_settings.json; start a fresh run")
before = json.loads(settings_path.read_text(encoding="utf-8"))
drift = sorted(k for k in set(before) | set(settings) if before.get(k) != settings.get(k))
if "method_params" in drift:
was, now = before.get("method_params") or {}, settings["method_params"]
drift = [k for k in drift if k != "method_params"] + sorted(
f"method_params.{k}" for k in set(was) | set(now) if was.get(k) != now.get(k))
if drift:
raise SystemExit(f"cannot resume {out_dir}: the settings differ in {drift}; use the same settings "
f"or start a fresh run")
else:
settings_path.write_text(json.dumps(settings, ensure_ascii=False, indent=2), encoding="utf-8")
nothing_to_run = not by_ckpt
agent = None
if not nothing_to_run:
agent = get_method(method_name, **method_params)
agent.set_context_id(patient.patient_id)
answerer = make_answerer(model=answer_model, max_tokens=answer_max_tokens,
temperature=answer_temperature, base_url=answer_base_url)
judge_client = make_judge(model=judge_model)
new_rows: list[dict] = []
if nothing_to_run:
print("[resume] nothing to run; recomputing the report from the stored answers")
sessions = []
partial_path = out_dir / "answers_partial.jsonl"
build_path = out_dir / "memory_build.jsonl"
if by_ckpt and build_path.exists():
os.replace(build_path,
out_dir / f"memory_build.pre_{time.strftime('%Y%m%d_%H%M%S', time.gmtime())}.jsonl")
build_log = open(build_path, "w" if by_ckpt else "a", encoding="utf-8")
partial_log = open(partial_path, "a", encoding="utf-8")
def _snapshot(tag: str) -> None:
snap_dir = out_dir / "memory_snapshots" / tag
try:
if snap_dir.exists():
shutil.rmtree(snap_dir)
snap_dir.mkdir(parents=True)
stats = agent.dump_memory(snap_dir)
if stats is None:
snap_dir.rmdir()
else:
(snap_dir / "stats.json").write_text(json.dumps(stats, default=str))
except Exception as e:
if getattr(agent, "FAIL_CLOSED_ON_MEMORY_ERROR", False):
raise RuntimeError(f"{method_name} snapshot {tag} failed; see the underlying error") from e
print(f" [snapshot] {tag} failed (run continues): {str(e)[:120]}", flush=True)
trace_log = (out_dir / "memory_trace.jsonl").open("a", encoding="utf-8")
def _trace_sink(rec: dict) -> None:
rec.setdefault("visit_number", _cur_visit[0])
trace_log.write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
trace_log.flush()
_cur_visit = [None]
if agent is not None:
try:
agent.set_trace_sink(_trace_sink)
except Exception:
pass
_memfail = {"n": 0, "streak": 0, "consolidate": 0}
try:
for s in sessions:
_cur_visit[0] = s["visit_number"]
USAGE.set_phase("memorize")
res = agent.memorize_session(s, _memorize_text(agent, s))
build_log.write(json.dumps(res.to_dict(), ensure_ascii=False) + "\n")
build_log.flush()
if not getattr(res, "success", True):
_memfail["n"] += 1
_memfail["streak"] += 1
if getattr(agent, "FAIL_CLOSED_ON_MEMORY_ERROR", False):
raise RuntimeError(
f"{method_name} ingestion failed at visit {s['visit_number']}; "
"memory is incomplete, see memory_build.jsonl")
if _memfail["streak"] in (5, 20) or _memfail["streak"] % 50 == 0:
detail = (res.extra or {}) if hasattr(res, "extra") else {}
print(" [memorize] failed on %d visits in a row (%d in total); last: %s"
% (_memfail["streak"], _memfail["n"],
str(detail.get("error") or detail or res.action)[:160]),
flush=True)
else:
_memfail["streak"] = 0
if s["visit_number"] in consolidate_at:
USAGE.set_phase("memorize")
try:
_cons = agent.consolidate()
except Exception as e:
_cons = {"error": str(e)[:200]}
if getattr(agent, "FAIL_CLOSED_ON_MEMORY_ERROR", False):
build_log.write(json.dumps({"method": method_name, "action": "consolidate",
"visit_number": s["visit_number"], "success": False, **_cons}) + "\n")
build_log.flush()
raise RuntimeError(
f"{method_name} consolidation failed at visit {s['visit_number']}") from e
print(f" [consolidate] v{s['visit_number']} failed (run continues): "
f"{str(e)[:120]}", flush=True)
if _cons and "error" in _cons:
_memfail["consolidate"] += 1
if _cons:
build_log.write(json.dumps(
{"method": agent.METHOD_NAME if hasattr(agent, "METHOD_NAME") else "",
"action": "consolidate", "visit_number": s["visit_number"],
"success": "error" not in _cons, **_cons},
ensure_ascii=False) + "\n")
build_log.flush()
print(f" [consolidate] v{s['visit_number']} {_cons}", flush=True)
if by_ckpt.get(s["visit_number"]):
_snapshot(f"ckpt_v{s['visit_number']:04d}")
for q in by_ckpt.get(s["visit_number"], []):
row = _answer_one(answerer, agent, q)
new_rows.append(row)
partial_log.write(json.dumps(row, ensure_ascii=False) + "\n")
partial_log.flush()
print(f" [q] qid={q.qid} {q.axis} @v{q.ckpt_visit} mem={row['n_tok']}tok "
f"ans={row['answer'][:60]!r}", flush=True)
if sessions:
_snapshot("final")
finally:
build_log.close()
partial_log.close()
trace_log.close()
close_agent = getattr(agent, "close", None)
if callable(close_agent):
close_agent()
new_rows.extend(carried)
_grade_rows(new_rows, qmap, judge_client, judge_workers)
fresh = {r["qid"] for r in new_rows}
rows = sorted([r for r in prev_rows if r["qid"] not in fresh] + new_rows, key=lambda r: r["qid"])
_write_rows(answers_path, rows)
if partial_path.exists():
partial_path.unlink()
prior_rep = {}
if (out_dir / "report.json").exists():
prior_rep = json.load(open(out_dir / "report.json", encoding="utf-8"))
prior = prior_rep if nothing_to_run else {}
rep = build_report(rows, {
"method": method_name, "patient": patient.name, "patient_id": patient.patient_id,
"answer_model": answer_model,
"judge_model": getattr(judge_client, "model", judge_model),
"method_params": {**{k: v for k, v in method_params.items() if not _endpoint_key(k)},
"working_dir": os.path.relpath(method_params["working_dir"], out_dir)},
"questions_file": Path(patient.questions_file).name if patient.questions_file else None,
"run_started_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(t_start)),
"internal_model": method_params.get("model") if use_internal else None,
"memorize_wrapper": agent.MEMORIZE_WITH_LLM if agent is not None else None,
"answers_internally": agent.ANSWERS_INTERNALLY if agent is not None else None,
"answer_temperature": answer_temperature, "answer_max_tokens": answer_max_tokens,
"n_sessions": len(sessions),
"memorize_failures": _memfail["n"], "consolidate_failures": _memfail["consolidate"],
"wall_time_s": round(time.time() - t_start, 1),
**{k: prior[k] for k in ("method_params", "answer_model", "answer_temperature",
"answer_max_tokens", "n_sessions", "wall_time_s",
"memorize_failures", "consolidate_failures",
"run_started_at",
"internal_model",
"questions_file", "judge_model",
"memorize_wrapper", "answers_internally")
if k in prior},
}, prior.get("llm_usage") or USAGE.totals)
if not nothing_to_run and resume and prev_rows:
rep["resumed_over"] = {"prior_run_started_at": prior_rep.get("run_started_at"),
"n_prior_rows": len(prev_rows)}
(out_dir / "report.json").write_text(json.dumps(rep, ensure_ascii=False, indent=2))
_mark_latest(run_root, out_dir)
print_report(rep)
print(f"wrote {answers_path}")
return rep