diff --git a/studio/bench/studio_rag_chat.py b/studio/bench/studio_rag_chat.py index af101da1c8..94e7d991b6 100644 --- a/studio/bench/studio_rag_chat.py +++ b/studio/bench/studio_rag_chat.py @@ -3,6 +3,7 @@ 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 @@ -12,21 +13,29 @@ 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"]), + ( + "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}) + 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 = 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 = 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 @@ -34,18 +43,23 @@ def token(c, base, pw): 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) + 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: + 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) + 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"] @@ -54,32 +68,48 @@ def main(): for q, phrases in QUESTIONS: payload = { "model": "local", - "messages": [{"role": "user", "content": q + - " Use the knowledge base; cite the source."}], + "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"}, + "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: + 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() + 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": + 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": @@ -92,20 +122,32 @@ def main(): 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"\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]}) + 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) + 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") + 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__": diff --git a/studio/bench/studio_rag_composer.py b/studio/bench/studio_rag_composer.py index 1a7c04fc46..eac0a001be 100644 --- a/studio/bench/studio_rag_composer.py +++ b/studio/bench/studio_rag_composer.py @@ -8,6 +8,7 @@ Usage: --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 @@ -20,16 +21,23 @@ 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}) + 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 = 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 = 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", "") @@ -43,21 +51,29 @@ def init_script(tok, refresh): async def run(args): out = Path(args.out) - (out / "video").mkdir(parents=True, exist_ok=True) + (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} + 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}) + 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.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")) + 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"): @@ -67,15 +83,15 @@ async def run(args): await btn.click() break await page.wait_for_timeout(700) - await page.screenshot(path=str(out / "02_rag_on.png")) + 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) + 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 + 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): @@ -90,30 +106,35 @@ async def run(args): 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")) + 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 + stop = page.locator( + 'button[aria-label="Stop generating"], button:has-text("Stop")' + ).first try: - await stop.wait_for(state="visible", timeout=30000) + await stop.wait_for(state = "visible", timeout = 30000) except Exception: pass try: - await stop.wait_for(state="hidden", timeout=180000) + 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) + 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')""") + .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 '')}") + print( + f"[{args.label}] indexed={indexed} in {result['index_seconds']}s; answer chars={len(answer or '')}" + ) await ctx.close() await browser.close() @@ -121,18 +142,18 @@ async def run(args): 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)) + 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) + 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()))