From ccebbed190b399f42613ee997e5c94ec1ee5be6d Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 18:15:02 +0400 Subject: [PATCH] Studio: always pre-fetch RAG + min-score threshold + retrieval logging - chat-adapter pre-fetches retrieval on every turn when RAG is on, regardless of provider. Users no longer have to phrase queries as 'the document I attached' for retrieval to fire. Local tool models still get search_knowledge_base registered as a refinement path. - New per-thread ragMinScore slider (Min relevance, 0..1) gates retrieved hits by dense cosine similarity. Hits below the floor (and BM25-only hits with no dense signal) are dropped server-side so unrelated docs don't get injected when the user's query is off-topic from what's indexed. - Backend logs at search start (scope, top_k, min_score, query preview), after retrieval (retrieved vs met_threshold counts), and on return (final hit count) for both /api/rag/search and the search_knowledge_base tool path. - System-prompt nudge prepended when pre-fetch returns hits so the model knows to cite [1], [2] rather than paraphrase silently. --- studio/backend/core/inference/tools.py | 1 + studio/backend/core/rag/retrieval.py | 28 +++++++++++++++++ studio/backend/core/rag/tool.py | 27 ++++++++++++++++- studio/backend/routes/rag.py | 27 +++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 30 +++++++++++-------- .../src/features/chat/chat-settings-sheet.tsx | 26 ++++++++++++++++ .../chat/stores/chat-runtime-store.ts | 16 +++++++++- .../frontend/src/features/rag/api/rag-api.ts | 2 ++ tests/python/test_rag_tool_handler.py | 6 +++- 9 files changed, 148 insertions(+), 15 deletions(-) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index d378e861d1..7af5eddd7f 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -569,6 +569,7 @@ def execute_tool( enable_rerank = bool(scope.get("enable_rerank")), reranker_model = scope.get("reranker_model"), default_top_k = int(scope.get("default_top_k") or 5), + min_score = float(scope.get("min_score") or 0.0), ) return f"Unknown tool: {name}" diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index 8fc3f203c2..7bec051e49 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -7,6 +7,10 @@ Reciprocal Rank Fusion is parameter-light: each candidate's fused score is the sum of ``1 / (rrf_k + rank)`` across rankers. It avoids the need to calibrate score scales between BM25 (raw, unbounded) and cosine similarity (-1..1). + +Hits also carry the raw dense cosine score when available so callers +can apply a meaningful similarity threshold (e.g., "drop hits below +0.3 cosine") — the fused RRF score isn't on a comparable scale. """ from __future__ import annotations @@ -30,6 +34,11 @@ class Hit: document_id: str | None = None chunk_index: int | None = None kind: str = "text" + # Raw cosine similarity from the dense retriever (0..1 for + # normalized embeddings). None when this chunk wasn't returned by + # the dense pass (BM25-only hit) — callers applying a similarity + # floor should treat None as "no signal" and exclude it. + dense_score: float | None = None def retrieve_bm25(scope: str, query: str, k: int | None = None) -> list[Hit]: @@ -62,6 +71,7 @@ def retrieve_dense( document_id = payload.get("document_id"), chunk_index = payload.get("chunk_index"), kind = payload.get("kind", "text"), + dense_score = r["score"], ) ) return out @@ -75,11 +85,16 @@ def _rrf_fuse( ) -> list[Hit]: fused: dict[str, float] = {} seen: dict[str, Hit] = {} + # Track the dense cosine score per chunk so it survives fusion — + # callers downstream filter on this, not the RRF score. + dense_scores: dict[str, float] = {} for ranking in rankings: for rank, hit in enumerate(ranking): fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / (rrf_k + rank + 1) if hit.chunk_id not in seen: seen[hit.chunk_id] = hit + if hit.dense_score is not None: + dense_scores[hit.chunk_id] = hit.dense_score ordered = sorted(fused.items(), key = lambda kv: kv[1], reverse = True)[:top_k] return [ Hit( @@ -88,6 +103,7 @@ def _rrf_fuse( document_id = seen[cid].document_id, chunk_index = seen[cid].chunk_index, kind = seen[cid].kind, + dense_score = dense_scores.get(cid), ) for cid, score in ordered ] @@ -114,3 +130,15 @@ def retrieve_hybrid( rrf_k = RAG_RRF_K, top_k = k or RAG_TOP_K_HYBRID, ) + + +def filter_by_min_score(hits: list[Hit], min_score: float) -> list[Hit]: + """Drop hits whose dense cosine score is below ``min_score``. + + Hits without a dense score (BM25-only) are dropped too — there's + no comparable signal to evaluate them against the similarity floor. + Use ``min_score = 0.0`` (or negative) to disable the filter. + """ + if min_score <= 0.0: + return hits + return [h for h in hits if h.dense_score is not None and h.dense_score >= min_score] diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 83a2599857..4916d77555 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -91,12 +91,14 @@ def search_knowledge_base( enable_rerank: bool = False, reranker_model: str | None = None, default_top_k: int = 5, + min_score: float = 0.0, ) -> str: """Execute the RAG search and return a tool-result string. `kb_id` takes precedence over `thread_id` when both are set — matches the create/upload contract that a document belongs to one - or the other, never both. + or the other, never both. ``min_score`` is a cosine-similarity + floor on dense hits; chunks below it are dropped. """ if not query or not query.strip(): return "Error: empty query." @@ -124,12 +126,35 @@ def search_knowledge_base( else: candidate_k = k + logger.info( + "search_knowledge_base: scope=%s top_k=%d min_score=%.3f rerank=%s query=%r", + scope, + k, + min_score, + enable_rerank, + query[:120], + ) + try: hits = retrieval.retrieve_hybrid(scope, query.strip(), k = candidate_k) except Exception as exc: # noqa: BLE001 logger.exception("search_knowledge_base retrieval failed") return f"Error: retrieval failed ({type(exc).__name__})." + retrieved_count = len(hits) + if min_score > 0.0: + hits = retrieval.filter_by_min_score(hits, min_score) + logger.info( + "search_knowledge_base: retrieved=%d met_threshold=%d (min_score=%.3f)", + retrieved_count, + len(hits), + min_score, + ) + else: + logger.info( + "search_knowledge_base: retrieved=%d (no threshold)", retrieved_count + ) + chunk_ids = [h.chunk_id for h in hits] lookup: dict[str, dict] = {} if chunk_ids: diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 88ce2d84b2..e35b0c0833 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -135,6 +135,10 @@ class SearchRequest(BaseModel): document_ids: list[str] | None = None enable_rerank: bool = False reranker_model: str | None = None + # Cosine-similarity floor on the dense retrieval score. Hits whose + # dense_score is below this (or absent — BM25-only hits) are + # dropped before the response is sent. 0.0 disables the filter. + min_score: float = Field(default = 0.0, ge = 0.0, le = 1.0) class SearchHit(BaseModel): @@ -1048,6 +1052,16 @@ def search( else: scope = thread_scope(payload.thread_id) + logger.info( + "RAG search: scope=%s mode=%s top_k=%d min_score=%.3f rerank=%s query=%r", + scope, + payload.mode, + payload.top_k, + payload.min_score, + payload.enable_rerank, + payload.query[:120], + ) + # When reranking is opt-in, pull a wider candidate pool so the # CrossEncoder has more to choose from before truncating to top_k. candidate_k = ( @@ -1073,6 +1087,18 @@ def search( document_ids = payload.document_ids, ) + retrieved_count = len(hits) + if payload.min_score > 0.0: + hits = retrieval.filter_by_min_score(hits, payload.min_score) + logger.info( + "RAG search: retrieved=%d met_threshold=%d (min_score=%.3f)", + retrieved_count, + len(hits), + payload.min_score, + ) + else: + logger.info("RAG search: retrieved=%d (no threshold)", retrieved_count) + chunk_ids = [h.chunk_id for h in hits] chunk_lookup: dict[str, dict] = {} if chunk_ids: @@ -1132,4 +1158,5 @@ def search( image_url = image_url, ) ) + logger.info("RAG search: returning %d hits", len(out)) return SearchResponse(hits = out) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index aa134d8b76..aebf3d4f61 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -94,12 +94,14 @@ function buildRagRequest( resolvedThreadId: string | undefined, enableRerank: boolean, topK: number, + minScore: number, ): SearchRequest | null { const base: SearchRequest = { query, top_k: topK, mode: "hybrid", enable_rerank: enableRerank, + min_score: minScore, }; if (source.kind === "thread") { if (!resolvedThreadId) return null; @@ -994,21 +996,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // don't break the chat — better to answer without context than // to drop a message the user just sent. // - // Phase 4: pre-fetch only fires when the new RAG button is on AND - // we can't register `search_knowledge_base` as a real tool — - // i.e., external providers or local models that don't expose - // tool-use. Local models with tool support take the tool-call - // path further down (see `enabled_tools` assembly), and the LLM - // decides per turn whether to invoke retrieval. + // Pre-fetch RAG context unconditionally when the RAG button is + // on and a source is selected. This runs for every provider — + // local-tool, local-no-tool, and external — so users don't have + // to phrase their query as "the document I attached" for + // retrieval to fire. On local tool-capable models the + // `search_knowledge_base` tool is *also* registered below as an + // optional refinement path (the LLM can run a second, narrower + // query if the pre-fetched chunks weren't enough). const ragSource = runtime.ragSource; const ragToolEnabled = runtime.ragToolEnabled; const ragToolPathTaken = ragToolEnabled && supportsTools && !isExternalRequest; - if ( - ragToolEnabled - && ragSource.kind !== "off" - && !ragToolPathTaken - ) { + if (ragToolEnabled && ragSource.kind !== "off") { const lastUser = [...outboundMessages] .reverse() .find((m) => m.role === "user"); @@ -1022,12 +1022,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { resolvedThreadId, runtime.enableRerank, runtime.ragTopK, + runtime.ragMinScore, ); if (ragReq) { try { const hits = await ragSearch(ragReq); if (hits.length > 0) { - const block = formatRagContext(hits); + const nudge = + "The following context was retrieved from the user's " + + "attached documents. Use it to answer; cite sources as " + + "[1], [2], etc. by source filename."; + const block = `${nudge}\n\n${formatRagContext(hits)}`; if ( outboundMessages[0]?.role === "system" && typeof outboundMessages[0].content === "string" @@ -1658,6 +1663,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { : null, enable_rerank: runtime.enableRerank, default_top_k: runtime.ragTopK, + min_score: runtime.ragMinScore, }, } : {}), diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 439b6088c7..606a94ef65 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -446,6 +446,8 @@ export function ChatSettingsPanel({ const setEnableRerank = useChatRuntimeStore((s) => s.setEnableRerank); const ragTopK = useChatRuntimeStore((s) => s.ragTopK); const setRagTopK = useChatRuntimeStore((s) => s.setRagTopK); + const ragMinScore = useChatRuntimeStore((s) => s.ragMinScore); + const setRagMinScore = useChatRuntimeStore((s) => s.setRagMinScore); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const { knowledgeBases, deleteKB } = useKnowledgeBases(); const { documents: threadDocs, remove: removeThreadDoc } = useThreadDocuments( @@ -1516,6 +1518,30 @@ export function ChatSettingsPanel({ grounding, more tokens.

+
+
+ + + {ragMinScore === 0 ? "off" : ragMinScore.toFixed(2)} + +
+ v != null && setRagMinScore(v)} + disabled={!ragEnabled} + /> +

+ Cosine-similarity floor for retrieved chunks (0 = off). + Set above 0 so unrelated docs are dropped — useful when + your query is off-topic from what's indexed. Try 0.3 as + a starting point. +

+
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index d1f84c42ae..5adc0e6589 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -300,6 +300,11 @@ type ChatRuntimeStore = { ragSource: RagSource; enableRerank: boolean; ragTopK: number; + // Cosine-similarity floor for RAG hits. Off (0) by default — set + // > 0 to drop chunks below the threshold so unrelated docs don't + // get injected into the prompt (e.g., asking about the weather + // when the indexed docs are about economics). + ragMinScore: number; hydratePersistedSettings: () => Promise; setModelLoading: (loading: boolean) => void; setModelRequiresTrustRemoteCode: (required: boolean) => void; @@ -346,6 +351,7 @@ type ChatRuntimeStore = { setRagSource: (source: RagSource) => void; setEnableRerank: (value: boolean) => void; setRagTopK: (value: number) => void; + setRagMinScore: (value: number) => void; setRagToolEnabled: (value: boolean) => void; }; @@ -365,7 +371,8 @@ type ScalarSettingKey = | "toolCallTimeout" | "ragSource" | "enableRerank" - | "ragTopK"; + | "ragTopK" + | "ragMinScore"; type PresetHydrationVersions = { customPresets: number; @@ -402,6 +409,7 @@ const SCALAR_SETTING_KEYS = [ "ragSource", "enableRerank", "ragTopK", + "ragMinScore", ] as const satisfies readonly ScalarSettingKey[]; const inferenceParamMutationVersions = Object.fromEntries( @@ -613,6 +621,7 @@ export const useChatRuntimeStore = create((set, get) => ({ ragSource: { kind: "off" }, enableRerank: false, ragTopK: 5, + ragMinScore: 0, hydratePersistedSettings: async () => { if (get().settingsHydrated) { return; @@ -846,6 +855,11 @@ export const useChatRuntimeStore = create((set, get) => ({ setScalarSettingVersion("ragTopK", ragTopK, state.ragTopK); return { ragTopK }; }), + setRagMinScore: (ragMinScore) => + set((state) => { + setScalarSettingVersion("ragMinScore", ragMinScore, state.ragMinScore); + return { ragMinScore }; + }), setToolsEnabled: (toolsEnabled, options) => set(() => { if (options?.persist !== false) { diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index e0ae25607d..de71717cab 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -61,6 +61,8 @@ export interface SearchRequest { document_ids?: string[]; enable_rerank?: boolean; reranker_model?: string; + /** Cosine-similarity floor (0..1). Hits below are dropped server-side. */ + min_score?: number; } export type JobEvent = diff --git a/tests/python/test_rag_tool_handler.py b/tests/python/test_rag_tool_handler.py index 552719ada8..18af735d1d 100644 --- a/tests/python/test_rag_tool_handler.py +++ b/tests/python/test_rag_tool_handler.py @@ -138,13 +138,15 @@ def test_execute_tool_dispatches_to_search_knowledge_base(): called = {} def _stub(*, query, top_k = None, scope_kb_id = None, scope_thread_id = None, - enable_rerank = False, reranker_model = None, default_top_k = 5): + enable_rerank = False, reranker_model = None, default_top_k = 5, + min_score = 0.0): called["query"] = query called["top_k"] = top_k called["scope_kb_id"] = scope_kb_id called["scope_thread_id"] = scope_thread_id called["enable_rerank"] = enable_rerank called["default_top_k"] = default_top_k + called["min_score"] = min_score return "stub-result" with patch("core.rag.tool.search_knowledge_base", _stub): @@ -156,6 +158,7 @@ def test_execute_tool_dispatches_to_search_knowledge_base(): "kb_id": "kb-1", "enable_rerank": True, "default_top_k": 3, + "min_score": 0.35, } }, ) @@ -166,6 +169,7 @@ def test_execute_tool_dispatches_to_search_knowledge_base(): assert called["scope_thread_id"] is None assert called["enable_rerank"] is True assert called["default_top_k"] == 3 + assert called["min_score"] == 0.35 def test_execute_tool_handles_missing_tool_context():