8848a310dfAdds a self-contained RAG stack to Studio: knowledge bases with chunked indexing, hybrid (dense + lexical) retrieval, and an automatic first-pass context inject into chat. Embeddings run through a local llama-server GGUF backend (default unsloth/bge-small-en-v1.5-GGUF) with a sentence-transformers fallback. The chat tool loop gains a search_knowledge_base tool, a per-turn re-search cap, and source citation, layered on top of the shared ToolLoopController.
456 lines
15 KiB
Python
456 lines
15 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""HTTP API for the RAG engine: KB CRUD, uploads, SSE ingestion, search.
|
|
|
|
Single-tenant: the subject gates access, not data. Without sqlite-vec the router
|
|
mounts but every endpoint returns 503.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import secrets
|
|
import time
|
|
import uuid
|
|
|
|
from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile
|
|
from fastapi.responses import FileResponse, StreamingResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
from auth.authentication import get_current_subject
|
|
from core.rag import config, ingestion, retrieval, store
|
|
from storage import rag_db
|
|
from utils.paths import ensure_dir, rag_uploads_root
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _require_rag() -> None:
|
|
if not rag_db.RAG_AVAILABLE:
|
|
raise HTTPException(
|
|
status_code = 503,
|
|
detail = "RAG is unavailable: the sqlite-vec extension could not be loaded.",
|
|
)
|
|
|
|
|
|
_SAFE = re.compile(r"[^A-Za-z0-9._-]+")
|
|
|
|
|
|
def _sanitize_filename(name: str) -> str:
|
|
base = os.path.basename(name or "").strip() or "document"
|
|
base = _SAFE.sub("_", base)
|
|
return base[:200]
|
|
|
|
|
|
def _save_upload(file: UploadFile) -> tuple[str, str]:
|
|
"""Persist an upload; returns (stored_path, filename)."""
|
|
filename = _sanitize_filename(file.filename or "document")
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
if ext not in config.UPLOAD_EXTS:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"Unsupported file type '{ext}'. Allowed: {sorted(config.UPLOAD_EXTS)}",
|
|
)
|
|
uploads = ensure_dir(rag_uploads_root())
|
|
stored_path = str(uploads / f"{uuid.uuid4().hex}{ext}")
|
|
size = 0
|
|
with open(stored_path, "wb") as out:
|
|
while True:
|
|
block = file.file.read(1 << 20)
|
|
if not block:
|
|
break
|
|
size += len(block)
|
|
out.write(block)
|
|
if size == 0:
|
|
os.remove(stored_path)
|
|
raise HTTPException(status_code = 400, detail = "Uploaded file is empty.")
|
|
return stored_path, filename
|
|
|
|
|
|
def _doc_view(row: dict) -> dict:
|
|
return {
|
|
"id": row["id"],
|
|
"filename": row["filename"],
|
|
"status": row["status"],
|
|
"error": row.get("error"),
|
|
"numChunks": row.get("num_chunks") or 0,
|
|
"kbId": row.get("kb_id"),
|
|
"threadId": row.get("thread_id"),
|
|
"createdAt": row.get("created_at"),
|
|
}
|
|
|
|
|
|
class CreateKbRequest(BaseModel):
|
|
name: str = Field(min_length = 1, max_length = 200)
|
|
description: str | None = None
|
|
|
|
|
|
class UpdateKbRequest(BaseModel):
|
|
name: str | None = Field(default = None, max_length = 200)
|
|
description: str | None = None
|
|
|
|
|
|
class SearchRequest(BaseModel):
|
|
query: str
|
|
kb_id: str | None = None
|
|
thread_id: str | None = None
|
|
top_k: int = Field(default = config.TOP_K_HYBRID, ge = 1, le = 50)
|
|
min_score: float = 0.0
|
|
mode: str = "hybrid" # hybrid | lexical | dense
|
|
|
|
|
|
@router.get("/knowledge-bases")
|
|
def list_knowledge_bases(subject: str = Depends(get_current_subject)) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
kbs = store.list_kbs(conn)
|
|
out = []
|
|
for kb in kbs:
|
|
docs = store.list_documents(conn, store.kb_scope(kb["id"]))
|
|
out.append(
|
|
{
|
|
"id": kb["id"],
|
|
"name": kb["name"],
|
|
"description": kb.get("description"),
|
|
"createdAt": kb.get("created_at"),
|
|
"documentCount": len(docs),
|
|
}
|
|
)
|
|
return {"knowledgeBases": out}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.post("/knowledge-bases")
|
|
def create_knowledge_base(
|
|
payload: CreateKbRequest, subject: str = Depends(get_current_subject)
|
|
) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
kb_id = store.create_kb(
|
|
conn,
|
|
name = payload.name.strip(),
|
|
description = (payload.description or None),
|
|
embedding_model = config.EMBEDDING_MODEL,
|
|
)
|
|
return {"id": kb_id, "name": payload.name.strip()}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.patch("/knowledge-bases/{kb_id}")
|
|
def update_knowledge_base(
|
|
kb_id: str,
|
|
payload: UpdateKbRequest,
|
|
subject: str = Depends(get_current_subject),
|
|
) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
if store.get_kb(conn, kb_id) is None:
|
|
raise HTTPException(status_code = 404, detail = "Knowledge base not found")
|
|
sets, params = [], []
|
|
if payload.name is not None:
|
|
sets.append("name=?")
|
|
params.append(payload.name.strip())
|
|
if payload.description is not None:
|
|
sets.append("description=?")
|
|
params.append(payload.description or None)
|
|
if sets:
|
|
params.append(kb_id)
|
|
conn.execute(f"UPDATE knowledge_bases SET {', '.join(sets)} WHERE id=?", params)
|
|
conn.commit()
|
|
return {"ok": True}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.delete("/knowledge-bases/{kb_id}")
|
|
def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject)) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
if store.get_kb(conn, kb_id) is None:
|
|
raise HTTPException(status_code = 404, detail = "Knowledge base not found")
|
|
store.delete_kb(conn, kb_id)
|
|
return {"ok": True}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.post("/knowledge-bases/{kb_id}/documents")
|
|
async def upload_kb_document(
|
|
kb_id: str,
|
|
file: UploadFile = File(...),
|
|
subject: str = Depends(get_current_subject),
|
|
) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
if store.get_kb(conn, kb_id) is None:
|
|
raise HTTPException(status_code = 404, detail = "Knowledge base not found")
|
|
finally:
|
|
conn.close()
|
|
stored_path, filename = _save_upload(file)
|
|
document_id, job_id = ingestion.start_ingestion(
|
|
store.kb_scope(kb_id), kb_id, None, filename, stored_path
|
|
)
|
|
return {"documentId": document_id, "jobId": job_id, "filename": filename}
|
|
|
|
|
|
@router.get("/knowledge-bases/{kb_id}/documents")
|
|
def list_kb_documents(kb_id: str, subject: str = Depends(get_current_subject)) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
docs = store.list_documents(conn, store.kb_scope(kb_id))
|
|
return {"documents": [_doc_view(d) for d in docs]}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.post("/threads/{thread_id}/documents")
|
|
async def upload_thread_document(
|
|
thread_id: str,
|
|
file: UploadFile = File(...),
|
|
subject: str = Depends(get_current_subject),
|
|
) -> dict:
|
|
_require_rag()
|
|
stored_path, filename = _save_upload(file)
|
|
document_id, job_id = ingestion.start_ingestion(
|
|
store.thread_scope(thread_id), None, thread_id, filename, stored_path
|
|
)
|
|
return {"documentId": document_id, "jobId": job_id, "filename": filename}
|
|
|
|
|
|
@router.get("/threads/{thread_id}/documents")
|
|
def list_thread_documents(thread_id: str, subject: str = Depends(get_current_subject)) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
docs = store.list_documents(conn, store.thread_scope(thread_id))
|
|
return {"documents": [_doc_view(d) for d in docs]}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.delete("/documents/{document_id}")
|
|
def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict:
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
if store.get_document(conn, document_id) is None:
|
|
raise HTTPException(status_code = 404, detail = "Document not found")
|
|
store.delete_document(conn, document_id)
|
|
return {"ok": True}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.get("/jobs/{job_id}")
|
|
def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict:
|
|
_require_rag()
|
|
row = ingestion.get_job_status(job_id)
|
|
if row is None:
|
|
raise HTTPException(status_code = 404, detail = "Job not found")
|
|
return {
|
|
"id": row["id"],
|
|
"documentId": row["document_id"],
|
|
"status": row["status"],
|
|
"stage": row.get("stage"),
|
|
"progress": row.get("progress") or 0.0,
|
|
"error": row.get("error"),
|
|
}
|
|
|
|
|
|
@router.get("/jobs/{job_id}/events")
|
|
def job_events(job_id: str, subject: str = Depends(get_current_subject)) -> StreamingResponse:
|
|
_require_rag()
|
|
|
|
def gen():
|
|
try:
|
|
for event in ingestion.job_events(job_id):
|
|
yield f"data: {json.dumps(event)}\n\n"
|
|
except Exception as exc: # noqa: BLE001
|
|
yield f"data: {json.dumps({'type': 'error', 'error': str(exc)})}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
return StreamingResponse(
|
|
gen(),
|
|
media_type = "text/event-stream",
|
|
headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
|
)
|
|
|
|
|
|
@router.post("/search")
|
|
def search(payload: SearchRequest, subject: str = Depends(get_current_subject)) -> dict:
|
|
_require_rag()
|
|
if payload.kb_id:
|
|
scope = store.kb_scope(payload.kb_id)
|
|
elif payload.thread_id:
|
|
scope = store.thread_scope(payload.thread_id)
|
|
else:
|
|
raise HTTPException(status_code = 400, detail = "Provide kb_id or thread_id")
|
|
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
if payload.mode == "lexical":
|
|
hits = retrieval.retrieve_lexical(conn, scope, payload.query, payload.top_k)
|
|
elif payload.mode == "dense":
|
|
hits = retrieval.retrieve_dense(conn, scope, payload.query, payload.top_k)
|
|
else:
|
|
hits = retrieval.retrieve_hybrid(conn, scope, payload.query, k = payload.top_k)
|
|
hits = retrieval.filter_min_score(hits, payload.min_score)
|
|
rows = store.chunks_by_id(conn, [h.chunk_id for h in hits])
|
|
results = []
|
|
for h in hits:
|
|
r = rows.get(h.chunk_id)
|
|
if r is None:
|
|
continue
|
|
results.append(
|
|
{
|
|
"chunkId": h.chunk_id,
|
|
"documentId": r["document_id"],
|
|
"filename": r["filename"],
|
|
"page": r["page_number"],
|
|
"score": h.score,
|
|
"text": r["text"],
|
|
}
|
|
)
|
|
return {"results": results}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# Per-process secret so pdf.js range requests fetch the file without a bearer
|
|
# header; tokens only work on this server instance.
|
|
_PREVIEW_SECRET = secrets.token_bytes(32)
|
|
_PREVIEW_TTL = 600 # seconds
|
|
|
|
_CONTENT_TYPES = {
|
|
".pdf": "application/pdf",
|
|
".txt": "text/plain; charset=utf-8",
|
|
".md": "text/markdown; charset=utf-8",
|
|
".markdown": "text/markdown; charset=utf-8",
|
|
".html": "text/html; charset=utf-8",
|
|
".htm": "text/html; charset=utf-8",
|
|
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
}
|
|
|
|
|
|
def _sign_document(document_id: str) -> str:
|
|
exp = int(time.time()) + _PREVIEW_TTL
|
|
payload = f"{document_id}.{exp}"
|
|
sig = hmac.new(_PREVIEW_SECRET, payload.encode(), hashlib.sha256).hexdigest()
|
|
return f"{payload}.{sig}"
|
|
|
|
|
|
def _verify_document_token(token: str) -> str | None:
|
|
try:
|
|
document_id, exp_s, sig = token.rsplit(".", 2)
|
|
except ValueError:
|
|
return None
|
|
expected = hmac.new(
|
|
_PREVIEW_SECRET, f"{document_id}.{exp_s}".encode(), hashlib.sha256
|
|
).hexdigest()
|
|
if not hmac.compare_digest(sig, expected):
|
|
return None
|
|
try:
|
|
if int(exp_s) < int(time.time()):
|
|
return None
|
|
except ValueError:
|
|
return None
|
|
return document_id
|
|
|
|
|
|
@router.get("/documents/{document_id}/preview-target")
|
|
def preview_target(
|
|
document_id: str,
|
|
chunk_id: str | None = Query(default = None),
|
|
subject: str = Depends(get_current_subject),
|
|
) -> dict:
|
|
"""Resolve a citation to filename, page, and highlight regions."""
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
doc = store.get_document(conn, document_id)
|
|
if doc is None:
|
|
raise HTTPException(status_code = 404, detail = "Document not found")
|
|
ext = os.path.splitext(doc["filename"])[1].lower()
|
|
out = {
|
|
"documentId": document_id,
|
|
"filename": doc["filename"],
|
|
"mediaKind": "pdf" if ext == ".pdf" else "text",
|
|
"targetPage": None,
|
|
"pdfRegions": [],
|
|
"text": None,
|
|
}
|
|
if chunk_id:
|
|
row = conn.execute(
|
|
"SELECT text, page_number, pdf_regions_json FROM chunks WHERE id=?",
|
|
(chunk_id,),
|
|
).fetchone()
|
|
if row is not None:
|
|
out["text"] = row["text"]
|
|
out["targetPage"] = row["page_number"]
|
|
if row["pdf_regions_json"]:
|
|
try:
|
|
out["pdfRegions"] = json.loads(row["pdf_regions_json"])
|
|
except Exception:
|
|
out["pdfRegions"] = []
|
|
return out
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@router.get("/documents/{document_id}/file-url")
|
|
def document_file_url(document_id: str, subject: str = Depends(get_current_subject)) -> dict:
|
|
"""Mint a short-lived signed URL for the source file."""
|
|
_require_rag()
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
doc = store.get_document(conn, document_id)
|
|
if doc is None or not doc.get("stored_path"):
|
|
raise HTTPException(status_code = 404, detail = "Document file not available")
|
|
finally:
|
|
conn.close()
|
|
token = _sign_document(document_id)
|
|
return {"url": f"/api/rag/documents/{document_id}/file-signed?token={token}"}
|
|
|
|
|
|
@router.get("/documents/{document_id}/file-signed", response_model = None)
|
|
def document_file_signed(document_id: str, token: str = Query(...)) -> FileResponse:
|
|
"""Serve the source file gated by the HMAC token (no bearer) so pdf.js range
|
|
requests work."""
|
|
_require_rag()
|
|
signed_id = _verify_document_token(token)
|
|
if signed_id != document_id:
|
|
raise HTTPException(status_code = 401, detail = "Invalid or expired token")
|
|
conn = rag_db.get_connection()
|
|
try:
|
|
doc = store.get_document(conn, document_id)
|
|
finally:
|
|
conn.close()
|
|
stored_path = (doc or {}).get("stored_path")
|
|
if not doc or not stored_path or not os.path.isfile(stored_path):
|
|
raise HTTPException(status_code = 404, detail = "Document file not found")
|
|
# Confine to the uploads root (defense in depth).
|
|
uploads = os.path.realpath(str(rag_uploads_root()))
|
|
if not os.path.realpath(stored_path).startswith(uploads):
|
|
raise HTTPException(status_code = 403, detail = "Forbidden")
|
|
ext = os.path.splitext(doc["filename"])[1].lower()
|
|
return FileResponse(
|
|
stored_path,
|
|
media_type = _CONTENT_TYPES.get(ext, "application/octet-stream"),
|
|
filename = doc["filename"],
|
|
)
|