diff --git a/studio/backend/main.py b/studio/backend/main.py index 16061abe91..d1a61d4417 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -260,16 +260,7 @@ async def lifespan(app: FastAPI): except Exception: pass # non-critical - def _precache_reranker(): - try: - from core.rag.reranker import precache_reranker - - precache_reranker() - except Exception: - pass # non-critical - threading.Thread(target = _precache, daemon = True).start() - threading.Thread(target = _precache_reranker, daemon = True).start() # Initialize RSA key pair for API key encryption (external providers) from core.inference.key_exchange import init_key_pair diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index beefc0b2f3..3f9febaf64 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -460,6 +460,32 @@ def warmup_rag_embedder( return {"ok": True, "model": model_name} +@router.post("/reranker/precache") +def precache_rag_reranker( + current_subject: str = Depends(get_current_subject), +) -> dict: + """Download the reranker weights (~1.1 GB) into the HF cache. + + Called from the frontend the moment the user flips the "Use + reranker" switch ON so the cost lands on the explicit toggle + instead of the first chat turn — where a multi-minute download + looks like a hung tool call. + """ + from core.rag.reranker import precache_reranker + from utils.rag.config import RAG_RERANKER_MODEL + + try: + precache_reranker() + except Exception as exc: # noqa: BLE001 + logger.warning( + "RAG reranker precache failed", + model = RAG_RERANKER_MODEL, + error = str(exc), + ) + return {"ok": False, "model": RAG_RERANKER_MODEL, "error": str(exc)} + return {"ok": True, "model": RAG_RERANKER_MODEL} + + @router.put("/defaults", response_model = RagDefaults) def set_rag_defaults( payload: UpdateRagDefaultsRequest, diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 089e90cf9b..ed658463d0 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -97,9 +97,10 @@ import { KBCreateDialog } from "@/features/rag/components/kb-create-dialog"; import { useKnowledgeBases } from "@/features/rag/hooks/use-knowledge-bases"; import { useThreadDocuments } from "@/features/rag/hooks/use-kb-documents"; import { useRagStore } from "@/features/rag/stores/rag-store"; -import type { - ChunkingStrategy as RagChunkingStrategy, - KBMode, +import { + type ChunkingStrategy as RagChunkingStrategy, + type KBMode, + precacheRagReranker, } from "@/features/rag/api/rag-api"; import { Add01Icon, Delete02Icon } from "@hugeicons/core-free-icons"; @@ -1577,7 +1578,37 @@ export function ChatSettingsPanel({ { + setEnableRerank(next); + if (!next) return; + // First flip-on may have to download ~1.1 GB; the + // toast covers the latency so the user doesn't think + // the next query is hung waiting on the reranker. + const toastId = toast.loading( + "Preparing reranker (one-time download)…", + ); + void precacheRagReranker() + .then((res) => { + if (res.ok) { + toast.success("Reranker ready", { id: toastId }); + } else { + toast.error( + `Reranker download failed: ${res.error ?? "unknown"}`, + { id: toastId }, + ); + setEnableRerank(false); + } + }) + .catch((err: unknown) => { + toast.error( + `Reranker download failed: ${ + err instanceof Error ? err.message : String(err) + }`, + { id: toastId }, + ); + setEnableRerank(false); + }); + }} disabled={!ragEnabled} /> diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index 9f56e20ac9..16b9bdfda5 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -311,6 +311,22 @@ export async function warmupRagEmbedder(): Promise { await authFetch("/api/rag/warmup", { method: "POST" }); } +/** Download reranker weights into the HF cache. ~1.1 GB on first call; + * no-op when cached. Called when the user flips the reranker toggle so + * the download lands on an explicit action, not the first chat turn. */ +export async function precacheRagReranker(): Promise<{ + ok: boolean; + model: string; + error?: string; +}> { + const response = await authFetch("/api/rag/reranker/precache", { + method: "POST", + }); + return parseJsonOrThrow<{ ok: boolean; model: string; error?: string }>( + response, + ); +} + // --- Search --- export async function search(req: SearchRequest): Promise {