unsloth/studio/bench/rag_studio_bench.py

#!/usr/bin/env python3
"""Benchmark a running Studio's RAG over its real HTTP API: indexing latency,
scaling, and retrieval accuracy. Same script drives baseline and improved.

Usage:
  python rag_studio_bench.py --base http://127.0.0.1:8901 --label baseline \
     --password "<bootstrap>" --corpus ../data/rag_corpus --out ../logs/bench_baseline.json
"""

import argparse
import json
import time
from pathlib import Path

import httpx

API = "/api/rag"

# Gold queries: (query, filename_substring, [accepted answer phrases any-of]).
# A hit counts when a retrieved chunk is FROM the right document AND contains an
# accepted phrase, so retrieval is scored independently of the chat model.
GOLD = [
    (
        "What are the sinusoidal positional encodings based on?",
        "attention",
        ["sine", "sinusoid", "wavelength"],
    ),
    (
        "How many attention heads does the base Transformer use?",
        "attention",
        ["eight", "h = 8", "8 parallel", "8 attention", "h=8"],
    ),
    (
        "What are BERT's two pre-training objectives?",
        "bert",
        ["masked", "next sentence"],
    ),
    (
        "How many Transformer layers does BERT-large have?",
        "bert",
        ["24", "l = 24", "l=24"],
    ),
    (
        "Difference between RAG-Sequence and RAG-Token models?",
        "rag",
        ["rag-token", "rag-sequence"],
    ),
    (
        "Which retriever does RAG use to fetch passages?",
        "rag",
        ["dpr", "dense passage", "bi-encoder", "mips"],
    ),
    (
        "What does the HTTP GET method do?",
        "rfc9110",
        ["transfer a current representation", "retrieve", "selector"],
    ),
    ("Which HTTP status code means the resource was not found?", "rfc9110", ["404"]),
]


def login(c, base, user, pw):
    new = pw + "Aa1!"
    # Re-run safe: bootstrap pw works first time; after we change it, the changed
    # pw works on later runs.
    r = c.post(f"{base}/api/auth/login", json = {"username": user, "password": pw})
    if r.status_code == 401:
        r = c.post(f"{base}/api/auth/login", json = {"username": user, "password": new})
    r.raise_for_status()
    body = r.json()
    tok = body["access_token"]
    if body.get("must_change_password"):
        r2 = c.post(
            f"{base}/api/auth/change-password",
            headers = H(tok),
            json = {"current_password": pw, "new_password": new},
        )
        r2.raise_for_status()
        tok = r2.json()["access_token"]
    return tok


def H(tok):
    return {"Authorization": f"Bearer {tok}"}


def warmup(c, base, tok):
    t = time.perf_counter()
    try:
        c.post(f"{base}{API}/warmup", headers = H(tok), timeout = 600)
    except Exception as e:
        print("warmup err:", e)
    return time.perf_counter() - t


def create_kb(c, base, tok, name):
    r = c.post(
        f"{base}{API}/knowledge-bases",
        headers = H(tok),
        json = {"name": name, "mode": "text", "chunking_strategy": "standard"},
    )
    r.raise_for_status()
    return r.json()["kb_id" if "kb_id" in r.json() else "id"]


def upload(c, base, tok, kb_id, path: Path):
    with open(path, "rb") as f:
        r = c.post(
            f"{base}{API}/knowledge-bases/{kb_id}/documents",
            headers = H(tok),
            files = {"file": (path.name, f, "application/octet-stream")},
            timeout = 600,
        )
    r.raise_for_status()
    return r.json()


def wait_indexed(c, base, tok, kb_id, document_id, timeout = 600):
    t0 = time.perf_counter()
    while time.perf_counter() - t0 < timeout:
        r = c.get(f"{base}{API}/knowledge-bases/{kb_id}/documents", headers = H(tok))
        r.raise_for_status()
        for d in r.json()["documents"]:
            if d["id"] == document_id:
                if d["status"] == "completed":
                    return time.perf_counter() - t0, d["num_chunks"], "completed"
                if d["status"] == "failed":
                    return time.perf_counter() - t0, 0, "failed"
        time.sleep(0.1)
    return timeout, 0, "timeout"


def index_doc(c, base, tok, kb_id, path):
    t0 = time.perf_counter()
    up = upload(c, base, tok, kb_id, path)
    if up.get("already_indexed"):
        return {"file": path.name, "elapsed_s": 0.0, "chunks": 0, "status": "dup"}
    elapsed, chunks, status = wait_indexed(c, base, tok, kb_id, up["document_id"])
    total = time.perf_counter() - t0
    return {
        "file": path.name,
        "elapsed_s": round(total, 3),
        "chunks": chunks,
        "status": status,
        "document_id": up["document_id"],
    }


def search(c, base, tok, kb_id, query, mode, top_k = 10):
    r = c.post(
        f"{base}{API}/search",
        headers = H(tok),
        json = {"query": query, "kb_id": kb_id, "mode": mode, "top_k": top_k},
    )
    r.raise_for_status()
    return r.json()["hits"]


def score_mode(c, base, tok, kb_id, mode):
    r1 = r3 = r5 = 0
    mrr = 0.0
    lat = []
    for query, fsub, phrases in GOLD:
        t = time.perf_counter()
        hits = search(c, base, tok, kb_id, query, mode, top_k = 10)
        lat.append((time.perf_counter() - t) * 1000)
        rank = None
        for i, h in enumerate(hits):
            fn = (h.get("filename") or "").lower()
            txt = (h.get("text") or "").lower()
            if fsub in fn and any(p in txt for p in phrases):
                rank = i
                break
        if rank is not None:
            if rank == 0:
                r1 += 1
            if rank < 3:
                r3 += 1
            if rank < 5:
                r5 += 1
            mrr += 1.0 / (rank + 1)
    n = len(GOLD)
    lat.sort()
    return {
        "recall@1": round(r1 / n, 3),
        "recall@3": round(r3 / n, 3),
        "recall@5": round(r5 / n, 3),
        "mrr": round(mrr / n, 3),
        "search_ms_median": round(lat[len(lat) // 2], 2),
    }


def make_synthetic(dirpath: Path, n: int):
    dirpath.mkdir(parents = True, exist_ok = True)
    paths = []
    for i in range(n):
        p = dirpath / f"syn_{i:02d}.txt"
        body = (
            f"Synthetic document number {i}. Project codename Orbit-{i} concerns "
            f"widget {i} calibration at {100 + i} hertz. Unique token zglyph{i} marks "
            f"this file. " * 8
        )
        p.write_text(body)
        paths.append(p)
    return paths


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required = True)
    ap.add_argument("--label", required = True)
    ap.add_argument("--password", required = True)
    ap.add_argument("--username", default = "unsloth")
    ap.add_argument("--corpus", required = True)
    ap.add_argument("--out", required = True)
    ap.add_argument("--scaling-n", type = int, default = 8)
    args = ap.parse_args()

    res = {"label": args.label, "base": args.base}
    with httpx.Client(timeout = 120) as c:
        tok = login(c, args.base, args.username, args.password)
        res["warmup_s"] = round(warmup(c, args.base, tok), 2)

        # --- Real corpus: index timing (first = cold, rest = warm) + accuracy ---
        corpus = sorted(Path(args.corpus).glob("*"))
        corpus = [
            p
            for p in corpus
            if p.suffix.lower() in (".pdf", ".txt", ".md", ".html", ".htm", ".docx")
        ]
        kb = create_kb(c, args.base, tok, f"{args.label}-corpus")
        res["kb"] = kb
        res["corpus_index"] = []
        for i, p in enumerate(corpus):
            r = index_doc(c, args.base, tok, kb, p)
            r["which"] = "cold" if i == 0 else "warm"
            res["corpus_index"].append(r)
            print(
                f"[{args.label}] index {p.name}: {r['elapsed_s']}s ({r['chunks']} chunks) {r['which']}"
            )

        res["accuracy"] = {
            m: score_mode(c, args.base, tok, kb, m) for m in ("bm25", "dense", "hybrid")
        }
        for m, s in res["accuracy"].items():
            print(
                f"[{args.label}] {m}: R@1={s['recall@1']} R@5={s['recall@5']} MRR={s['mrr']}"
            )

        # --- Scaling: N small docs into one fresh KB, per-doc index time ---
        kb2 = create_kb(c, args.base, tok, f"{args.label}-scaling")
        syn = make_synthetic(Path(args.corpus).parent / "rag_synthetic", args.scaling_n)
        res["scaling"] = []
        for i, p in enumerate(syn):
            r = index_doc(c, args.base, tok, kb2, p)
            res["scaling"].append({"n": i + 1, "elapsed_s": r["elapsed_s"]})
            print(f"[{args.label}] scaling doc {i+1}/{len(syn)}: {r['elapsed_s']}s")

    Path(args.out).write_text(json.dumps(res, indent = 2))
    print(f"[{args.label}] wrote {args.out}")


if __name__ == "__main__":
    main()