bench_vs_grep.py

#!/usr/bin/env python
"""Benchmark CodeRAG's indexed search against a raw grep baseline.

This makes the headline claim measurable: a warm, pre-indexed workspace answers a query
faster (one in-process call vs many ripgrep invocations) and more accurately on conceptual
/ natural-language queries than the agentic grep loop that tools like Claude Code and Codex
fall back to. It reuses the eval harness so the accuracy numbers are directly comparable to
``coderag eval``.

    python scripts/bench_vs_grep.py \
        --watched-dir . \
        --dataset coderag/eval/datasets/coderag_self.jsonl

What it reports, both for CodeRAG (hybrid retrieval) and for grep:
  * accuracy   — recall@k / nDCG@k / MRR at the file level (via coderag.eval.evaluate)
  * latency    — mean / p50 / p95 wall-clock per query
  * context    — approximate tokens needed to surface the top-k context (CodeRAG returns
                 compact chunks; the grep baseline must read whole matched files)

The grep baseline models the agent's behaviour: extract salient terms from the query, run
ripgrep for each, rank files by match frequency — i.e. the floor that semantic search is
meant to beat. As the project's own strategy notes, grep wins on exact identifiers and
persistent edit tasks; CodeRAG's edge is conceptual queries on larger repos plus BM25 for
identifiers, with no code leaving the machine.
"""

from __future__ import annotations

import argparse
import os
import re
import subprocess  # nosec B404 — benchmarking against the ripgrep CLI is the whole point
import time
from collections import Counter
from pathlib import Path
from statistics import mean
from typing import Callable, List, Sequence

from coderag.api import CodeRAG
from coderag.config import Config
from coderag.eval import EvalCase, evaluate, load_dataset
from coderag.eval.harness import EvalResult, format_table
from coderag.types import SearchHit

# Tiny stopword set so grep searches for content terms, not glue words.
_STOP = {
    "the",
    "and",
    "for",
    "with",
    "that",
    "this",
    "from",
    "into",
    "are",
    "was",
    "use",
    "add",
    "fix",
    "when",
    "where",
    "what",
    "how",
    "does",
    "should",
    "make",
    "now",
}
_TOKEN = re.compile(r"[A-Za-z_][A-Za-z0-9_]{2,}")


def _query_terms(query: str, limit: int = 8) -> List[str]:
    """Salient search terms from a natural-language query (what an agent would grep for)."""
    seen: List[str] = []
    for tok in _TOKEN.findall(query):
        low = tok.lower()
        if low in _STOP or low in {t.lower() for t in seen}:
            continue
        seen.append(tok)
        if len(seen) >= limit:
            break
    return seen


def make_grep_search(root: Path) -> Callable[[str, int], List[SearchHit]]:
    """A grep-backed retriever with the harness's ``(query, k) -> hits`` signature."""

    def search(query: str, k: int) -> List[SearchHit]:
        terms = _query_terms(query)
        if not terms:
            return []
        counts: Counter = Counter()
        for term in terms:
            try:
                proc = subprocess.run(  # nosec B603,B607 — fixed argv, no shell
                    [
                        "rg",
                        "--count-matches",
                        "--no-messages",
                        "-i",
                        "-e",
                        term,
                        str(root),
                    ],
                    capture_output=True,
                    text=True,
                    timeout=30,
                )
            except (FileNotFoundError, subprocess.TimeoutExpired):
                continue
            for line in proc.stdout.splitlines():
                path, _, num = line.rpartition(":")
                if not path:
                    continue
                try:
                    counts[path] += int(num)
                except ValueError:
                    # Not a "path:count" line (e.g. a path containing ':', or rg's
                    # summary output) — skip it rather than fail the whole query.
                    continue
        hits: List[SearchHit] = []
        for abs_path, score in counts.most_common(k):
            rel = os.path.relpath(abs_path, root)
            hits.append(
                SearchHit(
                    chunk_id=0,
                    path=Path(rel).as_posix(),
                    symbol=None,
                    kind="window",
                    language="",
                    start_line=1,
                    end_line=1,
                    text="",
                    score=float(score),
                    similarity=0.0,
                )
            )
        return hits

    return search


def _timed(
    fn: Callable[[str, int], List[SearchHit]], sink: List[float]
) -> Callable[[str, int], List[SearchHit]]:
    def wrapped(query: str, k: int) -> List[SearchHit]:
        start = time.perf_counter()
        try:
            return fn(query, k)
        finally:
            sink.append(time.perf_counter() - start)

    return wrapped


def _percentile(values: Sequence[float], pct: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    idx = min(len(ordered) - 1, int(round((pct / 100.0) * (len(ordered) - 1))))
    return ordered[idx]


def _fmt_ms(values: Sequence[float]) -> str:
    if not values:
        return "n/a"
    return (
        f"mean {mean(values) * 1000:7.1f}  "
        f"p50 {_percentile(values, 50) * 1000:7.1f}  "
        f"p95 {_percentile(values, 95) * 1000:7.1f}  (ms)"
    )


def _file_chars(path: Path) -> int:
    try:
        return len(path.read_text(encoding="utf-8", errors="replace"))
    except OSError:
        return 0


def _context_tokens(
    cr: CodeRAG,
    grep_search: Callable[[str, int], List[SearchHit]],
    cases: Sequence[EvalCase],
    root: Path,
    top_k: int,
) -> tuple[int, int]:
    """Approximate tokens (~chars/4) to surface top-k context for each retriever.

    CodeRAG returns the matched chunks; the grep baseline must read whole matched files —
    which is the token-cost argument in favour of indexed retrieval.
    """
    cr_tokens = grep_tokens = 0
    for case in cases:
        cr_tokens += sum(len(h.text) for h in cr.search(case.query, top_k)) // 4
        grep_tokens += (
            sum(_file_chars(root / h.path) for h in grep_search(case.query, top_k)) // 4
        )
    return cr_tokens, grep_tokens


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--watched-dir", default=".", help="Repo/directory to search.")
    ap.add_argument(
        "--dataset",
        default="coderag/eval/datasets/coderag_self.jsonl",
        help="JSONL query -> relevant_files dataset (file level).",
    )
    ap.add_argument(
        "--store-dir", default=None, help="Index location (default ./.coderag)."
    )
    ap.add_argument("--model", default="BAAI/bge-small-en-v1.5")
    ap.add_argument("--ks", default="1,5,10")
    ap.add_argument(
        "--top-k", type=int, default=10, help="k for latency/token sampling."
    )
    ap.add_argument("--no-index", action="store_true", help="Reuse the existing index.")
    args = ap.parse_args()

    root = Path(args.watched_dir).expanduser().resolve()
    ks = tuple(int(k) for k in args.ks.split(","))
    cases = load_dataset(args.dataset)
    if not cases:
        raise SystemExit(f"No eval cases in {args.dataset}.")

    cfg = Config.from_env(
        provider="fastembed",
        model=args.model,
        watched_dir=root,
        store_dir=Path(args.store_dir).expanduser()
        if args.store_dir
        else root / ".coderag",
    )
    cr = CodeRAG(cfg)
    if not args.no_index:
        stats = cr.index()
        print(f"Indexed {stats.total_files} files / {stats.total_chunks} chunks.\n")

    grep_search = make_grep_search(root)

    cr_times: List[float] = []
    grep_times: List[float] = []
    results: List[EvalResult] = [
        evaluate(
            _timed(cr.search, cr_times),
            cases,
            label="coderag (hybrid)",
            ks=ks,
            level="file",
        ),
        evaluate(
            _timed(grep_search, grep_times), cases, label="grep", ks=ks, level="file"
        ),
    ]

    cr_tok, grep_tok = _context_tokens(cr, grep_search, cases, root, args.top_k)

    print(f"Accuracy ({len(cases)} cases, file level)\n")
    print(format_table(results))
    print("\nLatency per query")
    print(f"  coderag (1 warm call)   : {_fmt_ms(cr_times)}")
    print(
        f"  grep ({len(_query_terms(cases[0].query)) or 'n'} rg calls/query) : {_fmt_ms(grep_times)}"
    )
    print(f"\nApprox context tokens for top-{args.top_k} (sum over cases)")
    print(f"  coderag (compact chunks): {cr_tok:>9,}")
    print(f"  grep (read whole files) : {grep_tok:>9,}")
    if cr_tok:
        print(f"  -> grep needs ~{grep_tok / max(cr_tok, 1):.1f}x the context tokens")
    cr.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())