diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bf8a3c04df..ba438d9ea6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 { diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 7e7d7026f6..4cfd87f1a1 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 73bb3d090a..edfae5ce16 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -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) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 0e9cce7c3e..d378e861d1 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -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}" diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py new file mode 100644 index 0000000000..83a2599857 --- /dev/null +++ b/studio/backend/core/rag/tool.py @@ -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) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b5626951c4..091dba4f8d 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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.", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 02270ab405..7c2a4e06bb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 "", diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 3928f79631..aa134d8b76 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -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: diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index cf5fb78677..dc57de469f 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -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({ ) : null} - + {ragToolEnabled ? ( +
+ ) : null} )} + {/* 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. */} +
{dictationSupported && ( 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 e7ce32e37d..d1f84c42ae 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -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::::`. 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((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((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((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); diff --git a/tests/python/test_rag_tool_handler.py b/tests/python/test_rag_tool_handler.py new file mode 100644 index 0000000000..552719ada8 --- /dev/null +++ b/tests/python/test_rag_tool_handler.py @@ -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