Benchmark run complete: 595/600 evaluations, grounded cuts misstatements 56% vs naive

scripts/run_benchmark.py

@@ -0,0 +1,189 @@
+"""
+Runs the full benchmark: all gold queries through all 5 RAG configs,
+scored by the two-layer judge in src/judge.py.
+
+Writes results/raw_results.jsonl (one line per (query, config) pair) and
+results/summary.json (taxonomy bucket counts per config).
+
+Resumable: on (re)start, already-written (query_id, config_name) pairs in
+raw_results.jsonl are skipped, so a crash mid-run doesn't redo finished work.
+"""
+
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+from db import get_connection
+from eval import evaluate_query
+from judge import evaluate_answer
+
+QUERIES_PATH = Path(__file__).resolve().parent.parent / "queries" / "gold_set.jsonl"
+RESULTS_DIR = Path(__file__).resolve().parent.parent / "results"
+RAW_RESULTS_PATH = RESULTS_DIR / "raw_results.jsonl"
+SUMMARY_PATH = RESULTS_DIR / "summary.json"
+
+CONFIG_NAMES = ["naive", "hybrid", "rerank", "kg_augmented", "grounded"]
+TAXONOMY_BUCKETS = [
+    "correct_and_faithful",
+    "fabricated_citation",
+    "misstated_obligation",
+    "missing_citation",
+    "real_but_irrelevant",
+]
+
+
+def load_gold_queries() -> list[dict]:
+    with open(QUERIES_PATH) as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def load_completed_pairs() -> set[tuple[str, str]]:
+    completed = set()
+    if not RAW_RESULTS_PATH.exists():
+        return completed
+
+    with open(RAW_RESULTS_PATH) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            record = json.loads(line)
+            if "taxonomy_bucket" in record:
+                completed.add((record["query_id"], record["config_name"]))
+
+    return completed
+
+
+def compute_summary() -> dict:
+    counts = {
+        config_name: {bucket: 0 for bucket in TAXONOMY_BUCKETS}
+        for config_name in CONFIG_NAMES
+    }
+    totals = {config_name: 0 for config_name in CONFIG_NAMES}
+
+    with open(RAW_RESULTS_PATH) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            record = json.loads(line)
+            bucket = record.get("taxonomy_bucket")
+            config_name = record["config_name"]
+            if bucket is None or config_name not in counts:
+                continue
+            totals[config_name] += 1
+            counts[config_name][bucket] += 1
+
+    return {
+        config_name: {"total": totals[config_name], **counts[config_name]}
+        for config_name in CONFIG_NAMES
+    }
+
+
+def run_benchmark() -> None:
+    RESULTS_DIR.mkdir(exist_ok=True)
+
+    gold_queries = load_gold_queries()
+    total_queries = len(gold_queries)
+    completed = load_completed_pairs()
+
+    with open(RAW_RESULTS_PATH, "a") as out_f:
+        for i, gold_query in enumerate(gold_queries, start=1):
+            pending_configs = [
+                config_name
+                for config_name in CONFIG_NAMES
+                if (gold_query["id"], config_name) not in completed
+            ]
+
+            if not pending_configs:
+                print(f"[{i}/{total_queries}] {gold_query['id']} | SKIPPED (already complete)")
+                continue
+
+            evaluation = evaluate_query(gold_query["question"])
+
+            conn = get_connection()
+            try:
+                for config_name in CONFIG_NAMES:
+                    if (gold_query["id"], config_name) in completed:
+                        print(f"[{i}/{total_queries}] {gold_query['id']} | {config_name} | SKIPPED (resume)")
+                        continue
+
+                    config_result = evaluation["results"][config_name]
+
+                    if "error" in config_result:
+                        record = {
+                            "query_id": gold_query["id"],
+                            "config_name": config_name,
+                            "error": config_result["error"],
+                        }
+                        out_f.write(json.dumps(record) + "\n")
+                        out_f.flush()
+                        print(f"[{i}/{total_queries}] {gold_query['id']} | {config_name} | ERROR")
+                        continue
+
+                    try:
+                        judged = evaluate_answer(
+                            query_id=gold_query["id"],
+                            question=gold_query["question"],
+                            gold_answer=gold_query["gold_answer"],
+                            gold_citations=gold_query["gold_citations"],
+                            rag_answer=config_result["answer"],
+                            config_name=config_name,
+                            conn=conn,
+                        )
+                    except Exception as e:
+                        record = {
+                            "query_id": gold_query["id"],
+                            "config_name": config_name,
+                            "error": str(e),
+                        }
+                        out_f.write(json.dumps(record) + "\n")
+                        out_f.flush()
+                        print(f"[{i}/{total_queries}] {gold_query['id']} | {config_name} | ERROR")
+                        continue
+
+                    bucket = judged["taxonomy_bucket"]
+
+                    record = {
+                        "query_id": gold_query["id"],
+                        "question": gold_query["question"],
+                        "gold_answer": gold_query["gold_answer"],
+                        "gold_citations": gold_query["gold_citations"],
+                        "config_name": config_name,
+                        "rag_answer": config_result["answer"],
+                        "taxonomy_bucket": bucket,
+                        "fabrication_rate": judged["layer1"]["fabrication_rate"],
+                        "layer2_judgment": judged["layer2"]["judgment"],
+                        "layer2_confidence": judged["layer2"]["confidence"],
+                        "layer2_reasoning": judged["layer2"]["reasoning"],
+                        "retrieved_paragraph_ids": [
+                            chunk["paragraph_id"] for chunk in config_result["retrieved_chunks"]
+                        ],
+                        "stratum": gold_query["stratum"],
+                    }
+                    out_f.write(json.dumps(record) + "\n")
+                    out_f.flush()
+
+                    print(f"[{i}/{total_queries}] {gold_query['id']} | {config_name} | {bucket}")
+            finally:
+                conn.close()
+
+    summary = compute_summary()
+
+    with open(SUMMARY_PATH, "w") as f:
+        json.dump(summary, f, indent=2)
+
+    print()
+    header = f"{'config':<15}{'total':>8}" + "".join(f"{b:>22}" for b in TAXONOMY_BUCKETS)
+    print(header)
+    for config_name in CONFIG_NAMES:
+        row = f"{config_name:<15}{summary[config_name]['total']:>8}" + "".join(
+            f"{summary[config_name][b]:>22}" for b in TAXONOMY_BUCKETS
+        )
+        print(row)
+
+
+if __name__ == "__main__":
+    run_benchmark()

src/db.py

@@ -1,4 +1,5 @@
 import os
+import time
 import psycopg2
 from dotenv import load_dotenv
 
@@ -8,7 +9,16 @@ _DATABASE_URL = os.environ["DATABASE_URL"]
 
 
 def get_connection() -> psycopg2.extensions.connection:
-    return psycopg2.connect(_DATABASE_URL)
+    last_error = None
+    for attempt in range(3):
+        try:
+            return psycopg2.connect(_DATABASE_URL)
+        except psycopg2.OperationalError as e:
+            last_error = e
+            if attempt < 2:
+                print(f"  DB connection attempt {attempt+1} failed, retrying in 5s...")
+                time.sleep(5)
+    raise last_error
 
 
 def _enable_pgvector(conn: psycopg2.extensions.connection) -> None: