Evaluate RAG retrieval, faithfulness, citations, and abstention

#!/usr/bin/env python3
import argparse
import json
import math
from pathlib import Path

def rows(path):
    with Path(path).open(encoding="utf-8") as handle:
        return [json.loads(line) for line in handle if line.strip()]

def by_id(items):
    result = {item["id"]: item for item in items}
    if len(result) != len(items):
        raise SystemExit("duplicate record id")
    return result

def percentile(values, fraction):
    ordered = sorted(values)
    if not ordered:
        return 0.0
    return ordered[min(len(ordered) - 1, math.ceil(len(ordered) * fraction) - 1)]

def labels_audit(args):
    data = rows(args.dataset)
    selected = [row for row in data if not args.critical_only or row.get("critical")]
    double = sum(len(row.get("reviewers", [])) >= 2 for row in selected)
    missing = sum(row.get("answerable") and not row.get("gold_evidence_ids") for row in selected)
    unresolved = sum(bool(row.get("review_disagreement")) and not row.get("adjudicated") for row in selected)
    result = {"records": len(selected), "double_reviewed": double, "missing_gold": missing, "unresolved": unresolved}
    print(json.dumps(result))
    if double != len(selected) or missing or unresolved:
        raise SystemExit(1)

def validate(args):
    data = by_id(rows(args.dataset))
    runs = rows(args.run)
    counts = {"responses": len(runs), "acl_violations": 0, "missing_ids": 0, "stale_revisions": 0, "unknown_citations": 0, "schema_errors": 0}
    for run in runs:
        case = data.get(run.get("id"))
        if case is None:
            counts["schema_errors"] += 1
            continue
        evidence = run.get("evidence", [])
        evidence_ids = {item.get("id") for item in evidence}
        for item in evidence:
            if item.get("tenant_id") != case.get("tenant_id") or item.get("authorized") is not True or item.get("state") != "active":
                counts["acl_violations"] += 1
            if not item.get("id"):
                counts["missing_ids"] += 1
            if item.get("corpus_revision") != case.get("corpus_revision"):
                counts["stale_revisions"] += 1
        counts["unknown_citations"] += len(set(run.get("cited_ids", [])) - evidence_ids)
    print(json.dumps(counts))
    if any(counts[key] for key in counts if key != "responses"):
        raise SystemExit(1)

def retrieval(args):
    data = by_id(rows(args.dataset))
    runs = rows(args.run)
    recalls, precisions, reciprocal, ndcg = [], [], [], []
    for run in runs:
        gold = set(data[run["id"]].get("gold_evidence_ids", []))
        ranked = list(dict.fromkeys(run.get("retrieved_ids", [])))[:args.k]
        hits = [1 if item in gold else 0 for item in ranked]
        recalls.append(1.0 if not gold else len(gold.intersection(ranked)) / len(gold))
        precisions.append(sum(hits) / max(1, len(ranked)))
        first = next((index + 1 for index, hit in enumerate(hits) if hit), None)
        reciprocal.append(0 if first is None else 1 / first)
        dcg = sum(hit / math.log2(index + 2) for index, hit in enumerate(hits))
        ideal = sum(1 / math.log2(index + 2) for index in range(min(len(gold), len(ranked))))
        ndcg.append(0 if ideal == 0 else dcg / ideal)
    result = {"records": len(runs), "recall": sum(recalls)/len(recalls), "precision": sum(precisions)/len(precisions), "mrr": sum(reciprocal)/len(reciprocal), "ndcg": sum(ndcg)/len(ndcg)}
    print(json.dumps(result))

def grounding(args):
    runs = rows(args.run)
    claims = [claim for run in runs for claim in run.get("claims", []) if claim.get("material", True)]
    supported = sum(claim.get("supported") is True for claim in claims)
    citations = sum(claim.get("citation_valid") is True for claim in claims)
    critical_failures = sum(claim.get("critical") and claim.get("supported") is not True for claim in claims)
    result = {"answers": len(runs), "material_claims": len(claims), "claim_support": supported/max(1,len(claims)), "citation_support": citations/max(1,len(claims)), "critical_unsupported": critical_failures}
    print(json.dumps(result))
    if critical_failures:
        raise SystemExit(1)

def answer_quality(args):
    candidate = rows(args.run)
    control = by_id(rows(args.control))
    metrics = ("correct", "complete", "useful")
    result = {name: sum(float(row.get(name, 0)) for row in candidate)/len(candidate) for name in metrics}
    result["pairwise_candidate_win"] = sum(float(row.get("pairwise_candidate_win", 0)) for row in candidate)/len(candidate)
    result["control_correct"] = sum(float(control[row["id"]].get("correct", 0)) for row in candidate)/len(candidate)
    print(json.dumps(result))

def abstention(args):
    data = by_id(rows(args.dataset))
    runs = rows(args.run)
    unanswerable = [run for run in runs if not data[run["id"]].get("answerable")]
    answerable = [run for run in runs if data[run["id"]].get("answerable")]
    true_abstain = sum(run.get("abstained") is True for run in unanswerable)
    all_abstain = sum(run.get("abstained") is True for run in runs)
    result = {"unanswerable": len(unanswerable), "abstention_recall": true_abstain/max(1,len(unanswerable)), "abstention_precision": true_abstain/max(1,all_abstain), "over_abstained": sum(run.get("abstained") is True for run in answerable), "forbidden_existence_leaks": sum(bool(run.get("forbidden_existence_leak")) for run in runs)}
    print(json.dumps(result))
    if result["forbidden_existence_leaks"]:
        raise SystemExit(1)

def report(args):
    candidate = rows(args.run)
    control = by_id(rows(args.control))
    blockers = {"acl_violations": sum(int(row.get("acl_violations",0)) for row in candidate), "unknown_citations": sum(int(row.get("unknown_citations",0)) for row in candidate), "critical_unsupported": sum(int(row.get("critical_unsupported",0)) for row in candidate)}
    result = {"records": len(candidate), "blockers": blockers, "p95_ms": percentile([float(row["latency_ms"]) for row in candidate], .95), "estimated_cost_per_query": sum(float(row.get("estimated_cost_usd",0)) for row in candidate)/len(candidate), "candidate_wins": sum(float(row.get("correct",0)) > float(control[row["id"]].get("correct",0)) for row in candidate)}
    result["release_gate"] = "PASS" if not any(blockers.values()) else "FAIL"
    Path(args.output).write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result))
    if result["release_gate"] != "PASS":
        raise SystemExit(1)

def calibrate(args):
    records = rows(args.set)
    repeats = [decision for row in records for decision in row.get("grader_passes", [])[:args.repeats]]
    paired = [(row["human_pass"], decision) for row in records for decision in row.get("grader_passes", [])[:args.repeats]]
    agreement = sum(human == grader for human, grader in paired) / max(1, len(paired))
    critical_false_pass = sum(row.get("critical") and not human and grader for row in records for human, grader in [(row["human_pass"], row.get("grader_passes", [False])[0])])
    result = {"records": len(records), "decisions": len(repeats), "agreement": agreement, "critical_false_pass": critical_false_pass}
    print(json.dumps(result))
    if agreement < .85 or critical_false_pass:
        raise SystemExit(1)

parser = argparse.ArgumentParser()
sub = parser.add_subparsers(dest="command", required=True)
for name in ("labels-audit", "validate", "retrieval", "grounding", "answer-quality", "abstention", "report", "calibrate"):
    command = sub.add_parser(name)
    command.add_argument("--dataset")
    command.add_argument("--run")
    command.add_argument("--control")
    command.add_argument("--output")
    command.add_argument("--set")
    command.add_argument("--critical-only", action="store_true")
    command.add_argument("--k", type=int, default=20)
    command.add_argument("--repeats", type=int, default=3)
args = parser.parse_args()
{"labels-audit": labels_audit, "validate": validate, "retrieval": retrieval, "grounding": grounding, "answer-quality": answer_quality, "abstention": abstention, "report": report, "calibrate": calibrate}[args.command](args)

python {{evaluationHarness}} validate --dataset {{evaluationDataset}} --run {{candidateRun}}

python {{evaluationHarness}} retrieval --dataset {{evaluationDataset}} --run {{candidateRun}} --k 20

python {{evaluationHarness}} grounding --run {{candidateRun}}

python {{evaluationHarness}} answer-quality --run {{candidateRun}} --control {{controlRun}}

python {{evaluationHarness}} abstention --dataset {{evaluationDataset}} --run {{candidateRun}}

python {{evaluationHarness}} report --run {{candidateRun}} --control {{controlRun}} --output {{evaluationReport}}

python {{evaluationHarness}} calibrate --set {{calibrationSet}} --repeats 3

sha256sum -c {{ragReleaseChecksums}} && sudo ln -sfn {{ragCandidateDir}} /srv/rag/candidate && sudo systemctl restart rag@candidate.service && curl -fsS http://127.0.0.1:8401/healthz | jq -e '.release == "rag-r18" and .ready == true' && jq -e '.blockers.acl_violations == 0 and .blockers.unknown_citations == 0 and .blockers.critical_unsupported == 0 and .p95_ms <= 900' {{evaluationReport}} && jq -e '.critical_probe_failures == 0 and .error_rate <= .01' {{ragCanaryReport}}