From 1f3a92cab91f13e929d2f747fa5f529f8f04b13e Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 26 May 2026 15:52:40 +0400 Subject: [PATCH] =?UTF-8?q?Studio:=20force=20RAG=20tool=20path=20=E2=80=94?= =?UTF-8?q?=20disable=20prefetch,=20RAG-first=20tool=20order,=20must-call?= =?UTF-8?q?=20directive?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- studio/backend/core/rag/tool.py | 8 ++-- studio/backend/routes/inference.py | 10 ++-- .../src/features/chat/api/chat-adapter.ts | 46 +++++++++++++------ 3 files changed, 44 insertions(+), 20 deletions(-) diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 6a5cb18486..913bddc5b3 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -21,9 +21,11 @@ SEARCH_KNOWLEDGE_BASE_TOOL = { "function": { "name": "search_knowledge_base", "description": ( - "Search the user's attached documents. Call this when the user " - "references content from their docs, asks fact-heavy questions, " - "or needs grounded citations. Returns chunks wrapped in " + "ALWAYS CALL THIS TOOL FIRST before answering any user question. " + "It searches the user's attached documents and returns the chunks " + "you must ground your reply in. Do not answer from your own " + "knowledge until you have called this tool with a focused query " + "derived from the user's latest message. Returns chunks wrapped in " '... ' "tags; cite them in your reply as [1], [2], etc." ), diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 53bac5d9ac..980794845f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2379,10 +2379,11 @@ async def openai_chat_completions( from core.inference.tools import ALL_TOOLS if payload.enabled_tools is not None: + # Preserve client-supplied order so prioritised tools + # (e.g. search_knowledge_base when RAG is on) appear first. + _by_name = {t["function"]["name"]: t for t in ALL_TOOLS} tools_to_use = [ - t - for t in ALL_TOOLS - if t["function"]["name"] in payload.enabled_tools + _by_name[name] for name in payload.enabled_tools if name in _by_name ] else: tools_to_use = ALL_TOOLS @@ -2877,8 +2878,9 @@ async def openai_chat_completions( from core.inference.tools import ALL_TOOLS if payload.enabled_tools is not None: + _by_name = {t["function"]["name"]: t for t in ALL_TOOLS} _sf_tools_to_use = [ - t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools + _by_name[name] for name in payload.enabled_tools if name in _by_name ] else: _sf_tools_to_use = ALL_TOOLS diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index a920b26594..e31043a435 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -983,23 +983,41 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { Boolean(message), ); - const safeSystemPrompt = - typeof params.systemPrompt === "string" ? params.systemPrompt : ""; - if (safeSystemPrompt.trim()) { - outboundMessages.unshift({ - role: "system", - content: safeSystemPrompt.trim(), - }); - } + // Temporary debug toggle: when false, the pre-fetch path is skipped + // entirely so retrieval only happens via the LLM-invoked + // search_knowledge_base tool. Flip back to true to restore the + // always-on grounding for external providers / non-tool models. + const RAG_PREFETCH_ENABLED = false; - // Pre-fetch RAG context for the last user turn; failures don't block chat. - // Runs for all providers; local tool-capable models also get the tool below - // for a narrower follow-up query if needed. const ragSource = runtime.ragSource; const ragToolEnabled = runtime.ragToolEnabled; const ragToolPathTaken = ragToolEnabled && supportsTools && !isExternalRequest; - if (ragToolEnabled && ragSource.kind !== "off") { + + const safeSystemPrompt = + typeof params.systemPrompt === "string" ? params.systemPrompt : ""; + const systemPromptParts: string[] = []; + if (safeSystemPrompt.trim()) { + systemPromptParts.push(safeSystemPrompt.trim()); + } + if (ragToolPathTaken && ragSource.kind !== "off") { + systemPromptParts.push( + "RAG retrieval is enabled for this conversation. You MUST call " + + "the `search_knowledge_base` tool before answering ANY user " + + "question — even short ones, follow-ups, clarifications, or " + + "questions you think you already know the answer to. Issue the " + + "tool call first, then ground your reply in the returned " + + " blocks and cite them as [1], [2], etc.", + ); + } + if (systemPromptParts.length > 0) { + outboundMessages.unshift({ + role: "system", + content: systemPromptParts.join("\n\n"), + }); + } + + if (RAG_PREFETCH_ENABLED && ragToolEnabled && ragSource.kind !== "off") { const lastUser = [...outboundMessages] .reverse() .find((m) => m.role === "user"); @@ -1633,9 +1651,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ? { enable_tools: true, enabled_tools: [ + // RAG goes first so the model sees it before any other + // tool when scanning the spec list. + ...(ragToolPathTaken ? ["search_knowledge_base"] : []), ...(toolsEnabled ? ["web_search"] : []), ...(codeToolsEnabled ? ["python", "terminal"] : []), - ...(ragToolPathTaken ? ["search_knowledge_base"] : []), ], // Per-request scope for the LLM-invoked tool; tool path only. ...(ragToolPathTaken