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:
parent
3f6a390df6
commit
8fb2fb9e2a
4 changed files with 77 additions and 13 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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[]> {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue