code_release/code/memebench/memory_runner.py

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