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.
This commit is contained in:
Roland Tannous
parent
e27c079e3d
commit
ccebbed190
9 changed files with 148 additions and 15 deletions
|
|
@ -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}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
</p>
|
||||
</div>
|
||||
<div className="flex flex-col gap-1.5">
|
||||
<div className="flex items-center justify-between">
|
||||
<label className="text-[12px] font-medium text-muted-foreground">
|
||||
Min relevance
|
||||
</label>
|
||||
<span className="text-[12px] tabular-nums text-muted-foreground">
|
||||
{ragMinScore === 0 ? "off" : ragMinScore.toFixed(2)}
|
||||
</span>
|
||||
</div>
|
||||
<Slider
|
||||
value={[ragMinScore]}
|
||||
min={0}
|
||||
max={1}
|
||||
step={0.05}
|
||||
onValueChange={([v]) => v != null && setRagMinScore(v)}
|
||||
disabled={!ragEnabled}
|
||||
/>
|
||||
<p className="text-[11px] text-muted-foreground">
|
||||
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.
|
||||
</p>
|
||||
</div>
|
||||
<div className="flex items-center justify-between">
|
||||
<div className="flex flex-col">
|
||||
<span className="text-[13px] font-medium">
|
||||
|
|
|
|||
|
|
@ -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<void>;
|
||||
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<ChatRuntimeStore>((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<ChatRuntimeStore>((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) {
|
||||
|
|
|
|||
|
|
@ -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 =
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue