diff --git a/studio/RAG_FAST_INDEXING.md b/studio/RAG_FAST_INDEXING.md index ff2de06bb3..e2a18783f4 100644 --- a/studio/RAG_FAST_INDEXING.md +++ b/studio/RAG_FAST_INDEXING.md @@ -86,4 +86,21 @@ improved path. Search latency on the improved path: 9-15 ms median (FTS5 + sqlit ### Summary Indexing a paper drops from ~23 s to ~7 s and a small document from ~18 s to ~0.12 s, with -retrieval accuracy held constant. UI walkthrough (Playwright) and the local-GGUF chat path follow. +retrieval accuracy held constant. + +### End-to-end validation in the real Studio UI (Playwright + local GGUF) + +Both Studios were driven through the actual web UI with Playwright, with the local +`unsloth/Qwen3.5-9B-GGUF` (Q4_K_M) loaded via llama-server. Flow per Studio: load model, enable the +RAG toggle, upload `bert_1810.04805.pdf` through the composer, then ask "What are BERT's two +pre-training objectives?". + +- **RAG works on both** (identical functionality): the model calls the `search_knowledge_base` + tool, retrieves from the uploaded PDF, and answers with source citations + (`[3][4] bert_1810.04805.pdf`). API-level tool-call test: 3/3 gold questions called the tool, + retrieved the correct source paper, and answered with the right fact (BERT -> MLM + NSP, + Transformer -> 8 heads, RAG -> DPR). +- **Indexing speed in the UI**: baseline 34.9 s vs improved 14.0 s for the same composer upload. + +This confirms the fast path keeps the PR's full RAG behavior (many document types, the RAG toggle, +and tool-call retrieval with citations) while indexing materially faster. diff --git a/studio/bench/studio_rag_chat.py b/studio/bench/studio_rag_chat.py new file mode 100644 index 0000000000..af101da1c8 --- /dev/null +++ b/studio/bench/studio_rag_chat.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""End-to-end RAG tool-call test against a running Studio with a GGUF loaded. +Asks gold questions grounded in an indexed KB, with rag_scope + tools enabled, +and reports whether the model called search_knowledge_base and answered correctly. +""" +import argparse +import json +import sys + +import httpx + +QUESTIONS = [ + ("What are BERT's two pre-training objectives?", ["masked", "next sentence"]), + ("How many attention heads does the base Transformer use?", ["8", "eight"]), + ("What retriever does the RAG paper use to fetch passages?", ["dpr", "dense passage", "bi-encoder"]), +] + + +def token(c, base, pw): + new = pw + "Aa1!" + r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": pw}) + if r.status_code == 401: + r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": new}) + r.raise_for_status() + b = r.json() + t = b["access_token"] + if b.get("must_change_password"): + r2 = c.post(f"{base}/api/auth/change-password", headers={"Authorization": f"Bearer {t}"}, + json={"current_password": pw, "new_password": new}) + r2.raise_for_status() + t = r2.json()["access_token"] + return t + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--base", required=True) + ap.add_argument("--password", required=True) + ap.add_argument("--kb-name-contains", default="corpus") + ap.add_argument("--out", default=None) + args = ap.parse_args() + + results = {"base": args.base, "qa": []} + with httpx.Client(timeout=300) as c: + tok = token(c, args.base, args.password) + H = {"Authorization": f"Bearer {tok}"} + kbs = c.get(f"{args.base}/api/rag/knowledge-bases", headers=H).json()["knowledge_bases"] + kb = next((k for k in kbs if args.kb_name_contains in k["name"]), kbs[0] if kbs else None) + if not kb: + sys.exit("no KB found") + kb_id = kb["id"] + print(f"KB: {kb['name']} ({kb_id})") + + for q, phrases in QUESTIONS: + payload = { + "model": "local", + "messages": [{"role": "user", "content": q + + " Use the knowledge base; cite the source."}], + "stream": False, + "enable_tools": True, + "enable_thinking": False, + "max_tokens": 400, + "temperature": 0.3, + "rag_scope": {"kb_id": kb_id, "default_top_k": 5, "min_score": 0.0, "mode": "hybrid"}, + } + # The endpoint streams SSE regardless of stream flag; parse it. + answer, called_tool, tool_query, src = "", False, "", "" + status = 0 + with c.stream("POST", f"{args.base}/api/inference/chat/completions", + headers=H, json=payload) as r: + status = r.status_code + for line in r.iter_lines(): + if not line or not line.startswith("data:"): + continue + data = line[len("data:"):].strip() + if data == "[DONE]": + break + try: + ev = json.loads(data) + except Exception: + continue + if ev.get("type") == "tool_start" and ev.get("tool_name") == "search_knowledge_base": + called_tool = True + tool_query = (ev.get("arguments") or {}).get("query", "") + elif ev.get("type") == "tool_end": + res = ev.get("result", "") + m = res.split('source="') + if len(m) > 1: + src = m[1].split('"')[0] + elif ev.get("choices"): + delta = ev["choices"][0].get("delta", {}) + if delta.get("content"): + answer += delta["content"] + hit = any(p in answer.lower() for p in phrases) + print(f"\nQ: {q}\n status={status} tool_called={called_tool} src={src} answer_has_fact={hit}") + print(f" tool_query={tool_query!r}") + print(f" A: {answer[:300].replace(chr(10),' ')}") + results["qa"].append({"q": q, "status": status, "tool_called": called_tool, + "retrieved_source": src, "tool_query": tool_query, + "answer_has_fact": hit, "answer": answer[:600]}) + + if args.out: + with open(args.out, "w") as f: + json.dump(results, f, indent=2) + print(f"\nwrote {args.out}") + n = len(results["qa"]) + print(f"\nSUMMARY: {sum(r['answer_has_fact'] for r in results['qa'])}/{n} answers contain the fact, " + f"{sum(r['tool_called'] for r in results['qa'])}/{n} called the RAG tool") + + +if __name__ == "__main__": + main() diff --git a/studio/bench/studio_rag_composer.py b/studio/bench/studio_rag_composer.py new file mode 100644 index 0000000000..1a7c04fc46 --- /dev/null +++ b/studio/bench/studio_rag_composer.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Full RAG-in-chat visual via Playwright: with a GGUF model already loaded, +enable the RAG toggle, upload a PDF through the composer, ask a grounded +question, and capture the answer + sources. Records screenshots + video. + +Usage: + python studio_rag_composer.py --base http://127.0.0.1:8912 --label improved \ + --password StudioBench2026! --doc ../data/rag_corpus/bert_1810.04805.pdf \ + --question "What are BERT's two pre-training objectives?" --out ../outputs/composer_improved +""" +import argparse +import asyncio +import json +import time +from pathlib import Path + +import httpx +from playwright.async_api import async_playwright + + +def get_token(base, pw): + new = pw + "Aa1!" + with httpx.Client(timeout=30) as c: + r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": pw}) + if r.status_code == 401: + r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": new}) + r.raise_for_status() + b = r.json() + t = b["access_token"] + if b.get("must_change_password"): + r2 = c.post(f"{base}/api/auth/change-password", headers={"Authorization": f"Bearer {t}"}, + json={"current_password": pw, "new_password": new}) + r2.raise_for_status() + t, b = r2.json()["access_token"], r2.json() + return t, b.get("refresh_token", "") + + +def init_script(tok, refresh): + seed = {"unsloth_auth_token": tok, "unsloth_refresh_token": refresh} + return f"""(() => {{ const s={json.dumps(seed)}; + for (const k of Object.keys(s)) {{ try {{ localStorage.setItem(k, s[k]); }} catch(e){{}} }} }})();""" + + +async def run(args): + out = Path(args.out) + (out / "video").mkdir(parents=True, exist_ok=True) + tok, refresh = get_token(args.base, args.password) + result = {"label": args.label, "doc": Path(args.doc).name, "question": args.question} + + async with async_playwright() as p: + browser = await p.chromium.launch(headless=True) + ctx = await browser.new_context(viewport={"width": 1440, "height": 900}, + record_video_dir=str(out / "video"), + record_video_size={"width": 1440, "height": 900}) + await ctx.add_init_script(init_script(tok, refresh)) + page = await ctx.new_page() + await page.goto(f"{args.base}/chat", wait_until="domcontentloaded") + await page.locator("form:has(textarea) textarea").first.wait_for(state="visible", timeout=30000) + await page.wait_for_timeout(1500) + await page.screenshot(path=str(out / "01_chat.png")) + + # Enable the RAG toggle (the pill is active once a model is loaded). + for lbl in ("Enable RAG", "Disable RAG"): + btn = page.locator(f'button[aria-label="{lbl}"]').first + if await btn.count(): + if lbl == "Enable RAG": + await btn.click() + break + await page.wait_for_timeout(700) + await page.screenshot(path=str(out / "02_rag_on.png")) + + # Upload a document through the composer attach control. + file_input = page.locator('input[type="file"][accept*=".pdf"]').first + await file_input.wait_for(state="attached", timeout=10000) + t0 = time.perf_counter() + await file_input.set_input_files(args.doc) + ready = page.get_by_text("Ready", exact=True).first + toast = page.get_by_text("RAG index ready", exact=False).first + indexed = False + while time.perf_counter() - t0 < 180: + for sig in (toast, ready): + try: + if await sig.is_visible(): + indexed = True + break + except Exception: + pass + if indexed: + break + await page.wait_for_timeout(250) + result["index_seconds"] = round(time.perf_counter() - t0, 2) + result["indexed"] = indexed + await page.screenshot(path=str(out / "03_indexed.png")) + + # Ask a grounded question. + box = page.locator("form:has(textarea) textarea").first + await box.click() + await box.fill(args.question) + await box.press("Enter") + stop = page.locator('button[aria-label="Stop generating"], button:has-text("Stop")').first + try: + await stop.wait_for(state="visible", timeout=30000) + except Exception: + pass + try: + await stop.wait_for(state="hidden", timeout=180000) + except Exception: + pass + await page.wait_for_timeout(1000) + await page.screenshot(path=str(out / "04_answer.png"), full_page=True) + + answer = await page.evaluate( + """() => Array.from(document.querySelectorAll('[data-role="assistant"], [data-message-role="assistant"]')) + .map(n => n.textContent || '').join('\\n')""") + result["answer_excerpt"] = (answer or "")[-600:] + print(f"[{args.label}] indexed={indexed} in {result['index_seconds']}s; answer chars={len(answer or '')}") + + await ctx.close() + await browser.close() + + webms = sorted((out / "video").glob("*.webm")) + if webms: + webms[-1].rename(out / "video" / f"{args.label}.webm") + Path(out / "result.json").write_text(json.dumps(result, indent=2)) + print(f"[{args.label}] wrote {out/'result.json'}") + + +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("--doc", required=True) + ap.add_argument("--question", required=True) + ap.add_argument("--out", required=True) + asyncio.run(run(ap.parse_args())) + + +if __name__ == "__main__": + main()