"""LLM-as-a-judge benchmark for the TraceRAG pipeline. python -m scripts.benchmark --db memory.lbug --out results.csv """ from __future__ import annotations import argparse import csv import logging import os import sys import time from dataclasses import dataclass from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from tracerag import config # noqa: E402 from tracerag.db import TraceDB # noqa: E402 from tracerag.router import TraceRouter # noqa: E402 from tracerag.llm import make_client # noqa: E402 logger = logging.getLogger("tracerag.benchmark") JUDGE_CONTEXT_CHARS = int(os.getenv("TRACERAG_JUDGE_CHARS", "100000")) JUDGE_SLEEP_SEC = float(os.getenv("TRACERAG_JUDGE_SLEEP", "1")) @dataclass(frozen=True) class TestQuery: query: str category: str TEST_SET: list[TestQuery] = [ # semantic / conceptual TestQuery("Explain the ShopFlow commerce platform architecture.", "semantic"), TestQuery("What is the responsibility of the PaymentService?", "semantic"), TestQuery("How does the AuthLayer handle login and JWT issuance?", "semantic"), TestQuery("What is the role of the InventoryService?", "semantic"), TestQuery("Summarize the ShopFlow microservices and what each one does.", "semantic"), # relational / multi-hop TestQuery("What was incident INC-4471 about?", "relational"), TestQuery("Which component was affected in incident INC-4480?", "relational"), TestQuery("Who is the tech lead and owning team for PaymentService?", "relational"), TestQuery("What does payment-service depend on?", "relational"), TestQuery("What is related to PR #847?", "relational"), ] def approx_tokens(text: str) -> int: """Cheap token estimate: words * 1.3.""" return int(len(text.split()) * 1.3) def baseline_context(db: TraceDB, node_ids: list[str]) -> str: """Pure-vector control: concatenate the unique chunk texts of the top-k hits.""" docs_by_entity = db.documents_for_entities(node_ids) seen, parts = set(), [] for nid in node_ids: for d in docs_by_entity.get(nid, []): text = (d.get("content") or "").strip() if text and text not in seen: seen.add(text) parts.append(text) return "\n\n".join(parts) def _safe_vector_search(db: TraceDB, embedding: list[float], k: int) -> list[dict]: try: return db.vector_search(embedding, k=k) except Exception as exc: # noqa: BLE001 logger.debug("baseline vector_search returned nothing (%s)", exc) return [] def judge_sufficient(client, query: str, context: str) -> tuple[int, str]: """LLM-as-judge: is the context relevant to answering the query? -> (1/0, raw).""" prompt = ( f"Does the provided context contain facts relevant to answering the query? " f"Reply YES if it includes information that helps answer it (even partially, " f"or via connected relationships); reply NO only if the context is completely " f"unrelated or off-topic.\n\n" f"Note: The context contains SYSTEM TRACES (hard factual relationships " f"between engineering entities) followed by standard DOCUMENT CHUNKS. Treat " f"the traces as absolute facts.\n\n" f"Output YES or NO as the very first word of your reply.\n\n" f"Query: {query}\n\n" f"Context:\n{context if context.strip() else '(no context retrieved)'}" ) try: resp = client.chat.completions.create( model=config.OPENROUTER_MODEL, messages=[{"role": "user", "content": prompt}], temperature=0, ) answer = resp.choices[0].message.content.strip().upper() except Exception as exc: # noqa: BLE001 logger.warning("Judge LLM failed (%s); scoring as 0.", exc) return 0, f"ERROR: {exc}" return (1 if answer.startswith("YES") else 0), answer def run(db_path: Path, out_path: Path, k: int) -> list[dict]: db = TraceDB(db_path) db.init_schema() router = TraceRouter(db) client = make_client() rows: list[dict] = [] try: for tq in TEST_SET: # hybrid (router) arm resp = router.route(tq.query, top_k=k) context = router.build_context(resp.results) intent = resp.trace_log.get("intent", {}) tokens = approx_tokens(context) # pure-vector baseline arm embedding = router.embed_query(tq.query) baseline_ids = [h["id"] for h in _safe_vector_search(db, embedding, k) if h.get("id")] baseline_tokens = approx_tokens(baseline_context(db, baseline_ids)) reduction = ( (1 - tokens / baseline_tokens) * 100 if baseline_tokens else 0.0 ) verdict, raw = judge_sufficient( client, tq.query, context[:JUDGE_CONTEXT_CHARS] ) time.sleep(JUDGE_SLEEP_SEC) rows.append({ "query": tq.query, "category": tq.category, "intent_type": intent.get("type"), "alpha": intent.get("alpha"), "beta": intent.get("beta"), "num_docs": len(resp.results), "context_tokens": tokens, "baseline_tokens": baseline_tokens, "token_reduction_pct": round(reduction, 1), "judge": verdict, "judge_raw": raw, }) logger.info("[%-10s] judge=%d tokens=%4d baseline=%4d reduction=%5.1f%% %s", tq.category, verdict, tokens, baseline_tokens, reduction, tq.query) finally: db.close() _write_csv(rows, out_path) return rows def _write_csv(rows: list[dict], out_path: Path) -> None: fields = ["query", "category", "intent_type", "alpha", "beta", "num_docs", "context_tokens", "baseline_tokens", "token_reduction_pct", "judge", "judge_raw"] with out_path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fields) writer.writeheader() writer.writerows(rows) logger.info("Wrote %d rows -> %s", len(rows), out_path) def print_summary(rows: list[dict]) -> None: header = (f"{'Category':<12}{'Queries':>9}{'Hybrid Tok':>12}" f"{'Baseline Tok':>14}{'Reduction':>11}{'Accuracy':>11}") width = len(header) def line(label: str, group: list[dict]) -> None: n = len(group) if not n: return hyb = sum(r["context_tokens"] for r in group) / n base = sum(r["baseline_tokens"] for r in group) / n red = sum(r["token_reduction_pct"] for r in group) / n acc = sum(r["judge"] for r in group) / n print(f"{label:<12}{n:>9}{hyb:>12.1f}{base:>14.1f}{red:>10.1f}%{acc:>11.1%}") print("\n" + "=" * width) print(header) print("-" * width) for category in ("semantic", "relational"): line(category, [r for r in rows if r["category"] == category]) print("-" * width) line("OVERALL", rows) print("=" * width + "\n") def parse_args(argv: list[str] | None = None) -> argparse.Namespace: p = argparse.ArgumentParser(description="TraceRAG LLM-as-judge benchmark.") p.add_argument("--db", type=Path, default=config.DB_PATH) p.add_argument("--out", type=Path, default=config.RESULTS_CSV) p.add_argument("--k", type=int, default=config.TOP_K_VECTOR, help="Top-k entities retrieved per query.") p.add_argument("-v", "--verbose", action="store_true") return p.parse_args(argv) def main(argv: list[str] | None = None) -> int: args = parse_args(argv) logging.basicConfig( level=logging.DEBUG if args.verbose else logging.INFO, format="%(asctime)s %(levelname)-7s %(name)s %(message)s", ) for noisy in ("httpx", "httpcore", "openai", "sentence_transformers"): logging.getLogger(noisy).setLevel(logging.WARNING) rows = run(args.db, args.out, args.k) print_summary(rows) return 0 if __name__ == "__main__": raise SystemExit(main())