Files
stack/dev/scripts/llm_golden.py

390 lines
14 KiB
Python

"""Golden longitudinal evaluation for the chat (P49 Task 7).
Streams the live ``/chat`` SSE endpoint for a fixed set of longitudinal
questions (``tests/llm/golden_lineage.yaml``) and checks each answer
against verified expectations: FR paragraph anchors that must appear in
``sources``, citation labels that must appear in the answer prose,
forbidden patterns (e.g. a bare dollar amount not backed by a valuation
row), lineage timeline events, comment dockets, and the number of
distinct rule eras the sources span. Regressions are filed/swept through
``dev/scripts/nb_issue_filer.py`` (copied beside this script in CI, and
imported by path — never re-implemented here).
Stdlib + httpx + PyYAML only — this script is docker-cp'd into the
``llm`` container and run there against its own ``http://localhost:8000``
(see ``dev/scripts/backends/gitea.py::_gen_llm_golden``); it must not
import ``llm.*``/``pfs.*``.
Usage::
uv run python dev/scripts/llm_golden.py run \\
--url http://localhost:8000 --set tests/llm/golden_lineage.yaml \\
--report report.json [--file-issues --source nightly-llm-golden] \\
[--only g2058-replacement] [--timeout 180]
"""
from __future__ import annotations
import argparse
import dataclasses
import importlib.util
import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import httpx
import yaml
_HERE = Path(__file__).resolve().parent
_EXPECTATION_KEYS = (
"expect_anchors",
"expect_labels",
"forbid",
"expect_events",
"expect_dockets",
"min_eras",
)
_SENT_SPLIT = re.compile(r"(?<=[.!?])\s+")
_BRACKET_RE = re.compile(r"\[([^\]]+)\]")
_DOCKET_YEAR_RE = re.compile(r"(19|20)\d{2}")
_DATE_YEAR_RE = re.compile(r"^(\d{4})")
# ── golden set loading ───────────────────────────────────────────
def load_set(path: Path) -> list[dict]:
"""Parse and validate the golden YAML set: a top-level list (or
``{questions: [...]}``) of entries, each with a unique ``id``, a
``question``, and at least one expectation key."""
raw = yaml.safe_load(path.read_text())
if isinstance(raw, dict) and "questions" in raw:
entries = raw["questions"]
elif isinstance(raw, list):
entries = raw
else:
raise ValueError(f"{path}: expected a list or {{questions: [...]}}")
ids: set[str] = set()
for e in entries:
if not isinstance(e, dict) or "id" not in e or "question" not in e:
raise ValueError(f"{path}: entry missing id/question: {e!r}")
if e["id"] in ids:
raise ValueError(f"{path}: duplicate id {e['id']!r}")
ids.add(e["id"])
if not any(k in e for k in _EXPECTATION_KEYS):
raise ValueError(f"{path}: entry {e['id']!r} has no expectations")
return entries
# ── transcript ───────────────────────────────────────────────────
@dataclass
class Transcript:
events: list[dict] = field(default_factory=list)
@property
def answer_text(self) -> str:
return "".join(
e.get("text", "") for e in self.events if e.get("type") == "token"
)
@property
def sources(self) -> list[dict]:
for e in self.events:
if e.get("type") == "sources":
return e.get("sources", [])
return []
@property
def lineage_events(self) -> list[dict]:
for e in self.events:
if e.get("type") == "lineage":
return e.get("events", [])
return []
@property
def error(self) -> str | None:
for e in self.events:
if e.get("type") == "error":
return e.get("message", "stream error")
return None
@dataclass(frozen=True)
class CheckResult:
name: str
passed: bool
detail: str
# ── checks (pure functions over a Transcript) ───────────────────
def _parse_pid_range(spec: Any) -> tuple[int, int]:
if isinstance(spec, (list, tuple)):
return int(spec[0]), int(spec[1])
return int(spec), int(spec)
def check_anchors(transcript: Transcript, expect_anchors: list[dict]) -> CheckResult:
sources = transcript.sources
missing = []
for a in expect_anchors:
lo, hi = _parse_pid_range(a["p_id"])
found = False
for s in sources:
if s.get("item_key") != a["item_key"]:
continue
try:
pid = int(s.get("p_id") or "")
except (TypeError, ValueError):
continue
if lo <= pid <= hi:
found = True
break
if not found:
label = (
f"{a['item_key']} p{lo}" if lo == hi else f"{a['item_key']} p{lo}-{hi}"
)
missing.append(label)
passed = not missing
detail = "ok" if passed else "missing: " + ", ".join(missing)
return CheckResult("anchors", passed, detail)
def check_labels(transcript: Transcript, expect_labels: list[str]) -> CheckResult:
text = transcript.answer_text
missing = [pat for pat in expect_labels if not re.search(pat, text)]
passed = not missing
detail = "ok" if passed else "missing: " + ", ".join(missing)
return CheckResult("labels", passed, detail)
def check_forbidden(transcript: Transcript, forbid: list[dict]) -> CheckResult:
sentences = _SENT_SPLIT.split(transcript.answer_text)
violations = []
for rule in forbid:
pat = re.compile(rule["pattern"])
unless = rule.get("unless_label")
unless_re = re.compile(unless) if unless else None
for sent in sentences:
if not pat.search(sent):
continue
labels = _BRACKET_RE.findall(sent)
ok = unless_re is not None and any(unless_re.search(l) for l in labels)
if not ok:
violations.append(f"{rule['pattern']!r} in {sent.strip()[:120]!r}")
passed = not violations
detail = "ok" if passed else "; ".join(violations)
return CheckResult("forbidden", passed, detail)
def check_events(transcript: Transcript, expect_events: list[dict]) -> CheckResult:
events = transcript.lineage_events
missing = []
for exp in expect_events:
kinds = set(str(exp["kind"]).split("|"))
code = str(exp["code"])
year = int(exp["year"])
found = any(
e.get("code") == code
and e.get("kind") in kinds
and int(e.get("year", -1)) == year
for e in events
)
if not found:
missing.append(f"{code} {exp['kind']} {year}")
passed = not missing
detail = "ok" if passed else "missing: " + ", ".join(missing)
return CheckResult("events", passed, detail)
def check_dockets(transcript: Transcript, expect_dockets: list[str]) -> CheckResult:
have = {
s.get("docket", "")
for s in transcript.sources
if s.get("kind") == "comment" and s.get("docket")
}
missing = [d for d in expect_dockets if d not in have]
passed = not missing
detail = "ok" if passed else "missing: " + ", ".join(missing)
return CheckResult("dockets", passed, detail)
def _rule_year(source: dict) -> int:
"""The rule year a source belongs to: the docket id's embedded year
for a comment, else the source's own date year — no imports, so this
is a simplification of ``llm.rag.era_of`` (which resolves a
comment's docket to its actual PFS rule year via the bib store)."""
if source.get("kind") == "comment":
m = _DOCKET_YEAR_RE.search(source.get("docket", "") or "")
return int(m.group(0)) if m else 0
m = _DATE_YEAR_RE.match(source.get("date", "") or "")
return int(m.group(1)) if m else 0
def check_eras(transcript: Transcript, min_eras: int) -> CheckResult:
eras = {_rule_year(s) for s in transcript.sources}
eras.discard(0)
passed = len(eras) >= int(min_eras)
detail = f"eras={sorted(eras)}"
return CheckResult("eras", passed, detail)
def evaluate(entry: dict, transcript: Transcript) -> list[CheckResult]:
if transcript.error is not None:
return [CheckResult("stream", False, transcript.error)]
results: list[CheckResult] = []
if "expect_anchors" in entry:
results.append(check_anchors(transcript, entry["expect_anchors"]))
if "expect_labels" in entry:
results.append(check_labels(transcript, entry["expect_labels"]))
if "forbid" in entry:
results.append(check_forbidden(transcript, entry["forbid"]))
if "expect_events" in entry:
results.append(check_events(transcript, entry["expect_events"]))
if "expect_dockets" in entry:
results.append(check_dockets(transcript, entry["expect_dockets"]))
if "min_eras" in entry:
results.append(check_eras(transcript, entry["min_eras"]))
return results
# ── streaming the live /chat endpoint ───────────────────────────
def stream_chat(url: str, question: str, mode: str, timeout: float) -> list[dict]:
"""POST ``{question, mode}`` to ``{url}/chat`` and parse the
``data: {json}\\n\\n`` SSE lines into a list of event dicts."""
events: list[dict] = []
with httpx.Client(timeout=timeout) as client:
with client.stream(
"POST", f"{url.rstrip('/')}/chat", json={"question": question, "mode": mode}
) as resp:
resp.raise_for_status()
for line in resp.iter_lines():
if not line or not line.startswith("data:"):
continue
payload = line[len("data:") :].strip()
if payload:
events.append(json.loads(payload))
return events
# ── issue filer (imported by path — nb_issue_filer.py sits beside
# this script both in-repo and when docker-cp'd into the container) ──
def _load_filer():
spec = importlib.util.spec_from_file_location(
"nb_issue_filer", _HERE / "nb_issue_filer.py"
)
assert spec and spec.loader
mod = importlib.util.module_from_spec(spec)
sys.modules["nb_issue_filer"] = mod
spec.loader.exec_module(mod)
return mod
# ── runner ───────────────────────────────────────────────────────
def cmd_run(args: argparse.Namespace) -> int:
entries = load_set(Path(args.set_path))
if args.only:
entries = [e for e in entries if e["id"] == args.only]
if not entries:
print(f"no entry with id {args.only!r} in {args.set_path}", file=sys.stderr)
return 2
rows: list[dict] = []
findings: list[dict] = []
any_fail = False
for entry in entries:
qid = entry["id"]
mode = entry.get("mode", "auto")
try:
events = stream_chat(args.url, entry["question"], mode, args.timeout)
results = evaluate(entry, Transcript(events))
except Exception as e: # noqa: BLE001 — a request failure is a finding, not a crash
results = [CheckResult("request", False, str(e))]
failing = [r for r in results if not r.passed]
if failing:
any_fail = True
rows.append(
{
"id": qid,
"passed": not failing,
"checks": [dataclasses.asdict(r) for r in results],
}
)
for r in failing:
findings.append(
{
"notebook": qid,
"ename": r.name,
"evalue": r.detail,
"detail": r.detail,
}
)
for row in rows:
status = "PASS" if row["passed"] else "FAIL"
failing_names = ",".join(c["name"] for c in row["checks"] if not c["passed"])
print(f"{row['id']:32} {status:4} {failing_names}")
report = {
"rows": rows,
"pass": sum(r["passed"] for r in rows),
"fail": sum(not r["passed"] for r in rows),
}
if args.report:
Path(args.report).write_text(json.dumps(report, indent=2))
if args.file_issues:
filer = _load_filer()
filer.cmd_report(findings, args.source)
active = {
filer.signature(f["notebook"], f["ename"], f["evalue"]) for f in findings
}
filer.cmd_sweep(active, args.source)
return 1 if any_fail else 0
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
sub = parser.add_subparsers(dest="cmd", required=True)
p_run = sub.add_parser("run", help="stream the golden set against a live /chat")
p_run.add_argument("--url", required=True, help="base URL of the chat service")
p_run.add_argument("--set", required=True, dest="set_path", help="golden YAML path")
p_run.add_argument("--report", help="write a JSON report here")
p_run.add_argument(
"--file-issues",
action="store_true",
help="file/sweep regressions via nb_issue_filer",
)
p_run.add_argument("--source", default="llm-golden", help="filer 'source' label")
p_run.add_argument("--only", help="run only this entry id")
p_run.add_argument("--timeout", type=float, default=180.0)
args = parser.parse_args(argv)
if args.cmd == "run":
return cmd_run(args)
parser.error(f"unknown command {args.cmd!r}")
return 2 # pragma: no cover — argparse exits before this
if __name__ == "__main__":
raise SystemExit(main())