Nomearod/agentbench

1"""Verify retrieval quality against golden dataset.2 3Runs the Day 4 gate check: for each positive golden question,4does hybrid retrieval return the expected source in top-5?5 6Usage:7    python scripts/verify_retrieval.py8    python scripts/verify_retrieval.py --store-path .cache/store --output docs/retrieval_gate.md9"""10 11from __future__ import annotations12 13import argparse14import json15import sys16from pathlib import Path17 18sys.path.insert(0, str(Path(__file__).resolve().parent.parent))19 20from agent_bench.rag.embedder import Embedder21from agent_bench.rag.store import HybridStore22 23 24def verify(25    store_path: str = ".cache/store",26    golden_path: str = "agent_bench/evaluation/datasets/tech_docs_golden.json",27    model_name: str = "all-MiniLM-L6-v2",28    cache_dir: str = ".cache/embeddings",29    output_path: str | None = None,30) -> bool:31    store = HybridStore.load(store_path)32    embedder = Embedder(model_name=model_name, cache_dir=cache_dir)33 34    with open(golden_path) as f:35        questions = json.load(f)36 37    lines: list[str] = []38    lines.append("# Retrieval Gate Check")39    lines.append("")40    lines.append(41        f"**Store:** {store.stats().total_chunks} chunks, "42        f"{store.stats().unique_sources} sources"43    )44    lines.append("")45    lines.append("| ID | Category | Expected Source | Top-5 Sources | Recall@5 | Result |")46    lines.append("|-----|----------|----------------|---------------|----------|--------|")47 48    total_recall = 0.049    scorable = 050 51    for q in questions:52        qid = q["id"]53        question = q["question"]54        expected = set(q["expected_sources"])55        category = q["category"]56 57        vec = embedder.embed(question)58        results = store.search(vec, question, top_k=5, strategy="hybrid")59        retrieved = [r.chunk.source for r in results]60        retrieved_set = set(retrieved)61 62        if expected:63            recall = len(expected & retrieved_set) / len(expected)64            total_recall += recall65            scorable += 166            result = "PASS" if recall >= 0.5 else "FAIL"67        else:68            recall = float("nan")69            result = "N/A"70 71        expected_str = ", ".join(sorted(expected)) if expected else "(none)"72        retrieved_str = ", ".join(dict.fromkeys(retrieved[:3]))  # dedup, first 373        recall_str = f"{recall:.2f}" if expected else "n/a"74        lines.append(75            f"| {qid} | {category} | {expected_str} | {retrieved_str} | {recall_str} | {result} |"76        )77 78    avg_recall = total_recall / max(scorable, 1)79    gate_pass = avg_recall >= 0.580 81    lines.append("")82    lines.append(f"**Avg Recall@5 (positive only):** {avg_recall:.2f}")83    lines.append(f"**Gate:** {'PASS' if gate_pass else 'FAIL'} (threshold >= 0.5)")84 85    report = "\n".join(lines)86    print(report)87 88    if output_path:89        Path(output_path).parent.mkdir(parents=True, exist_ok=True)90        Path(output_path).write_text(report + "\n")91        print(f"\nSaved to {output_path}")92 93    return gate_pass94 95 96def main() -> None:97    parser = argparse.ArgumentParser(description="Verify retrieval against golden dataset")98    parser.add_argument("--store-path", default=".cache/store")99    parser.add_argument("--golden-path", default="agent_bench/evaluation/datasets/tech_docs_golden.json")100    parser.add_argument("--output", default="docs/retrieval_gate.md")101    args = parser.parse_args()102 103    passed = verify(104        store_path=args.store_path,105        golden_path=args.golden_path,106        output_path=args.output,107    )108    sys.exit(0 if passed else 1)109 110 111if __name__ == "__main__":112    main()113