mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-02 22:02:13 +08:00
feat(tools): add tool_search for deferred MCP tool loading (#1176)
* feat(tools): add tool_search for deferred MCP tool loading When multiple MCP servers are enabled, total tool count can exceed 30-50, causing context bloat and degraded tool selection accuracy. This adds a deferred tool loading mechanism controlled by `tool_search.enabled` config. - Add ToolSearchConfig with single `enabled` field - Add DeferredToolRegistry with regex search (select:, +keyword, keyword) - Add tool_search tool returning OpenAI-compatible function JSON - Add DeferredToolFilterMiddleware to hide deferred schemas from bind_tools - Add <available-deferred-tools> section to system prompt - Enable MCP tool_name_prefix to prevent cross-server name collisions - Add 34 unit tests covering registry, tool, prompt, and middleware * fix: reset stale deferred registry and bump config_version - Reset deferred registry upfront in get_available_tools() to prevent stale tool entries when MCP servers are disabled between calls - Bump config_version to 2 for new tool_search config field Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(tests): mock get_app_config in prompt section tests for CI CI has no config.yaml, causing TestDeferredToolsPromptSection to fail with FileNotFoundError. Add autouse fixture to mock get_app_config. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -240,6 +240,11 @@ def _build_middlewares(config: RunnableConfig, model_name: str | None, agent_nam
|
||||
if model_config is not None and model_config.supports_vision:
|
||||
middlewares.append(ViewImageMiddleware())
|
||||
|
||||
# Add DeferredToolFilterMiddleware to hide deferred tool schemas from model binding
|
||||
if app_config.tool_search.enabled:
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
middlewares.append(DeferredToolFilterMiddleware())
|
||||
|
||||
# Add SubagentLimitMiddleware to truncate excess parallel task calls
|
||||
subagent_enabled = config.get("configurable", {}).get("subagent_enabled", False)
|
||||
if subagent_enabled:
|
||||
@@ -314,13 +319,11 @@ def make_lead_agent(config: RunnableConfig):
|
||||
|
||||
if is_bootstrap:
|
||||
# Special bootstrap agent with minimal prompt for initial custom agent creation flow
|
||||
system_prompt = apply_prompt_template(subagent_enabled=subagent_enabled, max_concurrent_subagents=max_concurrent_subagents, available_skills=set(["bootstrap"]))
|
||||
|
||||
return create_agent(
|
||||
model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled),
|
||||
tools=get_available_tools(model_name=model_name, subagent_enabled=subagent_enabled) + [setup_agent],
|
||||
middleware=_build_middlewares(config, model_name=model_name),
|
||||
system_prompt=system_prompt,
|
||||
system_prompt=apply_prompt_template(subagent_enabled=subagent_enabled, max_concurrent_subagents=max_concurrent_subagents, available_skills=set(["bootstrap"])),
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
|
||||
@@ -235,6 +235,8 @@ You: "Deploying to staging..." [proceed]
|
||||
|
||||
{skills_section}
|
||||
|
||||
{deferred_tools_section}
|
||||
|
||||
{subagent_section}
|
||||
|
||||
<working_directory existed="true">
|
||||
@@ -417,6 +419,31 @@ def get_agent_soul(agent_name: str | None) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
def get_deferred_tools_prompt_section() -> str:
|
||||
"""Generate <available-deferred-tools> block for the system prompt.
|
||||
|
||||
Lists only deferred tool names so the agent knows what exists
|
||||
and can use tool_search to load them.
|
||||
Returns empty string when tool_search is disabled or no tools are deferred.
|
||||
"""
|
||||
from deerflow.tools.builtins.tool_search import get_deferred_registry
|
||||
|
||||
try:
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
if not get_app_config().tool_search.enabled:
|
||||
return ""
|
||||
except FileNotFoundError:
|
||||
return ""
|
||||
|
||||
registry = get_deferred_registry()
|
||||
if not registry:
|
||||
return ""
|
||||
|
||||
names = "\n".join(e.name for e in registry.entries)
|
||||
return f"<available-deferred-tools>\n{names}\n</available-deferred-tools>"
|
||||
|
||||
|
||||
def apply_prompt_template(subagent_enabled: bool = False, max_concurrent_subagents: int = 3, *, agent_name: str | None = None, available_skills: set[str] | None = None) -> str:
|
||||
# Get memory context
|
||||
memory_context = _get_memory_context(agent_name)
|
||||
@@ -446,11 +473,15 @@ def apply_prompt_template(subagent_enabled: bool = False, max_concurrent_subagen
|
||||
# Get skills section
|
||||
skills_section = get_skills_prompt_section(available_skills)
|
||||
|
||||
# Get deferred tools section (tool_search)
|
||||
deferred_tools_section = get_deferred_tools_prompt_section()
|
||||
|
||||
# Format the prompt with dynamic skills and memory
|
||||
prompt = SYSTEM_PROMPT_TEMPLATE.format(
|
||||
agent_name=agent_name or "DeerFlow 2.0",
|
||||
soul=get_agent_soul(agent_name),
|
||||
skills_section=skills_section,
|
||||
deferred_tools_section=deferred_tools_section,
|
||||
memory_context=memory_context,
|
||||
subagent_section=subagent_section,
|
||||
subagent_reminder=subagent_reminder,
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Middleware to filter deferred tool schemas from model binding.
|
||||
|
||||
When tool_search is enabled, MCP tools are registered in the DeferredToolRegistry
|
||||
and passed to ToolNode for execution, but their schemas should NOT be sent to the
|
||||
LLM via bind_tools (that's the whole point of deferral — saving context tokens).
|
||||
|
||||
This middleware intercepts wrap_model_call and removes deferred tools from
|
||||
request.tools so that model.bind_tools only receives active tool schemas.
|
||||
The agent discovers deferred tools at runtime via the tool_search tool.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import override
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DeferredToolFilterMiddleware(AgentMiddleware[AgentState]):
|
||||
"""Remove deferred tools from request.tools before model binding.
|
||||
|
||||
ToolNode still holds all tools (including deferred) for execution routing,
|
||||
but the LLM only sees active tool schemas — deferred tools are discoverable
|
||||
via tool_search at runtime.
|
||||
"""
|
||||
|
||||
def _filter_tools(self, request: ModelRequest) -> ModelRequest:
|
||||
from deerflow.tools.builtins.tool_search import get_deferred_registry
|
||||
|
||||
registry = get_deferred_registry()
|
||||
if not registry:
|
||||
return request
|
||||
|
||||
deferred_names = {e.name for e in registry.entries}
|
||||
active_tools = [t for t in request.tools if getattr(t, "name", None) not in deferred_names]
|
||||
|
||||
if len(active_tools) < len(request.tools):
|
||||
logger.debug(f"Filtered {len(request.tools) - len(active_tools)} deferred tool schema(s) from model binding")
|
||||
|
||||
return request.override(tools=active_tools)
|
||||
|
||||
@override
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ModelCallResult:
|
||||
return handler(self._filter_tools(request))
|
||||
|
||||
@override
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
||||
) -> ModelCallResult:
|
||||
return await handler(self._filter_tools(request))
|
||||
@@ -17,6 +17,7 @@ from deerflow.config.subagents_config import load_subagents_config_from_dict
|
||||
from deerflow.config.summarization_config import load_summarization_config_from_dict
|
||||
from deerflow.config.title_config import load_title_config_from_dict
|
||||
from deerflow.config.tool_config import ToolConfig, ToolGroupConfig
|
||||
from deerflow.config.tool_search_config import ToolSearchConfig, load_tool_search_config_from_dict
|
||||
|
||||
load_dotenv()
|
||||
|
||||
@@ -32,6 +33,7 @@ class AppConfig(BaseModel):
|
||||
tool_groups: list[ToolGroupConfig] = Field(default_factory=list, description="Available tool groups")
|
||||
skills: SkillsConfig = Field(default_factory=SkillsConfig, description="Skills configuration")
|
||||
extensions: ExtensionsConfig = Field(default_factory=ExtensionsConfig, description="Extensions configuration (MCP servers and skills state)")
|
||||
tool_search: ToolSearchConfig = Field(default_factory=ToolSearchConfig, description="Tool search / deferred loading configuration")
|
||||
model_config = ConfigDict(extra="allow", frozen=False)
|
||||
checkpointer: CheckpointerConfig | None = Field(default=None, description="Checkpointer configuration")
|
||||
|
||||
@@ -101,6 +103,10 @@ class AppConfig(BaseModel):
|
||||
if "subagents" in config_data:
|
||||
load_subagents_config_from_dict(config_data["subagents"])
|
||||
|
||||
# Load tool_search config if present
|
||||
if "tool_search" in config_data:
|
||||
load_tool_search_config_from_dict(config_data["tool_search"])
|
||||
|
||||
# Load checkpointer config if present
|
||||
if "checkpointer" in config_data:
|
||||
load_checkpointer_config_from_dict(config_data["checkpointer"])
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Configuration for deferred tool loading via tool_search."""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ToolSearchConfig(BaseModel):
|
||||
"""Configuration for deferred tool loading via tool_search.
|
||||
|
||||
When enabled, MCP tools are not loaded into the agent's context directly.
|
||||
Instead, they are listed by name in the system prompt and discoverable
|
||||
via the tool_search tool at runtime.
|
||||
"""
|
||||
|
||||
enabled: bool = Field(
|
||||
default=False,
|
||||
description="Defer tools and enable tool_search",
|
||||
)
|
||||
|
||||
|
||||
_tool_search_config: ToolSearchConfig | None = None
|
||||
|
||||
|
||||
def get_tool_search_config() -> ToolSearchConfig:
|
||||
"""Get the tool search config, loading from AppConfig if needed."""
|
||||
global _tool_search_config
|
||||
if _tool_search_config is None:
|
||||
_tool_search_config = ToolSearchConfig()
|
||||
return _tool_search_config
|
||||
|
||||
|
||||
def load_tool_search_config_from_dict(data: dict) -> ToolSearchConfig:
|
||||
"""Load tool search config from a dict (called during AppConfig loading)."""
|
||||
global _tool_search_config
|
||||
_tool_search_config = ToolSearchConfig.model_validate(data)
|
||||
return _tool_search_config
|
||||
@@ -53,7 +53,7 @@ async def get_mcp_tools() -> list[BaseTool]:
|
||||
if oauth_interceptor is not None:
|
||||
tool_interceptors.append(oauth_interceptor)
|
||||
|
||||
client = MultiServerMCPClient(servers_config, tool_interceptors=tool_interceptors)
|
||||
client = MultiServerMCPClient(servers_config, tool_interceptors=tool_interceptors, tool_name_prefix=True)
|
||||
|
||||
# Get all tools from all servers
|
||||
tools = await client.get_tools()
|
||||
|
||||
168
backend/packages/harness/deerflow/tools/builtins/tool_search.py
Normal file
168
backend/packages/harness/deerflow/tools/builtins/tool_search.py
Normal file
@@ -0,0 +1,168 @@
|
||||
"""Tool search — deferred tool discovery at runtime.
|
||||
|
||||
Contains:
|
||||
- DeferredToolRegistry: stores deferred tools and handles regex search
|
||||
- tool_search: the LangChain tool the agent calls to discover deferred tools
|
||||
|
||||
The agent sees deferred tool names in <available-deferred-tools> but cannot
|
||||
call them until it fetches their full schema via the tool_search tool.
|
||||
Source-agnostic: no mention of MCP or tool origin.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
from langchain.tools import BaseTool
|
||||
from langchain_core.tools import tool
|
||||
from langchain_core.utils.function_calling import convert_to_openai_function
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_RESULTS = 5 # Max tools returned per search
|
||||
|
||||
|
||||
# ── Registry ──
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeferredToolEntry:
|
||||
"""Lightweight metadata for a deferred tool (no full schema in context)."""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
tool: BaseTool # Full tool object, returned only on search match
|
||||
|
||||
|
||||
class DeferredToolRegistry:
|
||||
"""Registry of deferred tools, searchable by regex pattern."""
|
||||
|
||||
def __init__(self):
|
||||
self._entries: list[DeferredToolEntry] = []
|
||||
|
||||
def register(self, tool: BaseTool) -> None:
|
||||
self._entries.append(
|
||||
DeferredToolEntry(
|
||||
name=tool.name,
|
||||
description=tool.description or "",
|
||||
tool=tool,
|
||||
)
|
||||
)
|
||||
|
||||
def search(self, query: str) -> list[BaseTool]:
|
||||
"""Search deferred tools by regex pattern against name + description.
|
||||
|
||||
Supports three query forms (aligned with Claude Code):
|
||||
- "select:name1,name2" — exact name match
|
||||
- "+keyword rest" — name must contain keyword, rank by rest
|
||||
- "keyword query" — regex match against name + description
|
||||
|
||||
Returns:
|
||||
List of matched BaseTool objects (up to MAX_RESULTS).
|
||||
"""
|
||||
if query.startswith("select:"):
|
||||
names = {n.strip() for n in query[7:].split(",")}
|
||||
return [e.tool for e in self._entries if e.name in names][:MAX_RESULTS]
|
||||
|
||||
if query.startswith("+"):
|
||||
parts = query[1:].split(None, 1)
|
||||
required = parts[0].lower()
|
||||
candidates = [e for e in self._entries if required in e.name.lower()]
|
||||
if len(parts) > 1:
|
||||
candidates.sort(
|
||||
key=lambda e: _regex_score(parts[1], e),
|
||||
reverse=True,
|
||||
)
|
||||
return [e.tool for e in candidates][:MAX_RESULTS]
|
||||
|
||||
# General regex search
|
||||
try:
|
||||
regex = re.compile(query, re.IGNORECASE)
|
||||
except re.error:
|
||||
regex = re.compile(re.escape(query), re.IGNORECASE)
|
||||
|
||||
scored = []
|
||||
for entry in self._entries:
|
||||
searchable = f"{entry.name} {entry.description}"
|
||||
if regex.search(searchable):
|
||||
score = 2 if regex.search(entry.name) else 1
|
||||
scored.append((score, entry))
|
||||
|
||||
scored.sort(key=lambda x: x[0], reverse=True)
|
||||
return [entry.tool for _, entry in scored][:MAX_RESULTS]
|
||||
|
||||
@property
|
||||
def entries(self) -> list[DeferredToolEntry]:
|
||||
return list(self._entries)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._entries)
|
||||
|
||||
|
||||
def _regex_score(pattern: str, entry: DeferredToolEntry) -> int:
|
||||
try:
|
||||
regex = re.compile(pattern, re.IGNORECASE)
|
||||
except re.error:
|
||||
regex = re.compile(re.escape(pattern), re.IGNORECASE)
|
||||
return len(regex.findall(f"{entry.name} {entry.description}"))
|
||||
|
||||
|
||||
# ── Singleton ──
|
||||
|
||||
_registry: DeferredToolRegistry | None = None
|
||||
|
||||
|
||||
def get_deferred_registry() -> DeferredToolRegistry | None:
|
||||
return _registry
|
||||
|
||||
|
||||
def set_deferred_registry(registry: DeferredToolRegistry) -> None:
|
||||
global _registry
|
||||
_registry = registry
|
||||
|
||||
|
||||
def reset_deferred_registry() -> None:
|
||||
"""Reset the deferred registry singleton. Useful for testing."""
|
||||
global _registry
|
||||
_registry = None
|
||||
|
||||
|
||||
# ── Tool ──
|
||||
|
||||
|
||||
@tool
|
||||
def tool_search(query: str) -> str:
|
||||
"""Fetches full schema definitions for deferred tools so they can be called.
|
||||
|
||||
Deferred tools appear by name in <available-deferred-tools> in the system
|
||||
prompt. Until fetched, only the name is known — there is no parameter
|
||||
schema, so the tool cannot be invoked. This tool takes a query, matches
|
||||
it against the deferred tool list, and returns the matched tools' complete
|
||||
definitions. Once a tool's schema appears in that result, it is callable.
|
||||
|
||||
Query forms:
|
||||
- "select:Read,Edit,Grep" — fetch these exact tools by name
|
||||
- "notebook jupyter" — keyword search, up to max_results best matches
|
||||
- "+slack send" — require "slack" in the name, rank by remaining terms
|
||||
|
||||
Args:
|
||||
query: Query to find deferred tools. Use "select:<tool_name>" for
|
||||
direct selection, or keywords to search.
|
||||
|
||||
Returns:
|
||||
Matched tool definitions as JSON array.
|
||||
"""
|
||||
registry = get_deferred_registry()
|
||||
if registry is None:
|
||||
return "No deferred tools available."
|
||||
|
||||
matched_tools = registry.search(query)
|
||||
if not matched_tools:
|
||||
return f"No tools found matching: {query}"
|
||||
|
||||
# Use LangChain's built-in serialization to produce OpenAI function format.
|
||||
# This is model-agnostic: all LLMs understand this standard schema.
|
||||
tool_defs = [convert_to_openai_function(t) for t in matched_tools[:MAX_RESULTS]]
|
||||
|
||||
return json.dumps(tool_defs, indent=2, ensure_ascii=False)
|
||||
@@ -5,6 +5,7 @@ from langchain.tools import BaseTool
|
||||
from deerflow.config import get_app_config
|
||||
from deerflow.reflection import resolve_variable
|
||||
from deerflow.tools.builtins import ask_clarification_tool, present_file_tool, task_tool, view_image_tool
|
||||
from deerflow.tools.builtins.tool_search import reset_deferred_registry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -42,27 +43,6 @@ def get_available_tools(
|
||||
config = get_app_config()
|
||||
loaded_tools = [resolve_variable(tool.use, BaseTool) for tool in config.tools if groups is None or tool.group in groups]
|
||||
|
||||
# Get cached MCP tools if enabled
|
||||
# NOTE: We use ExtensionsConfig.from_file() instead of config.extensions
|
||||
# to always read the latest configuration from disk. This ensures that changes
|
||||
# made through the Gateway API (which runs in a separate process) are immediately
|
||||
# reflected when loading MCP tools.
|
||||
mcp_tools = []
|
||||
if include_mcp:
|
||||
try:
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.mcp.cache import get_cached_mcp_tools
|
||||
|
||||
extensions_config = ExtensionsConfig.from_file()
|
||||
if extensions_config.get_enabled_mcp_servers():
|
||||
mcp_tools = get_cached_mcp_tools()
|
||||
if mcp_tools:
|
||||
logger.info(f"Using {len(mcp_tools)} cached MCP tool(s)")
|
||||
except ImportError:
|
||||
logger.warning("MCP module not available. Install 'langchain-mcp-adapters' package to enable MCP tools.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get cached MCP tools: {e}")
|
||||
|
||||
# Conditionally add tools based on config
|
||||
builtin_tools = BUILTIN_TOOLS.copy()
|
||||
|
||||
@@ -81,4 +61,41 @@ def get_available_tools(
|
||||
builtin_tools.append(view_image_tool)
|
||||
logger.info(f"Including view_image_tool for model '{model_name}' (supports_vision=True)")
|
||||
|
||||
# Get cached MCP tools if enabled
|
||||
# NOTE: We use ExtensionsConfig.from_file() instead of config.extensions
|
||||
# to always read the latest configuration from disk. This ensures that changes
|
||||
# made through the Gateway API (which runs in a separate process) are immediately
|
||||
# reflected when loading MCP tools.
|
||||
mcp_tools = []
|
||||
# Reset deferred registry upfront to prevent stale state from previous calls
|
||||
reset_deferred_registry()
|
||||
if include_mcp:
|
||||
try:
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.mcp.cache import get_cached_mcp_tools
|
||||
|
||||
extensions_config = ExtensionsConfig.from_file()
|
||||
if extensions_config.get_enabled_mcp_servers():
|
||||
mcp_tools = get_cached_mcp_tools()
|
||||
if mcp_tools:
|
||||
logger.info(f"Using {len(mcp_tools)} cached MCP tool(s)")
|
||||
|
||||
# When tool_search is enabled, register MCP tools in the
|
||||
# deferred registry and add tool_search to builtin tools.
|
||||
if config.tool_search.enabled:
|
||||
from deerflow.tools.builtins.tool_search import DeferredToolRegistry, set_deferred_registry
|
||||
from deerflow.tools.builtins.tool_search import tool_search as tool_search_tool
|
||||
|
||||
registry = DeferredToolRegistry()
|
||||
for t in mcp_tools:
|
||||
registry.register(t)
|
||||
set_deferred_registry(registry)
|
||||
builtin_tools.append(tool_search_tool)
|
||||
logger.info(f"Tool search active: {len(mcp_tools)} tools deferred")
|
||||
except ImportError:
|
||||
logger.warning("MCP module not available. Install 'langchain-mcp-adapters' package to enable MCP tools.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get cached MCP tools: {e}")
|
||||
|
||||
logger.info(f"Total tools loaded: {len(loaded_tools)}, built-in tools: {len(builtin_tools)}, MCP tools: {len(mcp_tools)}")
|
||||
return loaded_tools + builtin_tools + mcp_tools
|
||||
|
||||
363
backend/tests/test_tool_search.py
Normal file
363
backend/tests/test_tool_search.py
Normal file
@@ -0,0 +1,363 @@
|
||||
"""Tests for the tool_search (deferred tool loading) feature."""
|
||||
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from langchain_core.tools import tool as langchain_tool
|
||||
|
||||
from deerflow.config.tool_search_config import ToolSearchConfig, load_tool_search_config_from_dict
|
||||
from deerflow.tools.builtins.tool_search import (
|
||||
DeferredToolRegistry,
|
||||
get_deferred_registry,
|
||||
reset_deferred_registry,
|
||||
set_deferred_registry,
|
||||
)
|
||||
|
||||
# ── Fixtures ──
|
||||
|
||||
|
||||
def _make_mock_tool(name: str, description: str):
|
||||
"""Create a minimal LangChain tool for testing."""
|
||||
|
||||
@langchain_tool(name)
|
||||
def mock_tool(arg: str) -> str:
|
||||
"""Mock tool."""
|
||||
return f"{name}: {arg}"
|
||||
|
||||
mock_tool.description = description
|
||||
return mock_tool
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registry():
|
||||
"""Create a fresh DeferredToolRegistry with test tools."""
|
||||
reg = DeferredToolRegistry()
|
||||
reg.register(_make_mock_tool("github_create_issue", "Create a new issue in a GitHub repository"))
|
||||
reg.register(_make_mock_tool("github_list_repos", "List repositories for a GitHub user"))
|
||||
reg.register(_make_mock_tool("slack_send_message", "Send a message to a Slack channel"))
|
||||
reg.register(_make_mock_tool("slack_list_channels", "List available Slack channels"))
|
||||
reg.register(_make_mock_tool("sentry_list_issues", "List issues from Sentry error tracking"))
|
||||
reg.register(_make_mock_tool("database_query", "Execute a SQL query against the database"))
|
||||
return reg
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_singleton():
|
||||
"""Reset the module-level singleton before/after each test."""
|
||||
reset_deferred_registry()
|
||||
yield
|
||||
reset_deferred_registry()
|
||||
|
||||
|
||||
# ── ToolSearchConfig Tests ──
|
||||
|
||||
|
||||
class TestToolSearchConfig:
|
||||
def test_default_disabled(self):
|
||||
config = ToolSearchConfig()
|
||||
assert config.enabled is False
|
||||
|
||||
def test_enabled(self):
|
||||
config = ToolSearchConfig(enabled=True)
|
||||
assert config.enabled is True
|
||||
|
||||
def test_load_from_dict(self):
|
||||
config = load_tool_search_config_from_dict({"enabled": True})
|
||||
assert config.enabled is True
|
||||
|
||||
def test_load_from_empty_dict(self):
|
||||
config = load_tool_search_config_from_dict({})
|
||||
assert config.enabled is False
|
||||
|
||||
|
||||
# ── DeferredToolRegistry Tests ──
|
||||
|
||||
|
||||
class TestDeferredToolRegistry:
|
||||
def test_register_and_len(self, registry):
|
||||
assert len(registry) == 6
|
||||
|
||||
def test_entries(self, registry):
|
||||
names = [e.name for e in registry.entries]
|
||||
assert "github_create_issue" in names
|
||||
assert "slack_send_message" in names
|
||||
|
||||
def test_search_select_single(self, registry):
|
||||
results = registry.search("select:github_create_issue")
|
||||
assert len(results) == 1
|
||||
assert results[0].name == "github_create_issue"
|
||||
|
||||
def test_search_select_multiple(self, registry):
|
||||
results = registry.search("select:github_create_issue,slack_send_message")
|
||||
names = {t.name for t in results}
|
||||
assert names == {"github_create_issue", "slack_send_message"}
|
||||
|
||||
def test_search_select_nonexistent(self, registry):
|
||||
results = registry.search("select:nonexistent_tool")
|
||||
assert results == []
|
||||
|
||||
def test_search_plus_keyword(self, registry):
|
||||
results = registry.search("+github")
|
||||
names = {t.name for t in results}
|
||||
assert names == {"github_create_issue", "github_list_repos"}
|
||||
|
||||
def test_search_plus_keyword_with_ranking(self, registry):
|
||||
results = registry.search("+github issue")
|
||||
assert len(results) == 2
|
||||
# "github_create_issue" should rank higher (has "issue" in name)
|
||||
assert results[0].name == "github_create_issue"
|
||||
|
||||
def test_search_regex_keyword(self, registry):
|
||||
results = registry.search("slack")
|
||||
names = {t.name for t in results}
|
||||
assert "slack_send_message" in names
|
||||
assert "slack_list_channels" in names
|
||||
|
||||
def test_search_regex_description(self, registry):
|
||||
results = registry.search("SQL")
|
||||
assert len(results) == 1
|
||||
assert results[0].name == "database_query"
|
||||
|
||||
def test_search_regex_case_insensitive(self, registry):
|
||||
results = registry.search("GITHUB")
|
||||
assert len(results) == 2
|
||||
|
||||
def test_search_invalid_regex_falls_back_to_literal(self, registry):
|
||||
# "[" is invalid regex, should be escaped and used as literal
|
||||
results = registry.search("[")
|
||||
assert results == []
|
||||
|
||||
def test_search_name_match_ranks_higher(self, registry):
|
||||
# "issue" appears in both github_create_issue (name) and sentry_list_issues (name+desc)
|
||||
results = registry.search("issue")
|
||||
names = [t.name for t in results]
|
||||
# Both should be found (both have "issue" in name)
|
||||
assert "github_create_issue" in names
|
||||
assert "sentry_list_issues" in names
|
||||
|
||||
def test_search_max_results(self):
|
||||
reg = DeferredToolRegistry()
|
||||
for i in range(10):
|
||||
reg.register(_make_mock_tool(f"tool_{i}", f"Tool number {i}"))
|
||||
results = reg.search("tool")
|
||||
assert len(results) <= 5 # MAX_RESULTS = 5
|
||||
|
||||
def test_search_empty_registry(self):
|
||||
reg = DeferredToolRegistry()
|
||||
assert reg.search("anything") == []
|
||||
|
||||
def test_empty_registry_len(self):
|
||||
reg = DeferredToolRegistry()
|
||||
assert len(reg) == 0
|
||||
|
||||
|
||||
# ── Singleton Tests ──
|
||||
|
||||
|
||||
class TestSingleton:
|
||||
def test_default_none(self):
|
||||
assert get_deferred_registry() is None
|
||||
|
||||
def test_set_and_get(self, registry):
|
||||
set_deferred_registry(registry)
|
||||
assert get_deferred_registry() is registry
|
||||
|
||||
def test_reset(self, registry):
|
||||
set_deferred_registry(registry)
|
||||
reset_deferred_registry()
|
||||
assert get_deferred_registry() is None
|
||||
|
||||
|
||||
# ── tool_search Tool Tests ──
|
||||
|
||||
|
||||
class TestToolSearchTool:
|
||||
def test_no_registry(self):
|
||||
from deerflow.tools.builtins.tool_search import tool_search
|
||||
|
||||
result = tool_search.invoke({"query": "github"})
|
||||
assert result == "No deferred tools available."
|
||||
|
||||
def test_no_match(self, registry):
|
||||
from deerflow.tools.builtins.tool_search import tool_search
|
||||
|
||||
set_deferred_registry(registry)
|
||||
result = tool_search.invoke({"query": "nonexistent_xyz_tool"})
|
||||
assert "No tools found matching" in result
|
||||
|
||||
def test_returns_valid_json(self, registry):
|
||||
from deerflow.tools.builtins.tool_search import tool_search
|
||||
|
||||
set_deferred_registry(registry)
|
||||
result = tool_search.invoke({"query": "select:github_create_issue"})
|
||||
parsed = json.loads(result)
|
||||
assert isinstance(parsed, list)
|
||||
assert len(parsed) == 1
|
||||
assert parsed[0]["name"] == "github_create_issue"
|
||||
|
||||
def test_returns_openai_function_format(self, registry):
|
||||
from deerflow.tools.builtins.tool_search import tool_search
|
||||
|
||||
set_deferred_registry(registry)
|
||||
result = tool_search.invoke({"query": "select:slack_send_message"})
|
||||
parsed = json.loads(result)
|
||||
func_def = parsed[0]
|
||||
# OpenAI function format should have these keys
|
||||
assert "name" in func_def
|
||||
assert "description" in func_def
|
||||
assert "parameters" in func_def
|
||||
|
||||
def test_keyword_search_returns_json(self, registry):
|
||||
from deerflow.tools.builtins.tool_search import tool_search
|
||||
|
||||
set_deferred_registry(registry)
|
||||
result = tool_search.invoke({"query": "github"})
|
||||
parsed = json.loads(result)
|
||||
assert len(parsed) == 2
|
||||
names = {d["name"] for d in parsed}
|
||||
assert names == {"github_create_issue", "github_list_repos"}
|
||||
|
||||
|
||||
# ── Prompt Section Tests ──
|
||||
|
||||
|
||||
class TestDeferredToolsPromptSection:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _mock_app_config(self, monkeypatch):
|
||||
"""Provide a minimal AppConfig mock so tests don't need config.yaml."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from deerflow.config.tool_search_config import ToolSearchConfig
|
||||
|
||||
mock_config = MagicMock()
|
||||
mock_config.tool_search = ToolSearchConfig() # disabled by default
|
||||
monkeypatch.setattr("deerflow.config.get_app_config", lambda: mock_config)
|
||||
|
||||
def test_empty_when_disabled(self):
|
||||
from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section
|
||||
|
||||
# tool_search.enabled defaults to False
|
||||
section = get_deferred_tools_prompt_section()
|
||||
assert section == ""
|
||||
|
||||
def test_empty_when_enabled_but_no_registry(self, monkeypatch):
|
||||
from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
monkeypatch.setattr(get_app_config().tool_search, "enabled", True)
|
||||
section = get_deferred_tools_prompt_section()
|
||||
assert section == ""
|
||||
|
||||
def test_empty_when_enabled_but_empty_registry(self, monkeypatch):
|
||||
from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
monkeypatch.setattr(get_app_config().tool_search, "enabled", True)
|
||||
set_deferred_registry(DeferredToolRegistry())
|
||||
section = get_deferred_tools_prompt_section()
|
||||
assert section == ""
|
||||
|
||||
def test_lists_tool_names(self, registry, monkeypatch):
|
||||
from deerflow.agents.lead_agent.prompt import get_deferred_tools_prompt_section
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
monkeypatch.setattr(get_app_config().tool_search, "enabled", True)
|
||||
set_deferred_registry(registry)
|
||||
section = get_deferred_tools_prompt_section()
|
||||
assert "<available-deferred-tools>" in section
|
||||
assert "</available-deferred-tools>" in section
|
||||
assert "github_create_issue" in section
|
||||
assert "slack_send_message" in section
|
||||
assert "sentry_list_issues" in section
|
||||
# Should only have names, no descriptions
|
||||
assert "Create a new issue" not in section
|
||||
|
||||
|
||||
# ── DeferredToolFilterMiddleware Tests ──
|
||||
|
||||
|
||||
class TestDeferredToolFilterMiddleware:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _ensure_middlewares_package(self):
|
||||
"""Remove mock entries injected by test_subagent_executor.py.
|
||||
|
||||
That file replaces deerflow.agents and deerflow.agents.middlewares with
|
||||
MagicMock objects in sys.modules (session-scoped) to break circular imports.
|
||||
We must clear those mocks so real submodule imports work.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
mock_keys = [
|
||||
"deerflow.agents",
|
||||
"deerflow.agents.middlewares",
|
||||
"deerflow.agents.middlewares.deferred_tool_filter_middleware",
|
||||
]
|
||||
for key in mock_keys:
|
||||
if isinstance(sys.modules.get(key), MagicMock):
|
||||
del sys.modules[key]
|
||||
|
||||
def test_filters_deferred_tools(self, registry):
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
|
||||
set_deferred_registry(registry)
|
||||
middleware = DeferredToolFilterMiddleware()
|
||||
|
||||
# Build a mock tools list: 2 active + 1 deferred
|
||||
active_tool = _make_mock_tool("my_active_tool", "An active tool")
|
||||
deferred_tool = registry.entries[0].tool # github_create_issue
|
||||
|
||||
class FakeRequest:
|
||||
def __init__(self, tools):
|
||||
self.tools = tools
|
||||
|
||||
def override(self, **kwargs):
|
||||
return FakeRequest(kwargs.get("tools", self.tools))
|
||||
|
||||
request = FakeRequest(tools=[active_tool, deferred_tool])
|
||||
filtered = middleware._filter_tools(request)
|
||||
|
||||
assert len(filtered.tools) == 1
|
||||
assert filtered.tools[0].name == "my_active_tool"
|
||||
|
||||
def test_no_op_when_no_registry(self):
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
|
||||
middleware = DeferredToolFilterMiddleware()
|
||||
active_tool = _make_mock_tool("my_tool", "A tool")
|
||||
|
||||
class FakeRequest:
|
||||
def __init__(self, tools):
|
||||
self.tools = tools
|
||||
|
||||
def override(self, **kwargs):
|
||||
return FakeRequest(kwargs.get("tools", self.tools))
|
||||
|
||||
request = FakeRequest(tools=[active_tool])
|
||||
filtered = middleware._filter_tools(request)
|
||||
|
||||
assert len(filtered.tools) == 1
|
||||
assert filtered.tools[0].name == "my_tool"
|
||||
|
||||
def test_preserves_dict_tools(self, registry):
|
||||
"""Dict tools (provider built-ins) should not be filtered."""
|
||||
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
||||
|
||||
set_deferred_registry(registry)
|
||||
middleware = DeferredToolFilterMiddleware()
|
||||
|
||||
dict_tool = {"type": "function", "function": {"name": "some_builtin"}}
|
||||
active_tool = _make_mock_tool("my_active_tool", "Active")
|
||||
|
||||
class FakeRequest:
|
||||
def __init__(self, tools):
|
||||
self.tools = tools
|
||||
|
||||
def override(self, **kwargs):
|
||||
return FakeRequest(kwargs.get("tools", self.tools))
|
||||
|
||||
request = FakeRequest(tools=[dict_tool, active_tool])
|
||||
filtered = middleware._filter_tools(request)
|
||||
|
||||
# dict_tool has no .name attr → getattr returns None → not in deferred_names → kept
|
||||
assert len(filtered.tools) == 2
|
||||
@@ -12,7 +12,7 @@
|
||||
# ============================================================================
|
||||
# Bump this number when the config schema changes.
|
||||
# Run `make config-upgrade` to merge new fields into your local config.yaml.
|
||||
config_version: 1
|
||||
config_version: 2
|
||||
|
||||
# ============================================================================
|
||||
# Models Configuration
|
||||
@@ -224,6 +224,18 @@ tools:
|
||||
group: bash
|
||||
use: deerflow.sandbox.tools:bash_tool
|
||||
|
||||
# ============================================================================
|
||||
# Tool Search Configuration (Deferred Tool Loading)
|
||||
# ============================================================================
|
||||
# When enabled, MCP tools are not loaded into the agent's context directly.
|
||||
# Instead, they are listed by name in the system prompt and discoverable
|
||||
# via the tool_search tool at runtime.
|
||||
# This reduces context usage and improves tool selection accuracy when
|
||||
# multiple MCP servers expose a large number of tools.
|
||||
|
||||
tool_search:
|
||||
enabled: false
|
||||
|
||||
# ============================================================================
|
||||
# Sandbox Configuration
|
||||
# ============================================================================
|
||||
|
||||
Reference in New Issue
Block a user