ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: RAG-as-tool composer button (Phase 4)

Browse source 
Promotes RAG to a first-class composer toggle alongside Think / Web
Search / Code, with tool-use semantics on local models that support
tools and a pre-fetch fallback on external providers. The model
decides when to call `search_knowledge_base` on local inference; on
external providers retrieval still fires before each message (the
existing pre-fetch path), gated on the same button.

Backend
- core/rag/tool.py (new): search_knowledge_base handler + JSON-schema
  tool spec. Resolves scope (kb_id wins over thread_id) from the
  request's rag_scope, runs retrieve_hybrid + optional rerank, then
  hydrates filename / page_number / text from sqlite and formats as
  numbered Markdown citations ('[1] file.pdf (page 5): ...') for the
  LLM to cite. Empty scope returns a user-facing hint; empty results
  return a clear no-match message instead of an empty string.
- core/inference/tools.py: SEARCH_KNOWLEDGE_BASE_TOOL added to
  ALL_TOOLS (lazy import keeps tools.py importable on inference
  paths that never touch RAG). execute_tool() gains a tool_context
  parameter that carries per-request extras the LLM doesn't see
  (currently just rag_scope). The new 'search_knowledge_base' branch
  dispatches to the handler with scope unpacked from tool_context.
- core/inference/llama_cpp.py + safetensors_agentic.py +
  orchestrator.py: thread tool_context through generate_chat_completion_
  with_tools / run_safetensors_tool_loop / execute_tool. Both local
  backends (GGUF llama-server and safetensors agentic) carry the same
  context object.
- models/inference.py: ChatCompletionRequest gains optional
  rag_scope: dict ({kb_id?, thread_id?, enable_rerank?, default_top_k?,
  reranker_model?}). Ignored unless 'search_knowledge_base' is in
  enabled_tools.
- routes/inference.py: both the GGUF and safetensors call sites for
  generate_chat_completion_with_tools forward payload.rag_scope into
  tool_context.

Frontend
- chat-runtime-store.ts: global ragToolEnabled boolean + setter +
  CHAT_RAG_TOOL_ENABLED_KEY localStorage, mirroring toolsEnabled /
  codeToolsEnabled. Settings-hydration migration auto-flips
  ragToolEnabled=true for pre-Phase-4 users who already had ragSource
  set, so existing RAG users don't silently lose retrieval on upgrade.
- shared-composer.tsx: new 'RAG' pill button after Images (uses
  lucide BookOpenIcon, composer-pill-btn style, data-active toggle).
  Disabled when no model is loaded. Toggling on from ragSource='off'
  auto-flips source to 'thread' so the sidebar lands ready-to-go.
- chat-adapter.ts:
  * The existing pre-fetch block is now gated on ragToolEnabled AND
    only fires when the tool path isn't viable (external provider OR
    local model without tool-use support). Tool-capable local models
    skip pre-fetch and let the LLM decide.
  * The local-model body assembly adds 'search_knowledge_base' to
    enabled_tools and packs ragSource + enableRerank + ragTopK into a
    rag_scope object the backend tool handler consumes.
- chat-settings-sheet.tsx: entire Retrieval CollapsibleSection is
  wrapped in {ragToolEnabled && ...} so it hides when the button is
  off — the button is now the single on/off control. The 'Off'
  option is removed from the Source dropdown (the button handles
  that). Default open when shown so settings are one click away.

Tests
- test_rag_tool_handler.py: handler covers empty query, missing
  scope, kb_id > thread_id precedence, thread-only path, citation
  formatting (numbered + page numbers + unknown source); tool spec
  shape (function/name/required); execute_tool dispatch with and
  without tool_context; ALL_TOOLS includes the new spec without
  dropping the existing ones.

Verification scope
- Local GGUF with tools: toggle button on, upload doc, ask about
  doc content → assistant emits a search_knowledge_base tool call
  card (rendered by the existing ToolFallback component since no
  custom UI exists yet — that's a v2 nice-to-have).
- External provider (Anthropic / OpenAI / etc.): same button, same
  UX, but uses the pre-fetch path under the hood.
- Migration: pre-existing ragSource != off → button initializes ON
  so retrieval keeps working.
This commit is contained in:
Roland Tannous 2026-05-24 13:41:45 +04:00
commit c74fc13ebc
12 changed files with 531 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/backend/core/inference/llama_cpp.py
View file

@ -4428,6 +4428,7 @@ class LlamaCppBackend:
auto_heal_tool_calls: bool = True,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
tool_context: Optional[dict] = None,
) -> Generator[dict, None, None]:
"""
Agentic loop: let the model call tools, execute them, and continue.
@ -5080,6 +5081,7 @@ class LlamaCppBackend:
cancel_event = cancel_event,
timeout = _effective_timeout,
session_id = session_id,
tool_context = tool_context,
)
yield {

2
studio/backend/core/inference/orchestrator.py
View file

@ -838,6 +838,7 @@ class InferenceOrchestrator:
auto_heal_tool_calls: bool = True,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
tool_context: Optional[dict] = None,
use_adapter: Optional[Union[bool, str]] = None,
**_unused,
):
@ -895,6 +896,7 @@ class InferenceOrchestrator:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
tool_context = tool_context,
)
def generate_with_adapter_control(

2
studio/backend/core/inference/safetensors_agentic.py
View file

@ -105,6 +105,7 @@ def run_safetensors_tool_loop(
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
tool_context: Optional[dict] = None,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
@ -340,6 +341,7 @@ def run_safetensors_tool_loop(
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
tool_context = tool_context,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", tool_name, exc)

33
studio/backend/core/inference/tools.py
View file

@ -502,7 +502,20 @@ TERMINAL_TOOL = {
},
}
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL]
# Lazy import — keeps studio.db init lazy so tools.py doesn't pull in
# the whole rag stack on inference paths that never see RAG.
def _get_rag_tool_spec():
from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL
return SEARCH_KNOWLEDGE_BASE_TOOL
# RAG_SEARCH_TOOL is included in ALL_TOOLS; routes/inference.py filters
# the list against payload.enabled_tools so each request only sees the
# tools the frontend explicitly enabled. When the RAG button is off
# the tool name won't be in enabled_tools and the LLM will never see
# the spec.
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, _get_rag_tool_spec()]
_TIMEOUT_UNSET = object()
@ -514,12 +527,17 @@ def execute_tool(
cancel_event = None,
timeout: int | None = _TIMEOUT_UNSET,
session_id: str | None = None,
tool_context: dict | None = None,
) -> str:
"""Execute a tool by name with the given arguments. Returns result as a string.
``timeout``: int sets per-call limit in seconds, ``None`` means no limit,
unset (default) uses ``_EXEC_TIMEOUT`` (300 s).
``session_id``: optional thread/session ID for per-conversation sandbox isolation.
``tool_context``: optional per-request extras the LLM does not see (RAG scope,
future per-tool overrides). Keys consumed:
- ``rag_scope``: ``{kb_id?, thread_id?, enable_rerank?, default_top_k?,
reranker_model?}`` — consumed by ``search_knowledge_base``.
"""
logger.info(
f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
@ -539,6 +557,19 @@ def execute_tool(
return _bash_exec(
arguments.get("command", ""), cancel_event, effective_timeout, session_id
)
if name == "search_knowledge_base":
from core.rag.tool import search_knowledge_base
scope = (tool_context or {}).get("rag_scope") or {}
return search_knowledge_base(
query = arguments.get("query", ""),
top_k = arguments.get("top_k"),
scope_kb_id = scope.get("kb_id"),
scope_thread_id = scope.get("thread_id"),
enable_rerank = bool(scope.get("enable_rerank")),
reranker_model = scope.get("reranker_model"),
default_top_k = int(scope.get("default_top_k") or 5),
)
return f"Unknown tool: {name}"

179
studio/backend/core/rag/tool.py Normal file
View file

@ -0,0 +1,179 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""`search_knowledge_base` tool — RAG retrieval surfaced to the LLM.
Invoked from `core/inference/tools.execute_tool` when the local model
emits a `search_knowledge_base` call. The handler runs the existing
hybrid retrieval, hydrates chunk text + filename + page number from
sqlite, and returns a Markdown-with-numbered-citations string that
the LLM consumes as the tool-result message.
Scope (`kb_id` / `thread_id`) is not exposed as a tool argument — it
comes from the chat-completions request body (`rag_scope`) so the
LLM doesn't need to know about KB UUIDs.
"""
from __future__ import annotations
import logging
from typing import Any
logger = logging.getLogger(__name__)
SEARCH_KNOWLEDGE_BASE_TOOL = {
"type": "function",
"function": {
"name": "search_knowledge_base",
"description": (
"Search the user's attached documents for information relevant to "
"the user's question. Call this when the user references content "
"from their docs, asks fact-heavy questions, or needs grounded "
"citations. Returns numbered chunks with source filenames; cite "
"them in your reply as [1], [2], etc."
),
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": (
"A focused search query — phrase it as the question "
"you want answered, not as a keyword list."
),
},
"top_k": {
"type": "integer",
"minimum": 1,
"maximum": 20,
"description": (
"How many chunks to retrieve (default 5). Higher = "
"more grounding, more tokens."
),
},
},
"required": ["query"],
},
},
}
def _format_hits_for_llm(hits: list[Any]) -> str:
"""Render hits as numbered Markdown citations for the LLM.
Empty results produce a one-line message rather than an empty
string — the model needs to know the search ran but found nothing
so it can fall back to its own knowledge or ask the user.
"""
if not hits:
return (
"No matching chunks were found in the attached documents. "
"Either nothing in this scope is relevant, or no documents "
"have been ingested yet."
)
lines: list[str] = []
for index, hit in enumerate(hits, start = 1):
name = hit.get("filename") or "unknown source"
page = hit.get("page_number")
suffix = f" (page {page})" if page is not None else ""
text = (hit.get("text") or "").strip()
lines.append(f"[{index}] {name}{suffix}: {text}")
return "\n\n".join(lines)
def search_knowledge_base(
*,
query: str,
top_k: int | None = None,
scope_kb_id: str | None = None,
scope_thread_id: str | None = None,
enable_rerank: bool = False,
reranker_model: str | None = None,
default_top_k: int = 5,
) -> 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.
"""
if not query or not query.strip():
return "Error: empty query."
if not scope_kb_id and not scope_thread_id:
return (
"No knowledge base or thread documents are configured for "
"retrieval. Ask the user to upload a document or select a "
"knowledge base in the chat settings."
)
from core.rag import retrieval
from core.rag.vector_store import kb_scope, thread_scope
from storage.studio_db import get_connection
scope = (
kb_scope(scope_kb_id) if scope_kb_id else thread_scope(scope_thread_id)
)
k = top_k if top_k is not None else default_top_k
if enable_rerank:
from utils.rag.config import RAG_RERANK_CANDIDATE_K
candidate_k = max(k, RAG_RERANK_CANDIDATE_K)
else:
candidate_k = k
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__})."
chunk_ids = [h.chunk_id for h in hits]
lookup: dict[str, dict] = {}
if chunk_ids:
placeholders = ",".join("?" for _ in chunk_ids)
with get_connection() as conn:
rows = conn.execute(
f"""
SELECT c.id AS chunk_id, c.text, c.page_number,
c.kind, d.filename
FROM rag_chunks c
JOIN rag_documents d ON d.id = c.document_id
WHERE c.id IN ({placeholders})
""",
chunk_ids,
).fetchall()
for row in rows:
lookup[row["chunk_id"]] = dict(row)
if enable_rerank and hits:
from core.rag import reranker
pairs = [
(hit, lookup[hit.chunk_id]["text"])
for hit in hits
if hit.chunk_id in lookup
]
try:
hits = reranker.rerank(
query.strip(),
pairs,
model_name = reranker_model,
top_k = k,
)
except Exception as exc: # noqa: BLE001
logger.warning("rerank failed in search_knowledge_base: %s", exc)
hits = hits[:k]
else:
hits = hits[:k]
# Image-kind hits don't carry LLM-friendly text — skip them. The
# paired caption (linked_chunk_id) usually surfaces separately.
formatted = [
lookup[hit.chunk_id]
for hit in hits
if hit.chunk_id in lookup and lookup[hit.chunk_id].get("kind") != "image"
]
return _format_hits_for_llm(formatted)

10
studio/backend/models/inference.py
View file

@ -686,6 +686,16 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
)
rag_scope: Optional[dict] = Field(
None,
description = (
"[x-unsloth] Per-request context the `search_knowledge_base` tool "
"consumes when the LLM invokes it. Shape: "
"{kb_id?: str, thread_id?: str, enable_rerank?: bool, "
"default_top_k?: int, reranker_model?: str}. Ignored unless "
"'search_knowledge_base' is in enabled_tools."
),
)
cancel_id: Optional[str] = Field(
None,
description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.",

10
studio/backend/routes/inference.py
View file

@ -2483,6 +2483,11 @@ async def openai_chat_completions(
if payload.tool_call_timeout is not None
else 300,
session_id = payload.session_id,
tool_context = (
{"rag_scope": payload.rag_scope}
if payload.rag_scope
else None
),
)
_tool_sentinel = object()
@ -2950,6 +2955,11 @@ async def openai_chat_completions(
def sf_generate_with_tools():
return backend.generate_chat_completion_with_tools(
tool_context = (
{"rag_scope": payload.rag_scope}
if payload.rag_scope
else None
),
messages = _sf_chat_messages,
tools = _sf_tools_to_use,
system_prompt = _sf_system_prompt or "",

41
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -993,8 +993,22 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// prepend it as a system-role block. Failures are logged but
// 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.
const ragSource = runtime.ragSource;
if (ragSource.kind !== "off") {
const ragToolEnabled = runtime.ragToolEnabled;
const ragToolPathTaken =
ragToolEnabled && supportsTools && !isExternalRequest;
if (
ragToolEnabled
&& ragSource.kind !== "off"
&& !ragToolPathTaken
) {
const lastUser = [...outboundMessages]
.reverse()
.find((m) => m.role === "user");
@ -1617,13 +1631,36 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
...(supportsPreserveThinking
? { preserve_thinking: preserveThinking }
: {}),
...(supportsTools && (toolsEnabled || codeToolsEnabled)
...(supportsTools
&& (toolsEnabled || codeToolsEnabled || ragToolPathTaken)
? {
enable_tools: true,
enabled_tools: [
...(toolsEnabled ? ["web_search"] : []),
...(codeToolsEnabled ? ["python", "terminal"] : []),
...(ragToolPathTaken ? ["search_knowledge_base"] : []),
],
// Phase 4: per-request RAG context the backend's
// `search_knowledge_base` handler reads when the LLM
// invokes the tool. Only sent when the tool path
// is taken — external providers fall through to the
// pre-fetch block above.
...(ragToolPathTaken
? {
rag_scope: {
kb_id:
ragSource.kind === "kb"
? ragSource.kbId
: null,
thread_id:
ragSource.kind === "thread"
? (resolvedThreadId ?? null)
: null,
enable_rerank: runtime.enableRerank,
default_top_k: runtime.ragTopK,
},
}
: {}),
auto_heal_tool_calls:
useChatRuntimeStore.getState().autoHealToolCalls,
max_tool_calls_per_message:

10
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -441,6 +441,7 @@ export function ChatSettingsPanel({
const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
const ragSource = useChatRuntimeStore((s) => s.ragSource);
const setRagSource = useChatRuntimeStore((s) => s.setRagSource);
const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled);
const enableRerank = useChatRuntimeStore((s) => s.enableRerank);
const setEnableRerank = useChatRuntimeStore((s) => s.setEnableRerank);
const ragTopK = useChatRuntimeStore((s) => s.ragTopK);
@ -1246,7 +1247,8 @@ export function ChatSettingsPanel({
</CollapsibleSection>
) : null}
<CollapsibleSection label="Retrieval" defaultOpen={false}>
{ragToolEnabled ? (
<CollapsibleSection label="Retrieval" defaultOpen={true}>
<div className="flex flex-col gap-3 pt-1">
<div className="flex flex-col gap-1.5">
<label className="text-[12px] font-medium text-muted-foreground">
@ -1265,11 +1267,6 @@ export function ChatSettingsPanel({
align="start"
className="w-[var(--radix-dropdown-menu-trigger-width)]"
>
<DropdownMenuItem
onSelect={() => setRagSource({ kind: "off" })}
>
Off
</DropdownMenuItem>
<DropdownMenuItem
onSelect={() => setRagSource({ kind: "thread" })}
>
@ -1527,6 +1524,7 @@ export function ChatSettingsPanel({
</div>
</div>
</CollapsibleSection>
) : null}
<CollapsibleSection label="System Prompt" defaultOpen={true}>
<button

32
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -21,7 +21,7 @@ import { isTauri } from "@/lib/api-base";
import { isMultimodalResponse } from "./types/api";
import { getImageInputUnavailableReason } from "./utils/image-input-support";
import { useAui } from "@assistant-ui/react";
import { ArrowUpIcon, FileTextIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
import { ArrowUpIcon, BookOpenIcon, FileTextIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
import { useRagStore } from "@/features/rag/stores/rag-store";
import { subscribeToJobEvents } from "@/features/rag/api/rag-api";
import { toast } from "@/lib/toast";
@ -368,6 +368,8 @@ export function SharedComposer({
const setImageToolsEnabled = useChatRuntimeStore(
(s) => s.setImageToolsEnabled,
);
const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled);
const setRagToolEnabled = useChatRuntimeStore((s) => s.setRagToolEnabled);
const lastOpenRouterChosenModel = useChatRuntimeStore(
(s) => s.lastOpenRouterChosenModel,
);
@ -1275,6 +1277,34 @@ export function SharedComposer({
<span>Images</span>
</button>
)}
{/* RAG: master switch for retrieval. On local models with
tool-use support, registers `search_knowledge_base` as a
tool the LLM can call. On external providers, falls back
to the pre-fetch path. The sidebar Retrieval section
configures the source / top-K / reranker; the button is
the only on/off control. */}
<button
type="button"
disabled={!modelLoaded}
onClick={() => {
const next = !ragToolEnabled;
setRagToolEnabled(next);
if (next && ragSource.kind === "off") {
setRagSource({ kind: "thread" });
}
}}
className="composer-pill-btn"
data-active={ragToolEnabled && modelLoaded ? "true" : "false"}
aria-label={ragToolEnabled ? "Disable RAG" : "Enable RAG"}
title={
ragToolEnabled
? "RAG on — the model can search your attached documents"
: "Enable RAG — let the model search your documents"
}
>
<BookOpenIcon className="size-3.5" />
<span>RAG</span>
</button>
</div>
<div className="flex items-center gap-1">
{dictationSupported && (

27
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -26,6 +26,7 @@ export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled";
export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled";
export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled";
export const CHAT_RAG_TOOL_ENABLED_KEY = "unsloth_chat_rag_tool_enabled";
// External provider selection is encoded into `params.checkpoint` as
// `external::<providerId>::<modelId>`. PersistedChatSettings deliberately
@ -264,6 +265,7 @@ type ChatRuntimeStore = {
*/
supportsBuiltinImageGeneration: boolean;
toolsEnabled: boolean;
ragToolEnabled: boolean;
codeToolsEnabled: boolean;
imageToolsEnabled: boolean;
toolStatus: string | null;
@ -344,6 +346,7 @@ type ChatRuntimeStore = {
setRagSource: (source: RagSource) => void;
setEnableRerank: (value: boolean) => void;
setRagTopK: (value: number) => void;
setRagToolEnabled: (value: boolean) => void;
};
type PersistedChatSettings = Awaited<
@ -578,6 +581,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
supportsBuiltinCodeExecution: false,
supportsBuiltinImageGeneration: false,
toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false),
// Phase 4: RAG button defaults off. Migration nudge happens after
// settings hydration, when the persisted ragSource becomes visible —
// see hydratePersistedSettings.
ragToolEnabled: loadBool(CHAT_RAG_TOOL_ENABLED_KEY, false),
codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false),
imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false),
toolStatus: null,
@ -630,6 +637,21 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
),
...getHydratedSettingsState(settings, state, hydrationVersions),
};
// Phase 4 migration: pre-existing users with ragSource set
// before the RAG button shipped should keep getting RAG —
// auto-flip ragToolEnabled so the button starts ON for them.
// The CHAT_RAG_TOOL_ENABLED_KEY localStorage write makes the
// migration stick across reloads.
const hydratedRagSource =
(nextState.ragSource as RagSource | undefined) ?? state.ragSource;
if (
hydratedRagSource &&
hydratedRagSource.kind !== "off" &&
!state.ragToolEnabled
) {
nextState.ragToolEnabled = true;
saveBool(CHAT_RAG_TOOL_ENABLED_KEY, true);
}
return nextState;
});
} catch {
@ -831,6 +853,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
}
return { toolsEnabled };
}),
setRagToolEnabled: (ragToolEnabled) =>
set(() => {
saveBool(CHAT_RAG_TOOL_ENABLED_KEY, ragToolEnabled);
return { ragToolEnabled };
}),
setCodeToolsEnabled: (codeToolsEnabled) =>
set(() => {
saveBool(CHAT_CODE_TOOLS_ENABLED_KEY, codeToolsEnabled);

193
tests/python/test_rag_tool_handler.py Normal file
View file

@ -0,0 +1,193 @@
"""Unit tests for the `search_knowledge_base` tool handler (Phase 4)."""
import sys
from pathlib import Path
from unittest.mock import patch
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
def _make_hit(chunk_id: str):
"""Minimal stand-in for retrieval.Hit — just needs .chunk_id."""
class _Hit:
pass
h = _Hit()
h.chunk_id = chunk_id
h.score = 1.0
h.kind = "text"
h.document_id = None
h.chunk_index = 0
return h
def test_empty_query_returns_error():
from core.rag.tool import search_knowledge_base
result = search_knowledge_base(query = "", scope_thread_id = "t-1")
assert result.startswith("Error:")
assert "empty" in result.lower()
def test_missing_scope_returns_user_facing_hint():
from core.rag.tool import search_knowledge_base
result = search_knowledge_base(
query = "anything",
scope_kb_id = None,
scope_thread_id = None,
)
assert "No knowledge base" in result
assert "thread documents" in result
def test_kb_takes_precedence_over_thread():
"""When both kb_id and thread_id are passed, kb_id wins."""
from core.rag import tool
captured = {}
def _stub_retrieve(scope, query, k):
captured["scope"] = scope
return []
with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
"retrieve_hybrid",
_stub_retrieve):
result = tool.search_knowledge_base(
query = "x",
scope_kb_id = "kb-abc",
scope_thread_id = "thread-xyz",
)
assert captured["scope"].startswith("kb_")
assert "kb-abc" in captured["scope"]
assert "thread" not in captured["scope"].split("kb_")[1]
def test_thread_scope_when_only_thread_set():
from core.rag import tool
captured = {}
def _stub_retrieve(scope, query, k):
captured["scope"] = scope
return []
with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
"retrieve_hybrid",
_stub_retrieve):
tool.search_knowledge_base(
query = "x",
scope_thread_id = "thread-xyz",
)
assert captured["scope"].startswith("thread_")
def test_empty_results_message_is_user_facing():
from core.rag.tool import _format_hits_for_llm
result = _format_hits_for_llm([])
assert "No matching chunks" in result
def test_format_hits_produces_numbered_citations():
from core.rag.tool import _format_hits_for_llm
hits = [
{"filename": "alpha.pdf", "page_number": 3, "text": "first body"},
{"filename": "beta.md", "page_number": None, "text": "second body"},
]
result = _format_hits_for_llm(hits)
assert "[1] alpha.pdf (page 3): first body" in result
assert "[2] beta.md: second body" in result
# Each hit on its own paragraph so the LLM can cite cleanly.
assert "\n\n" in result
def test_format_hits_handles_unknown_source():
from core.rag.tool import _format_hits_for_llm
hits = [{"filename": None, "page_number": None, "text": "orphan"}]
result = _format_hits_for_llm(hits)
assert "[1] unknown source: orphan" in result
def test_tool_spec_shape_is_openai_compatible():
from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL
assert SEARCH_KNOWLEDGE_BASE_TOOL["type"] == "function"
fn = SEARCH_KNOWLEDGE_BASE_TOOL["function"]
assert fn["name"] == "search_knowledge_base"
assert "query" in fn["parameters"]["required"]
assert "top_k" in fn["parameters"]["properties"]
# Description should hint at when to call so the LLM picks it up
# appropriately. Don't lock the exact wording.
assert "documents" in fn["description"].lower()
def test_execute_tool_dispatches_to_search_knowledge_base():
"""tools.execute_tool should route 'search_knowledge_base' correctly."""
from core.inference import tools
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):
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
return "stub-result"
with patch("core.rag.tool.search_knowledge_base", _stub):
result = tools.execute_tool(
"search_knowledge_base",
{"query": "hello", "top_k": 7},
tool_context = {
"rag_scope": {
"kb_id": "kb-1",
"enable_rerank": True,
"default_top_k": 3,
}
},
)
assert result == "stub-result"
assert called["query"] == "hello"
assert called["top_k"] == 7
assert called["scope_kb_id"] == "kb-1"
assert called["scope_thread_id"] is None
assert called["enable_rerank"] is True
assert called["default_top_k"] == 3
def test_execute_tool_handles_missing_tool_context():
"""tool_context=None should still dispatch without crashing."""
from core.inference import tools
def _stub(*, query, **_kwargs):
return f"got: {query}"
with patch("core.rag.tool.search_knowledge_base", _stub):
result = tools.execute_tool(
"search_knowledge_base",
{"query": "ping"},
tool_context = None,
)
assert result == "got: ping"
def test_all_tools_includes_rag():
from core.inference.tools import ALL_TOOLS
names = [t["function"]["name"] for t in ALL_TOOLS]
assert "search_knowledge_base" in names
assert "web_search" in names # regression — we shouldn't have removed the others
assert "python" in names