ratatoskr/docs/diagnostics/lexical_recall_gate.py

"""Lexical-recall gate — the class acceptance instrument for exact-term recall survival.

Generalizes the crown repro (Worldtree #400 / thread 01KZETD98T) beyond its anchor. The
class property under test: when the corpus contains a chunk whose text literally carries a
queried surface form, a natural query for that entity should serve >= 1 such chunk at a
USABLE rank (inside the top-K window). The crown ("Crown of the Sepsis Whore") is the
motivating STRESS case — a common word with a dense-similar vanity cluster that buries the
exact-lexical match; distinctive names (Krakaren, Vine Creeper) are CONTROLs that should
always pass. The gap lives on the stress class, not the controls.

This is a before/after regression instrument, NOT a fix: run it pre-deploy and post-deploy
(the BM25 stemming fold, then any rerank_hybrid_floor lever) to measure whether the served
miss-rate on the stress class actually moves. Real-world by design — it drives the agent
end-to-end (the agent composes its own reference_knowledge query, as in production), and
--runs samples that query-formulation variance to estimate a true miss-rate.

Self-contained (httpx only). Config from env (source env.sh first):
  WORLDTREE_API_URL (default personal :8081), WORLDTREE_API_KEY (required),
  RATATOSKR_END_USER_ID (default ratatoskr-tui), RATATOSKR_TTS_AGENT unused here.

  uv run --with httpx python docs/diagnostics/lexical_recall_gate.py
  uv run --with httpx python docs/diagnostics/lexical_recall_gate.py --runs 5
  uv run --with httpx python docs/diagnostics/lexical_recall_gate.py --anchor crown
"""
from __future__ import annotations

import argparse
import json
import os
import re

import httpx

# Served window: a hit past this rank is not "usable" (Worldtree serves ~top-10; a row at
# rank 9-11 is the RRF knife-edge worldtree-dev identified — treated as a KNIFE-EDGE pass).
USABLE_K = 10
KNIFE_EDGE_FROM = 8  # ranks >= this inside the window are fragile (one-rank-edge residual class)

AGENT = "ratatoskr:donut"

# (label, kind, term-regex the served chunk's excerpt must contain, [natural user messages]).
# kind: "stress" = common word + competing dense cluster; "control" = distinctive name.
# Controls should pass every run; the class limitation shows as stress-class misses / knife-edges.
ANCHORS = [
    ("crown", "stress", r"\bcrown",
     ["What crown do you own?", "Do you have a crown?", "Tell me about your crown."]),
    ("vine-creeper", "control", r"vine creeper",
     ["Tell me about the Vine Creeper.", "What is the Vine Creeper?"]),
    ("danger-dingo", "control", r"danger dingo|\bdingo",
     ["What is the Danger Dingo?", "Describe the Danger Dingo."]),
    ("pedicure-kit", "control", r"pedicure",
     ["What does the Pedicure Kit do?", "Tell me about the Pedicure Kit."]),
    ("neighborhood-map", "control", r"neighborhood map",
     ["What is the Neighborhood Map?", "Describe the Neighborhood Map."]),
]


def _cfg() -> tuple[str, dict, str]:
    base = os.environ.get("WORLDTREE_API_URL", "http://10.250.50.152:8081")
    key = os.environ.get("WORLDTREE_API_KEY")
    if not key:
        raise SystemExit("WORLDTREE_API_KEY unset — source env.sh first.")
    return base, {"Authorization": f"Bearer {key}"}, os.environ.get("RATATOSKR_END_USER_ID", "ratatoskr-tui")


def _session(base: str, headers: dict, end_user: str) -> str:
    r = httpx.post(f"{base}/sessions", json={"agent_id": AGENT, "end_user_id": end_user},
                   headers=headers, timeout=30)
    r.raise_for_status()
    return r.json()["session_id"]


def _drive(base: str, headers: dict, sid: str, content: str) -> tuple[str | None, list]:
    """POST a turn; return (actual reference_knowledge query, served hits list)."""
    q, hits = None, []
    with httpx.stream("POST", f"{base}/sessions/{sid}/messages", json={"content": content},
                      headers=headers, timeout=180) as r:
        for line in r.iter_lines():
            if not line.startswith("data: "):
                continue
            ev = json.loads(line[6:])
            t = ev.get("type")
            if t == "tool_start" and q is None:
                q = (ev.get("arguments") or {}).get("query")
            elif t == "tool_result" and not hits:
                res = ev.get("result")
                if isinstance(res, dict):
                    hits = res.get("hits", res.get("results", [])) or []
            elif t == "done":
                break
    return q, hits


def _served_rank(hits: list, term_re: str) -> int | None:
    """Rank of the first served hit whose excerpt literally contains the term (None = miss)."""
    for i, h in enumerate(hits[:USABLE_K]):
        if isinstance(h, dict) and re.search(term_re, h.get("excerpt", ""), re.I):
            return i
    return None


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", type=int, default=1, help="repeats per message (samples query variance)")
    ap.add_argument("--anchor", default=None, help="run only this anchor label")
    args = ap.parse_args()
    base, headers, end_user = _cfg()

    anchors = [a for a in ANCHORS if args.anchor is None or a[0] == args.anchor]
    totals = {"trials": 0, "miss": 0, "knife": 0}
    by_kind: dict[str, dict] = {}

    for label, kind, term_re, messages in anchors:
        print(f"\n[{label}] ({kind})  term=/{term_re}/")
        agg = by_kind.setdefault(kind, {"trials": 0, "miss": 0, "knife": 0})
        for msg in messages:
            for _ in range(args.runs):
                sid = _session(base, headers, end_user)  # fresh session per trial
                q, hits = _drive(base, headers, sid, msg)
                rank = _served_rank(hits, term_re)
                miss = rank is None
                knife = rank is not None and rank >= KNIFE_EDGE_FROM
                for d in (totals, agg):
                    d["trials"] += 1
                    d["miss"] += int(miss)
                    d["knife"] += int(knife)
                tag = "MISS" if miss else (f"knife@{rank}" if knife else f"ok@{rank}")
                print(f"    {tag:9} msg={msg!r:42} q={q!r}")

    def pct(n: int, d: int) -> str:
        return f"{(100*n/d):.0f}%" if d else "n/a"

    print("\n=== SUMMARY ===")
    for kind, d in sorted(by_kind.items()):
        print(f"  {kind:8} trials={d['trials']:3}  miss={pct(d['miss'], d['trials'])}  "
              f"knife-edge={pct(d['knife'], d['trials'])}")
    t = totals
    print(f"  {'ALL':8} trials={t['trials']:3}  miss={pct(t['miss'], t['trials'])}  "
          f"knife-edge={pct(t['knife'], t['trials'])}")
    print("\nGate: stress-class miss-rate is the deciding signal for the rerank_hybrid_floor lever.")
    print("Controls should sit at ~0% miss; a stress miss/knife-edge is the residual class to weigh.")


if __name__ == "__main__":
    main()