"""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 /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