benchmark.py

"""
Benchmark runner -- runs all 18 queries against a saved pickle run, saves results.

Usage:
  python benchmark.py --run baseline                    # Run all queries against a pickle
  python benchmark.py --run baseline --query-id 3       # Run a single query
  python benchmark.py --list-results                    # List all saved benchmark results
  python benchmark.py --compare                         # Compare results across runs

Results saved to: benchmark_results/<embed_model>/<run_name>.json
"""

import argparse
import json
import math
import os
import pickle
import subprocess
import sys
import time
from datetime import datetime

# ============================================================================
# LOAD .env
# ============================================================================
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_env_path = os.path.join(SCRIPT_DIR, ".env")
if os.path.isfile(_env_path):
    with open(_env_path) as _f:
        for _line in _f:
            _line = _line.strip()
            if not _line or _line.startswith("#"):
                continue
            key, _, val = _line.partition("=")
            os.environ.setdefault(key.strip(), val.strip().strip("'\""))

# ============================================================================
# CONFIG
# ============================================================================
RUNS_DIR = os.path.join(SCRIPT_DIR, "runs")
RESULTS_BASE_DIR = os.path.join(SCRIPT_DIR, "benchmark_results")
QUERIES_PATH = os.path.join(SCRIPT_DIR, "benchmark_queries.json")
os.makedirs(RESULTS_BASE_DIR, exist_ok=True)

C = {
    "RESET": "\033[0m",
    "DIM": "\033[2m",
    "BOLD": "\033[1m",
    "CYAN": "\033[36m",
    "GREEN": "\033[32m",
    "YELLOW": "\033[33m",
    "MAGENTA": "\033[35m",
    "RED": "\033[31m",
}

BOILERPLATE_PATTERNS = [
    "equal opportunity", "sexual orientation", "national origin", "gender identity",
    "without regard to", "race, color", "religion, sex", "marital status",
    "protected veteran", "reasonable accommodation", "genetic information", "EEO",
    "affirmative action", "we do not discriminate", "e-verify", "Americans with Disabilities",
]


def log(tag, msg):
    ts = datetime.now().strftime("%H:%M:%S")
    colors = {"BENCH": C["CYAN"], "QUERY": C["YELLOW"], "INFO": C["DIM"], "ERROR": C["RED"], "RESULT": C["GREEN"]}
    color = colors.get(tag, C["RESET"])
    print(f"{C['DIM']}{ts}{C['RESET']} {color}[{tag}]{C['RESET']}  {msg}")


# ============================================================================
# LOAD DATA
# ============================================================================

def load_queries():
    with open(QUERIES_PATH) as f:
        data = json.load(f)
    return data["queries"]


def load_run(name):
    path = os.path.join(RUNS_DIR, f"{name}.pkl")
    if not os.path.exists(path):
        log("ERROR", f"Run not found: runs/{name}.pkl")
        available = [f[:-4] for f in os.listdir(RUNS_DIR) if f.endswith(".pkl")]
        if available:
            log("INFO", f"Available: {', '.join(available)}")
        sys.exit(1)
    with open(path, "rb") as f:
        data = pickle.load(f)
    pkl_size_mb = os.path.getsize(path) / (1024 * 1024)
    log("BENCH", f"Loaded run: {name} | {data['num_vectors']} vectors | {data['dims']} dims | {pkl_size_mb:.1f} MB")
    return data, pkl_size_mb


# ============================================================================
# COSINE SIMILARITY
# ============================================================================

def cosine_sim(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(x * x for x in b))
    return dot / (na * nb) if na and nb else 0


def is_boilerplate(text):
    text_lower = text.lower()
    count = sum(1 for p in BOILERPLATE_PATTERNS if p.lower() in text_lower)
    return count >= 3


# ============================================================================
# LLM via opencode (anthropic/claude-opus-4-6)
# ============================================================================
LLM_MODEL = "anthropic/claude-opus-4-6"
LLM_TIMEOUT = 120

def call_llm(prompt):
    """Call opencode with a prompt via stdin, return the text response."""
    proc = subprocess.run(
        ["opencode", "run", "--model", LLM_MODEL, "--format", "json"],
        input=prompt,
        capture_output=True,
        text=True,
        timeout=LLM_TIMEOUT,
    )
    if proc.returncode != 0:
        raise RuntimeError(f"opencode exited {proc.returncode}: {proc.stderr[:500]}")

    # Parse JSON events, extract text parts
    content = ""
    for line in proc.stdout.strip().split("\n"):
        try:
            event = json.loads(line)
            if event.get("type") == "text" and event.get("part", {}).get("text"):
                content += event["part"]["text"]
        except json.JSONDecodeError:
            pass

    if not content:
        raise RuntimeError(f"opencode returned no text content. stdout: {proc.stdout[:500]}")
    return content


# ============================================================================
# RUN SINGLE QUERY
# ============================================================================

def run_query(query_obj, vectors, texts, metadatas, embeddings_model, max_chunks=200):
    query_text = query_obj["query"]
    result = {
        "query_id": query_obj["id"],
        "category": query_obj["category"],
        "query": query_text,
        "tests": query_obj["tests"],
    }

    # Embed query
    q_start = time.time()
    query_vector = embeddings_model.embed_query(query_text)
    embed_latency = time.time() - q_start

    # Cosine similarity against all chunks
    sim_start = time.time()
    similarities = []
    for i, vec in enumerate(vectors):
        sim = cosine_sim(query_vector, vec)
        similarities.append((i, sim))
    similarities.sort(key=lambda x: x[1], reverse=True)
    search_latency = time.time() - sim_start

    # Top-N selection: take the highest similarity chunks, capped at max_chunks
    selected = similarities[:max_chunks]
    was_capped = len(similarities) > max_chunks

    # Collect detailed info for top-20 (for the results JSON)
    top_20 = []
    unique_sources_5 = set()
    unique_sources_10 = set()
    unique_sources_20 = set()
    boilerplate_in_top_20 = 0

    for rank, (idx, sim) in enumerate(similarities[:20]):
        src = metadatas[idx].get("source", "?")
        text = texts[idx]
        bp = is_boilerplate(text)
        if bp:
            boilerplate_in_top_20 += 1

        entry = {
            "rank": rank + 1,
            "chunk_index": idx,
            "source": src,
            "similarity": round(sim, 6),
            "char_count": len(text),
            "is_boilerplate": bp,
            "preview": text[:200].replace("\n", " "),
        }
        top_20.append(entry)

        if rank < 5:
            unique_sources_5.add(src)
        if rank < 10:
            unique_sources_10.add(src)
        unique_sources_20.add(src)

    # Similarity score stats (based on all selected chunks, not just top-20)
    selected_sims = [sim for _, sim in selected]
    sim_stats = {
        "top_1": selected_sims[0] if selected_sims else 0,
        "min": selected_sims[-1] if selected_sims else 0,
        "spread": round(selected_sims[0] - selected_sims[-1], 6) if selected_sims else 0,
        "median": selected_sims[len(selected_sims) // 2] if selected_sims else 0,
    }

    # Build context from top-N chunks
    context_chunks = [texts[idx] for idx, sim in selected]
    context = "\n\n---\n\n".join(context_chunks)
    unique_context_sources = len({metadatas[idx].get("source", "?") for idx, sim in selected})
    boilerplate_in_context = sum(1 for idx, sim in selected if is_boilerplate(texts[idx]))
    min_sim = selected[-1][1] if selected else 0

    log("QUERY", f"  top-{len(selected)} chunks (sim {selected[0][1]:.4f} -> {min_sim:.4f}){' (CAPPED)' if was_capped else ''} | {unique_context_sources} sources | {boilerplate_in_context} boilerplate")

    augmented_prompt = f"""Based on the following context from job listings, answer the question.
You are receiving {len(context_chunks)} relevant chunks from a database of ~6500 job listings.

CONTEXT:
{context}

QUESTION: {query_text}

Answer based only on the provided context. If the context doesn't contain enough information, say so."""

    llm_start = time.time()
    response_text = call_llm(augmented_prompt)
    llm_latency = time.time() - llm_start

    prompt_chars = len(augmented_prompt)
    prompt_est_tokens = prompt_chars // 4
    context_chars = len(context)

    result["top_20"] = top_20
    result["similarity_stats"] = sim_stats
    result["retrieval_stats"] = {
        "max_chunks": max_chunks,
        "chunks_used": len(selected),
        "was_capped": was_capped,
        "unique_sources_in_context": unique_context_sources,
        "boilerplate_in_context": boilerplate_in_context,
    }
    result["unique_sources"] = {
        "top_5": len(unique_sources_5),
        "top_10": len(unique_sources_10),
        "top_20": len(unique_sources_20),
    }
    result["boilerplate_in_top_20"] = boilerplate_in_top_20
    result["llm_prompt"] = {
        "total_chars": prompt_chars,
        "est_tokens": prompt_est_tokens,
        "context_chars": context_chars,
        "query_chars": len(query_text),
        "num_context_chunks": len(context_chunks),
    }
    result["llm_response"] = response_text
    result["llm_response_chars"] = len(response_text)
    result["latency"] = {
        "embed_query_ms": round(embed_latency * 1000, 1),
        "cosine_search_ms": round(search_latency * 1000, 1),
        "llm_generation_ms": round(llm_latency * 1000, 1),
        "total_ms": round((embed_latency + search_latency + llm_latency) * 1000, 1),
    }

    return result


# ============================================================================
# DISPLAY RESULT
# ============================================================================

def display_result(r):
    log("QUERY", f"Q{r['query_id']}: \"{r['query']}\"")
    rs = r["retrieval_stats"]
    ss = r["similarity_stats"]
    log("RESULT", f"Context: {rs['chunks_used']} chunks | {rs['unique_sources_in_context']} sources | {rs['boilerplate_in_context']} boilerplate")
    log("RESULT", f"Sims: top1={ss['top_1']:.4f} floor={ss['min']:.4f} spread={ss['spread']:.4f}")
    prompt = r["llm_prompt"]
    log("RESULT", f"Prompt: ~{prompt['est_tokens']} tok | context={prompt['context_chars']} chars ({prompt['num_context_chunks']} chunks)")
    lat = r["latency"]
    log("RESULT", f"Latency: embed={lat['embed_query_ms']:.0f}ms search={lat['cosine_search_ms']:.0f}ms llm={lat['llm_generation_ms']:.0f}ms total={lat['total_ms']:.0f}ms")
    log("RESULT", f"LLM ({r['llm_response_chars']} chars): {r['llm_response'][:150]}...")
    print()


# ============================================================================
# COMPARE RESULTS
# ============================================================================

def compare_results(results_dir):
    files = sorted([f for f in os.listdir(results_dir) if f.endswith(".json")])
    if not files:
        log("INFO", "No benchmark results yet.")
        return

    all_results = {}
    for f in files:
        with open(os.path.join(results_dir, f)) as fh:
            data = json.load(fh)
        all_results[data["run_name"]] = data

    # Header
    run_names = list(all_results.keys())
    header = f"{'Query':<60s}"
    for name in run_names:
        header += f" | {name:<20s}"
    print(f"\n{C['BOLD']}{header}{C['RESET']}")
    print("-" * len(header))

    # Per-query comparison (top-1 similarity + spread)
    first_run = all_results[run_names[0]]
    for qr in first_run["query_results"]:
        qid = qr["query_id"]
        line = f"Q{qid}: {qr['query'][:55]:<58s}"
        for name in run_names:
            run_data = all_results[name]
            match = next((r for r in run_data["query_results"] if r["query_id"] == qid), None)
            if match:
                top1 = match["similarity_stats"]["top_1"]
                spread = match["similarity_stats"]["spread"]
                line += f" | {top1:.4f} ({spread:.3f}){' ':>5s}"
            else:
                line += f" | {'N/A':>20s}"
        print(line)

    # Summary row
    print("-" * len(header))
    line = f"{'AVERAGE':.<60s}"
    for name in run_names:
        run_data = all_results[name]
        top1s = [r["similarity_stats"]["top_1"] for r in run_data["query_results"]]
        spreads = [r["similarity_stats"]["spread"] for r in run_data["query_results"]]
        avg_top1 = sum(top1s) / len(top1s) if top1s else 0
        avg_spread = sum(spreads) / len(spreads) if spreads else 0
        line += f" | {avg_top1:.4f} ({avg_spread:.3f}){' ':>5s}"
    print(f"{C['BOLD']}{line}{C['RESET']}")

    # Efficiency summary
    print(f"\n{C['BOLD']}Efficiency:{C['RESET']}")
    for name in run_names:
        r = all_results[name]
        meta = r["run_metadata"]
        avg_lat = sum(q["latency"]["total_ms"] for q in r["query_results"]) / len(r["query_results"]) if r["query_results"] else 0
        print(f"  {name}: cost=${meta.get('est_cost_usd', 0):.4f} | embed={meta.get('embed_seconds', 0):.0f}s | dims={meta.get('dims', '?')} | pkl={meta.get('pkl_size_mb', '?'):.1f}MB | avg_query={avg_lat:.0f}ms")
    print()


# ============================================================================
# MAIN
# ============================================================================

def main():
    parser = argparse.ArgumentParser(description="Benchmark runner for RAG pipeline")
    parser.add_argument("--run", type=str, help="Name of the pickle run to benchmark")
    parser.add_argument("--query-id", type=int, help="Run only this query ID (default: all)")
    parser.add_argument("--max-chunks", type=int, default=200, help="Max chunks sent to LLM per query, ranked by similarity (default: 200)")
    parser.add_argument("--embed-model", type=str, default="gemini", help="Embedding model name, used as subfolder in benchmark_results/ (default: gemini)")
    parser.add_argument("--list-results", action="store_true", help="List all saved benchmark results")
    parser.add_argument("--compare", action="store_true", help="Compare results across runs")
    args = parser.parse_args()

    results_dir = os.path.join(RESULTS_BASE_DIR, args.embed_model)
    os.makedirs(results_dir, exist_ok=True)

    if args.list_results:
        files = sorted([f for f in os.listdir(results_dir) if f.endswith(".json")])
        if not files:
            log("INFO", "No benchmark results yet.")
        for f in files:
            with open(os.path.join(results_dir, f)) as fh:
                data = json.load(fh)
            n = len(data["query_results"])
            log("INFO", f"{data['run_name']}: {n} queries | {data['run_metadata'].get('num_vectors', '?')} vectors | saved {data['timestamp']}")
        return

    if args.compare:
        compare_results(results_dir)
        return

    if not args.run:
        parser.error("--run is required (or use --list-results / --compare)")

    # Load run and queries
    run_data, pkl_size_mb = load_run(args.run)
    vectors = run_data["vectors"]
    texts = run_data["texts"]
    metadatas = run_data["metadatas"]
    queries = load_queries()

    if args.query_id:
        queries = [q for q in queries if q["id"] == args.query_id]
        if not queries:
            log("ERROR", f"Query ID {args.query_id} not found in benchmark_queries.json")
            sys.exit(1)

    # Initialize embedding model (must match the model that produced the pickle vectors)
    from embed_model import get_embeddings, get_model_info

    # Short name -> registry name mapping so --embed-model gemini still works
    EMBED_SHORT_NAMES = {
        "gemini": "gemini-embedding-001",
        "openai_small": "text-embedding-3-small",
        "openai_large": "text-embedding-3-large",
        "e5": "e5-large-instruct",
        "bge": "bge-base-en-v1.5",
        "qwen3": "qwen3-embedding-0.6b",
        "gemma": "embeddinggemma-300m",
    }
    embed_model_name = EMBED_SHORT_NAMES.get(args.embed_model, args.embed_model)
    model_info = get_model_info(embed_model_name)
    if not model_info:
        log("ERROR", f"Unknown embedding model: {embed_model_name}")
        log("INFO", f"Short names: {', '.join(EMBED_SHORT_NAMES.keys())}")
        sys.exit(1)

    embeddings_model = get_embeddings(embed_model_name)

    log("BENCH", f"Running {len(queries)} queries against run '{args.run}' | top-{args.max_chunks} chunks per query")
    log("BENCH", f"Embedding: {embed_model_name} ({model_info['dims']}d, {model_info['provider']}) | LLM: {LLM_MODEL} (via opencode)")

    out_path = os.path.join(results_dir, f"{args.run}.json")

    # Build the output shell -- query_results grows after each query and gets flushed to disk
    run_metadata = {
        "num_vectors": run_data.get("num_vectors", len(vectors)),
        "dims": run_data.get("dims", len(vectors[0]) if vectors else 0),
        "est_cost_usd": run_data.get("est_cost_usd", 0),
        "embed_seconds": run_data.get("embed_seconds", 0),
        "pkl_size_mb": round(pkl_size_mb, 1),
        "params": run_data.get("params", {}),
        "embed_model": embed_model_name,
        "benchmark_params": {
            "max_chunks": args.max_chunks,
        },
    }

    # Load existing partial results if resuming after a crash
    all_query_results = []
    completed_ids = set()
    if os.path.exists(out_path):
        try:
            with open(out_path) as f:
                existing = json.load(f)
            all_query_results = existing.get("query_results", [])
            completed_ids = {r["query_id"] for r in all_query_results}
            if completed_ids:
                log("BENCH", f"Resuming -- {len(completed_ids)} queries already done, skipping them")
        except (json.JSONDecodeError, KeyError):
            pass

    total_start = time.time()

    for i, q in enumerate(queries):
        if q["id"] in completed_ids:
            log("BENCH", f"[{i+1}/{len(queries)}] Q{q['id']}: SKIP (already done)")
            continue

        log("BENCH", f"[{i+1}/{len(queries)}] Q{q['id']}: {q['query'][:70]}")
        result = run_query(q, vectors, texts, metadatas, embeddings_model, max_chunks=args.max_chunks)
        display_result(result)
        all_query_results.append(result)

        # Flush to disk after every query
        total_duration = time.time() - total_start
        output = {
            "run_name": args.run,
            "timestamp": datetime.now().isoformat(),
            "total_duration_s": round(total_duration, 1),
            "num_queries": len(all_query_results),
            "run_metadata": run_metadata,
            "query_results": all_query_results,
        }
        with open(out_path, "w") as f:
            json.dump(output, f, indent=2)
        log("INFO", f"Saved {len(all_query_results)}/{len(queries)} -> benchmark_results/{args.embed_model}/{args.run}.json")

    total_duration = time.time() - total_start
    log("BENCH", f"Done. {len(all_query_results)} queries in {total_duration:.1f}s -> benchmark_results/{args.embed_model}/{args.run}.json")

    # Summary
    avg_top1 = sum(r["similarity_stats"]["top_1"] for r in all_query_results) / len(all_query_results)
    avg_spread = sum(r["similarity_stats"]["spread"] for r in all_query_results) / len(all_query_results)
    avg_lat = sum(r["latency"]["total_ms"] for r in all_query_results) / len(all_query_results)
    log("BENCH", f"Avg top1: {avg_top1:.4f} | Avg spread: {avg_spread:.4f} | Avg latency: {avg_lat:.0f}ms")


# Graceful shutdown
import signal

def _shutdown(signum, frame):
    name = signal.Signals(signum).name
    sys.stdout.write("\033[?25h")
    sys.stdout.flush()
    log("INFO", f"Received {name}, exiting.")
    sys.exit(0)

signal.signal(signal.SIGINT, _shutdown)
signal.signal(signal.SIGTERM, _shutdown)

if __name__ == "__main__":
    main()