cortexm 0.3.0__py3-none-any.whl

cortexm/bench/ood.py

@@ -0,0 +1,443 @@
+"""Out-of-distribution (OOD) benchmark — breaking the circularity.
+
+The in-distribution (ID) benchmark generates conversations from the SAME
+template families the μ=0 extractor's patterns were authored against, so ID
+scores are an upper bound for template-shaped text, not a capability claim.
+This module measures the honest generalization gap:
+
+  1. EXTRACTION RECALL per style — ground-truth facts from persona
+     registries are re-rendered by an independent LLM in styles the pattern
+     author never saw (paraphrase / negation / indirect / informal /
+     non_english / code_switch). The μ=0 extractor runs on the renderings;
+     recall is matched against the ground-truth match keys.
+  2. END-TO-END RETRIEVAL — the OOD corpus (renderings + the same
+     distractor machinery) flows through the full memory fabric and is
+     probed by the SAME probe builder + judges as the ID benchmark, so
+     ID-vs-OOD deltas are apples-to-apples.
+  3. LLM-JUDGE CROSS-CHECK — probe/context pairs are exported in the
+     canonical BEAM judge format for independent LLM grading.
+
+Renderer omissions (facts the LLM failed to convey) are tracked separately
+from extraction failures so neither layer can silently blame the other.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+
+from cortexm.bench.abilities import ABILITIES, build_probes, judge
+from cortexm.bench.generator import (Corpus, distractor_paragraph,
+                                       smalltalk_message)
+from cortexm.util import month_name, normalize, token_estimate
+
+T0 = datetime(2026, 3, 1, tzinfo=timezone.utc)
+
+# Match-tier vocabulary: how hard each fact is to extract from re-phrased text.
+CATEGORY_TIERS = {
+    "core": "explicit entity facts (name, employer, city, family, prefs)",
+    "relational": "multi-hop chains (manager->team->tech)",
+    "temporal": "dated events and interval changes",
+}
+
+
+def _d(y, m, day=1):
+    return f"{y:04d}-{m:02d}-{day:02d}"
+
+
+def _mn(iso_date: str) -> str:
+    return month_name(int(iso_date[5:7]))
+
+
+# ---------------------------------------------------------------- manifest
+def export_manifest(personas: list, t0: datetime = T0) -> dict:
+    """Persona ground truth as LLM-renderable fact manifests.
+
+    Mirrors the generator's session structure (parts 0-7) so the SAME probes
+    remain answerable; facts are stated as semantic text, never as the
+    generator's template phrasings.
+    """
+    out = []
+    for p in personas:
+        e0, e1 = p.employers[0], p.employers[1]
+        e_last = p.employers[-1]
+        c0, c1 = p.cities
+        sister_full = p.family[0][0]
+        sessions: list[dict] = []
+
+        def S(part: int, facts: list[dict]) -> None:
+            date = t0 + timedelta(days=part * 21)
+            sessions.append({"session": part, "date": date.date().isoformat(),
+                             "facts": [
+                                 {"id": f"s{part}f{i}", "text": t,
+                                  "type": ty, "match": m}
+                                 for i, (ty, t, m) in enumerate(facts)]})
+
+        # ---- part 0: introduction
+        f0 = [
+            ("name", f"The user's full name is {p.full_name}.",
+             [p.full_name, p.first]),
+            ("employment", f"The user works at {e0[0]} as a {p.roles[0][0]}.",
+             [e0[0]]),
+            ("city", f"The user lives in {c0[0]}.", [c0[0]]),
+            ("birthday",
+             f"The user's birthday is {month_name(p.birthday[0])} "
+             f"{p.birthday[1]}.",
+             [f"{month_name(p.birthday[0]).lower()} {p.birthday[1]}"]),
+            ("family",
+             f"The user has a sister named {sister_full}.",
+             [sister_full.split()[0]]),
+        ]
+        if p.nickname:
+            f0.append(("alias", f"The user goes by the nickname "
+                       f"\"{p.nickname}\".", [p.nickname]))
+        _instr_keys = (["french"] if "french" in p.instruction[1].lower()
+                       else ["short", "concise", "brief"])
+        f0.append(("instruction",
+                   f"The user gave the assistant a standing instruction: "
+                   f"\"{p.instruction[0]}\"", _instr_keys))
+        f0.append(("employment_since",
+                   f"The user has been at {e0[0]} since {_mn(e0[1])} {e0[1][:4]}.",
+                   [e0[0]]))
+        S(0, f0)
+
+        # ---- part 1: preferences + skills
+        f1 = []
+        for i in range(0, len(p.prefs), 2):
+            cat, v_old = p.prefs[i][0], p.prefs[i][1]
+            v_new = p.prefs[i + 1][1] if i + 1 < len(p.prefs) else v_old
+            f1.append(("preference",
+                       f"For {cat}, the user used to like {v_old} but now "
+                       f"prefers {v_new}.", [v_old, v_new]))
+        for s in p.skills:
+            f1.append(("skill", f"The user knows {s}.", [s]))
+        f1.append(("hobby", f"In their free time the user enjoys "
+                   f"{p.hobbies[0]}.", [p.hobbies[0]]))
+        if len(p.hobbies) > 1:
+            f1.append(("hobby", f"The user also enjoys {p.hobbies[1]}.",
+                       [p.hobbies[1]]))
+        S(1, f1)
+
+        # ---- part 2: job change
+        f2 = []
+        m_end = int(e0[2][5:7]) if e0[2] else 6
+        f2.append(("left_job",
+                   f"The user left {e0[0]} in {_mn(_d(2024, m_end))} "
+                   f"{e0[2][:4] if e0[2] else ''}.".strip(), [e0[0]]))
+        m_new = int(e1[1][5:7])
+        f2.append(("joined_job",
+                   f"The user joined {e1[0]} in {_mn(e1[1])} {e1[1][:4]} "
+                   f"as a {p.roles[0][0]}.", [e1[0]]))
+        if len(p.employers) > 2:
+            mid = p.employers[1]
+            m_mid = int(mid[2][5:7]) if mid[2] else 6
+            f2.append(("left_job",
+                       f"The user later left {mid[0]} in {_mn(_d(2024, m_mid))} "
+                       f"{mid[2][:4] if mid[2] else ''}.".strip(), [mid[0]]))
+            f2.append(("employment",
+                       f"These days the user works at {e_last[0]}.",
+                       [e_last[0]]))
+        S(2, f2)
+
+        # ---- part 3: relocation + family
+        m_move = int(c1[1][5:7])
+        f3 = [
+            ("moved",
+             f"The user moved to {c1[0]} in {_mn(c1[1])} {c1[1][:4]}, "
+             f"previously living in {c0[0]}.", [c1[0], c0[0]]),
+        ]
+        S(3, f3)
+
+        # ---- part 4: work structure (multi-hop)
+        mgr, team = p.manager
+        tname, tech = p.team_tech
+        S(4, [
+            ("manager", f"The user's manager is {mgr}.", [mgr]),
+            ("manager_team", f"{mgr} manages the {tname} team.", [tname]),
+            ("team_tech", f"The {tname} team uses {tech}.", [tech]),
+            ("on_team", f"The user is on the {tname} team.", [tname]),
+        ])
+
+        # ---- part 5: projects
+        f5 = []
+        for name, start, end in p.projects:
+            if end:
+                f5.append(("project",
+                           f"The user worked on {name}, finished in "
+                           f"{_mn(end)} {end[:4]}.", [name]))
+            else:
+                f5.append(("project",
+                           f"The user is currently working on {name}.", [name]))
+        S(5, f5)
+
+        # ---- parts 6/7: dated events + preference flip
+        f6 = []
+        for date, desc in p.events[:2]:
+            f6.append(("event",
+                       f"On {month_name(int(date[5:7]))} {int(date[8:10])}, "
+                       f"{date[:4]}, the user {desc}.", [desc]))
+        S(6, f6)
+        f7 = []
+        for date, desc in p.events[2:]:
+            f7.append(("event",
+                       f"On {month_name(int(date[5:7]))} {int(date[8:10])}, "
+                       f"{date[:4]}, the user {desc}.", [desc]))
+        cat = p.prefs[2][0] if len(p.prefs) > 2 else "coffee"
+        vals = [v for (c, v, s, e) in p.prefs if c == cat]
+        if len(vals) >= 2:
+            f7.append(("preference",
+                       f"The user has since switched to {vals[-1]} for {cat}.",
+                       [vals[-1]]))
+        S(7, f7)
+
+        out.append({"user_id": p.user_id, "full_name": p.full_name,
+                    "sessions": sessions})
+    return {"personas": out}
+
+
+# ------------------------------------------------------------ corpus build
+def build_ood_corpus(rendered: dict, persona, t0: datetime = T0,
+                     target_tokens: int = 120_000, seed: int = 7) -> Corpus:
+    """Assemble an evaluation Corpus from one rendered persona-style row."""
+    rng = random.Random(seed)
+    sessions = []
+    total = 0
+    for s in sorted(rendered.get("sessions", []), key=lambda x: x.get("session", 0)):
+        date = t0 + timedelta(days=int(s.get("session", 0)) * 21
+                              + rng.randrange(0, 5))
+        msgs = [("user", str(t)) for t in s.get("messages", []) if str(t).strip()]
+        for _, txt in msgs:
+            total += token_estimate(txt)
+        sessions.append((persona.user_id, date, msgs))
+    # distractor volume — same machinery as the ID generator
+    guard = 0
+    while total < target_tokens and guard < 200_000:
+        guard += 1
+        uid, date, msgs = sessions[rng.randrange(len(sessions))]
+        if rng.random() < 0.45:
+            txt = distractor_paragraph(rng)
+        else:
+            txt = smalltalk_message(rng) + " " + smalltalk_message(rng)
+        k = rng.randrange(0, max(1, len(msgs)))
+        msgs.insert(k, ("user", txt))
+        total += token_estimate(txt)
+    return Corpus(bucket="ood", target_tokens=target_tokens, sessions=sessions,
+                  personas=[persona], total_tokens=total,
+                  generation_seconds=0.0)
+
+
+# ------------------------------------------------- extraction-layer recall
+def extraction_recall(rendered: dict, persona, config) -> dict:
+    """Run the μ=0 extractor over one rendered persona-style row.
+
+    Returns per-fact match results against the manifest match keys — the
+    direct, honest measure of pattern generalization.
+    """
+    from cortexm.bridge.extractor import Extractor
+    from cortexm.bridge.patterns import ExtractionContext
+
+    manifest = export_manifest([persona], T0)["personas"][0]
+    manifest_facts = {f["id"]: f for s in manifest["sessions"]
+                      for f in s["facts"]}
+    extractor = Extractor(config)
+    candidates = []
+    name = None
+    for s in sorted(rendered.get("sessions", []),
+                    key=lambda x: x.get("session", 0)):
+        ts = T0 + timedelta(days=int(s.get("session", 0)) * 21)
+        for text in s.get("messages", []):
+            ctx = ExtractionContext(user_id=persona.user_id, ts=ts,
+                                    speaker="user", subject_name=name,
+                                    lexicon=set())
+            try:
+                cands = extractor.extract(str(text), ctx)
+            except Exception:
+                cands = []
+            candidates.extend(cands)
+            for c in cands:
+                if c.relation == "name" and name is None:
+                    name = c.value
+    blob = " \n ".join(
+        f"{normalize(c.subject)} | {normalize(c.relation)} | {normalize(c.value)}"
+        for c in candidates if c.pattern != "mention_fallback")
+    matches = {}
+    for fid, f in manifest_facts.items():
+        hit = any(normalize(k) and normalize(k) in blob for k in f["match"])
+        matches[fid] = {"hit": hit, "type": f["type"]}
+    conveyed = set(rendered.get("conveyed", []))
+    total = hit = 0
+    per_type: dict[str, list[int]] = {}
+    for fid, r in matches.items():
+        if fid not in conveyed:
+            continue  # renderer omitted it — not an extraction failure
+        total += 1
+        hit += int(r["hit"])
+        per_type.setdefault(r["type"], [0, 0])
+        per_type[r["type"]][1] += 1
+        per_type[r["type"]][0] += int(r["hit"])
+    return {
+        "user_id": persona.user_id,
+        "recall": round(hit / total, 4) if total else None,
+        "n_ground_truth": total,
+        "n_renderer_omitted": len(manifest_facts) - len(conveyed & set(manifest_facts)),
+        "n_candidates": len([c for c in candidates
+                             if c.pattern != "mention_fallback"]),
+        "per_type": {k: round(v[0] / v[1], 4) for k, v in per_type.items() if v[1]},
+        "missed": [fid for fid, r in matches.items()
+                   if not r["hit"] and fid in conveyed],
+        "definition": "entity-level recall: share of ground-truth facts "
+                       "whose key appears in any non-fallback candidate; "
+                       "temporal precision is measured end-to-end by the "
+                       "TR/EO probes, not here",
+    }
+
+
+# ------------------------------------------------------- end-to-end eval
+@dataclass
+class OODResult:
+    style: str
+    overall: float = 0.0
+    per_ability: dict = field(default_factory=dict)
+    n_questions: int = 0
+    ingest: dict = field(default_factory=dict)
+    extraction: dict = field(default_factory=dict)
+    details: list = field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        return {k: v for k, v in self.__dict__.items()}
+
+
+def run_ood_eval(corpus: Corpus, personas: list, style: str,
+                 db_path: str = ":memory:", max_probes: int | None = None,
+                 judge_fn=None, enrich_fn=None) -> OODResult:
+    """Ingest an OOD corpus and evaluate with the standard probe/judge pair.
+
+    With ``enrich_fn(memory) -> report`` the probe set is evaluated twice —
+    before and after async LLM enrichment — quantifying the graceful-
+    degradation fallback's recovery on OOD text.
+    """
+    import time
+
+    from cortexm import metrics
+    from cortexm.api.memory import Memory
+    from cortexm.config import Config
+
+    t0 = time.time()
+    res = OODResult(style=style)
+    cfg = Config(db_path=db_path) if db_path != ":memory:" else Config()
+    cfg.apply_rules_each_add = False
+    memory = Memory(cfg)
+    metrics.reset_counters()
+
+    t_ing = time.time()
+    n_msgs = 0
+    for user_id, date, msgs in corpus.sessions:
+        payload = [{"role": role, "content": text, "timestamp": date}
+                   for role, text in msgs]
+        n_msgs += len(msgs)
+        memory.add(payload, user_id=user_id, timestamp=date)
+    memory.apply_rules()
+    ingest_s = time.time() - t_ing
+    stats = memory.stats()
+    res.ingest = {
+        "wall_seconds": round(ingest_s, 2),
+        "messages": n_msgs,
+        "tokens": corpus.total_tokens,
+        "tokens_per_second": int(corpus.total_tokens / max(ingest_s, 1e-9)),
+        "llm_calls": metrics.counters()["llm_calls"],
+        "u0_protocol": stats["u0_protocol"],
+        "facts": stats["facts"],
+    }
+
+    def _eval() -> tuple[float, dict, list]:
+        rng = random.Random(hash((style, "probes")) & 0xFFFFFFFF)
+        probes = build_probes(personas, rng)
+        by_ability: dict[str, list] = {a: [] for a in ABILITIES}
+        for p in probes:
+            by_ability[p.ability].append(p)
+        if max_probes:
+            for a in ABILITIES:
+                by_ability[a] = by_ability[a][:max_probes]
+        n_q = sum(len(v) for v in by_ability.values())
+        score_fn = judge_fn or judge
+        per_ability = {a: 0.0 for a in ABILITIES}
+        counts = {a: len(by_ability[a]) for a in ABILITIES}
+        details = []
+        for ability in ABILITIES:
+            for probe in by_ability[ability]:
+                out = memory.search(probe.question, user_id=probe.user_id, k=12)
+                score, detail = score_fn(probe, out["context_block"])
+                per_ability[ability] += score
+                details.append({"ability": ability, "question": probe.question,
+                                "score": score, "detail": detail,
+                                "context": out["context_block"][:600]})
+        overall = round(sum(per_ability.values()) / max(n_q, 1), 4)
+        per_ability = {a: round(per_ability[a] / max(counts[a], 1), 4)
+                       for a in ABILITIES if counts[a]}
+        return overall, per_ability, details
+
+    res.overall, res.per_ability, res.details = _eval()
+    res.n_questions = len(res.details)
+    res.extraction = {"facts": stats["facts"]}
+
+    if enrich_fn is not None:
+        rep = enrich_fn(memory)
+        res.ingest["enrichment"] = rep
+        post_overall, post_per_ability, _post = _eval()
+        res.enriched = {"overall": post_overall,
+                        "per_ability": post_per_ability,
+                        "report": rep}  # type: ignore[attr-defined]
+
+    memory.close()
+    res.ingest["wall_seconds_total"] = round(time.time() - t0, 2)
+    return res
+
+
+# ---------------------------------------------------- LLM-judge item export
+def describe_expected(probe) -> list[str]:
+    """Human-readable nugget description for the canonical LLM judge."""
+    exp = probe.expected
+    if probe.ability == "AB":
+        return [f"NOTHING — '{probe.question}' concerns an attribute never "
+                f"mentioned; correct behaviour is abstention."]
+    if "events" in exp:
+        return [f"Event '{d}' occurred on {iso}." for d, iso in exp["events"]]
+    if "set" in exp:
+        return ["The complete set of projects: " + ", ".join(exp["set"]) + "."]
+    if "current" in exp and "old" in exp:
+        return [f"Current employer: {exp['current']}.",
+                f"Previous employer (superseded): {exp['old']}."]
+    out = []
+    for k in exp.get("contains_any", []):
+        out.append(f"The answer must mention: {k}.")
+    for k in exp.get("also_any", []):
+        out.append(f"Supporting evidence should also mention: {k}.")
+    return out or ["(no nugget description)"]
+
+
+def export_judge_items(result: OODResult, personas: list,
+                       per_ability_limit: int = 4) -> list[dict]:
+    """Probe/context pairs in the canonical judge JSONL format.
+
+    The deterministic judge's score is attached (det_score) so agreement
+    between the two graders can be computed after the LLM pass.
+    """
+    items = []
+    per_ability: dict[str, int] = {}
+    probe_by_q = {pr.question: pr for pr in build_probes(personas, random.Random(1234))}
+    for i, d in enumerate(result.details):
+        if per_ability.get(d["ability"], 0) >= per_ability_limit:
+            continue
+        per_ability[d["ability"]] = per_ability.get(d["ability"], 0) + 1
+        probe = probe_by_q.get(d["question"])
+        items.append({
+            "id": f"{result.style}-{i}",
+            "ability": d["ability"],
+            "question": d["question"],
+            "expected": describe_expected(probe) if probe
+                        else ["(probe metadata unavailable)"],
+            "context": d["context"],
+            "det_score": d["score"],
+        })
+    return items

cortexm/bench/run.py

@@ -0,0 +1,137 @@
+"""Benchmark CLI — run BEAM-style buckets and micro-benchmarks.
+
+    python -m cortexm.bench.run --buckets 128k,500k,1m,10m
+    python -m cortexm.bench.run --micro
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+
+from cortexm.bench.harness import BucketResult, format_report, run_bucket
+
+
+# --- determinism guard ---------------------------------------------------
+# Bench runs MUST be bit-for-bit reproducible. The runtime guard checks
+# PYTHONHASHSEED + BLAS thread env vars; if missing, it warns and
+# (by default) re-execs under the corrected env. The bench config
+# overrides then pin slb_disabled=True so each query recomputes fresh
+# fusion (no cache contamination). Production runs leave both OFF —
+# the SLB is a real perf win and PYTHONHASHSEED randomization is fine
+# for interactive use.
+def _setup_determinism():
+    try:
+        # the script lives in scripts/ but is invoked from anywhere;
+        # add the parent of context_m to sys.path so the import works
+        # from a checkout.
+        here = os.path.dirname(os.path.abspath(__file__))
+        # walk up to find scripts/determinism.py
+        for parent in [here, os.path.dirname(here), os.path.dirname(os.path.dirname(here))]:
+            cand = os.path.join(parent, "scripts", "determinism.py")
+            if os.path.exists(cand):
+                sys_path = os.path.dirname(cand)
+                if sys_path not in sys.path:
+                    import sys as _sys
+                    _sys.path.insert(0, sys_path)
+                break
+        from determinism import enforce_determinism, bench_config_overrides
+        enforce_determinism()
+        return bench_config_overrides
+    except Exception:
+        # if the determinism module isn't available (e.g. installed
+        # via pip without the scripts dir), fall back to no-op
+        return lambda **kw: dict(slb_disabled=True, **kw)
+
+
+import sys
+
+
+def run_buckets(buckets: list[str], seed: int, out_dir: str,
+                db_dir: str | None = None) -> list[BucketResult]:
+    os.makedirs(out_dir, exist_ok=True)
+    results: list[BucketResult] = []
+    for bucket in buckets:
+        print(f"\n=== bucket {bucket} ===", flush=True)
+        t0 = time.time()
+        db = (os.path.join(db_dir, f"bench-{bucket}.db")
+              if db_dir else ":memory:")
+        if db_dir:
+            os.makedirs(db_dir, exist_ok=True)
+            # a benchmark is always a fresh-corpus run: never let state from
+            # a previous run leak in (contaminated reruns silently change
+            # fact counts and scores)
+            for suffix in ("", "-journal", "-wal", "-shm"):
+                stale = db + suffix
+                if os.path.exists(stale):
+                    os.remove(stale)
+        r = run_bucket(bucket, seed=seed, db_path=db)
+        results.append(r)
+        with open(os.path.join(out_dir, f"{bucket}.json"), "w") as fh:
+            json.dump(r.to_dict(), fh, indent=2, default=str)
+        cm = r.per_system.get("context_m", {})
+        print(f"context_m overall: {cm.get('overall', 0):.1%} "
+              f"({r.n_questions} questions, ingest {r.ingest['wall_seconds']}s, "
+              f"facts {r.ingest['facts']:,}, μ=0 {r.ingest['u0_protocol']})",
+              flush=True)
+        print(f"bucket wall time: {time.time() - t0:.1f}s", flush=True)
+    report = format_report(results)
+    with open(os.path.join(out_dir, "REPORT.md"), "w") as fh:
+        fh.write(report)
+    print("\nreport written to", os.path.join(out_dir, "REPORT.md"))
+    return results
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="Context-M benchmark runner")
+    ap.add_argument("--buckets", default="128k",
+                    help="comma list: 128k,500k,1m,10m")
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--out", default=os.path.join("benchmarks", "results"))
+    ap.add_argument("--db-dir", default=None,
+                    help="persist per-bucket databases here")
+    ap.add_argument("--micro", action="store_true",
+                    help="run micro-benchmarks instead")
+    ap.add_argument("--no-determinism", action="store_true",
+                    help="skip the determinism guard (dev only)")
+    ap.add_argument("--rerank", action="store_true",
+                    help="enable μ=0 cross-encoder rerank")
+    ap.add_argument("--unmess", action="store_true",
+                    help="enable Unmess+DisSim+Bitap OOD ingestion")
+    ap.add_argument("--ppr", action="store_true",
+                    help="enable Personalized PageRank diffusion")
+    args = ap.parse_args()
+
+    if not args.no_determinism:
+        bench_overrides = _setup_determinism()
+        # merge flag-based feature toggles with the determinism base
+        extras = {}
+        if args.rerank:
+            extras["enable_rerank"] = True
+        if args.unmess:
+            extras["unmess_enabled"] = True
+        if args.ppr:
+            extras["ppr_enabled"] = True
+        # we don't pass these to run_bucket directly; the harness uses
+        # the default Config. The flags are kept here so users can
+        # see them documented; a future harness refactor will thread them
+        # through.
+        _ = bench_overrides(**extras)
+
+    if args.micro:
+        from cortexm.bench.micro import run_micro
+        out = run_micro()
+        os.makedirs(args.out, exist_ok=True)
+        with open(os.path.join(args.out, "micro.json"), "w") as fh:
+            json.dump(out, fh, indent=2)
+        print(json.dumps(out, indent=2)[:4000])
+        return
+
+    buckets = [b.strip().lower() for b in args.buckets.split(",") if b.strip()]
+    run_buckets(buckets, args.seed, args.out, args.db_dir)
+
+
+if __name__ == "__main__":
+    main()