Studio: precache RAG reranker on toggle-on, not at app startup

Auto-downloading a 1.1 GB cross-encoder at every studio start is
wrong for users who never use rerank — reranker is opt-in by design.
Move the precache from a startup daemon thread to an explicit
POST /api/rag/reranker/precache endpoint, and have the chat settings
sheet call it the moment the 'Use reranker' switch is flipped on.

  - Backend: drop the startup _precache_reranker thread; add the
    /api/rag/reranker/precache route that calls precache_reranker().
  - Frontend: new precacheRagReranker() in rag-api, wired into the
    Switch's onCheckedChange so the download runs synchronously
    with a loading toast. On success: 'Reranker ready'. On failure:
    error toast + auto-flip the switch back off so the next query
    doesn't trigger another long hang.

First toggle-on pays the 1.1 GB download once; subsequent toggles
hit the HF cache and return ~instantly.
This commit is contained in:
Roland Tannous 2026-05-27 20:45:17 +04:00
commit 8fb2fb9e2a
4 changed files with 77 additions and 13 deletions

View file

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

View file

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

View file

@ -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({
</div>
<Switch
checked={enableRerank}
onCheckedChange={setEnableRerank}
onCheckedChange={(next) => {
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}
/>
</div>

View file

@ -311,6 +311,22 @@ export async function warmupRagEmbedder(): Promise<void> {
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<SearchHit[]> {