ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: RAG end-to-end UI + local GGUF tool-call validation (both variants)

Browse source
This commit is contained in:
Daniel Han 2026-05-31 14:07:00 +00:00
commit 4f20be0d2f
3 changed files with 270 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

19
studio/RAG_FAST_INDEXING.md
View file

@ -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.

112
studio/bench/studio_rag_chat.py Normal file
View file

@ -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()

140
studio/bench/studio_rag_composer.py Normal file
View file

@ -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()