Studio: force RAG tool path — disable prefetch, RAG-first tool order, must-call directive

This commit is contained in:
Roland Tannous 2026-05-26 15:52:40 +04:00
commit 1f3a92cab9
3 changed files with 41 additions and 17 deletions

View file

@ -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 "
'<chunk id="N" source="..." page="..." score="...">...</chunk> '
"tags; cite them in your reply as [1], [2], etc."
),

View file

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

View file

@ -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 " +
"<chunk> 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