mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-06-09 17:12:01 +00:00
2bbc7879fa
Follow-up to #3342 (deferred MCP tool loading). Maintainability cleanup plus hardening of malformed/empty tool_search queries; no change to the deferral mechanism or search ranking. - Add deerflow/tools/mcp_metadata.py as the single source of truth for the "deerflow_mcp" tag (MCP_TOOL_METADATA_KEY + tag_mcp_tool + public is_mcp_tool). Removes the duplicated magic string and the private, cross-module _is_mcp_tool import. - tool_search.search: never raise on model-generated input. Extract _compile_catalog_regex (shared compile-with-literal-fallback); return empty for empty/whitespace queries and a bare "+" instead of matching everything or raising IndexError. - DeferredToolSetup: document the empty-vs-populated invariant. - build_deferred_tool_setup: comment the two distinct empty-return branches. - _assemble_deferred: add return type, rename local to deferred_setup, build the final list with an explicit append. - Tests: use tag_mcp_tool instead of per-file tag helpers; cover empty and bare-"+" queries.
74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
"""End-to-end: tool_search promotes a deferred tool into the next model turn.
|
|
|
|
Locks the full loop through a real ``create_agent`` graph:
|
|
turn 1 -> deferred MCP tools hidden from bind_tools; model calls tool_search
|
|
ToolNode-> tool_search returns Command(update={"promoted": {...}}) -> state
|
|
turn 2 -> middleware reads state["promoted"] (hash-scoped) -> the searched
|
|
tool's schema is now bound; un-searched deferred tools stay hidden
|
|
|
|
This is the behavior #3272's redesign depends on (no ContextVar): promotion
|
|
flows through graph state, so it works regardless of build/execute context.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
from langchain_core.tools import tool as as_tool
|
|
|
|
from deerflow.agents.middlewares.deferred_tool_filter_middleware import DeferredToolFilterMiddleware
|
|
from deerflow.agents.thread_state import ThreadState
|
|
from deerflow.tools.builtins.tool_search import build_deferred_tool_setup
|
|
from deerflow.tools.mcp_metadata import tag_mcp_tool
|
|
|
|
|
|
@as_tool
|
|
def active_tool(x: str) -> str:
|
|
"An always-active tool."
|
|
return x
|
|
|
|
|
|
@as_tool
|
|
def mcp_calc(expression: str) -> str:
|
|
"Evaluate arithmetic."
|
|
return expression
|
|
|
|
|
|
@as_tool
|
|
def mcp_other(x: str) -> str:
|
|
"Another deferred MCP tool."
|
|
return x
|
|
|
|
|
|
def test_tool_search_promotes_into_next_turn():
|
|
bound: list[list[str]] = []
|
|
|
|
class RecordingModel(GenericFakeChatModel):
|
|
def bind_tools(self, tools, **kwargs):
|
|
bound.append([getattr(t, "name", None) for t in tools])
|
|
return self
|
|
|
|
setup = build_deferred_tool_setup([active_tool, tag_mcp_tool(mcp_calc), tag_mcp_tool(mcp_other)], enabled=True)
|
|
turn1 = AIMessage(content="", tool_calls=[{"name": "tool_search", "args": {"query": "select:mcp_calc"}, "id": "c1", "type": "tool_call"}])
|
|
turn2 = AIMessage(content="done")
|
|
model = RecordingModel(messages=iter([turn1, turn2]))
|
|
|
|
graph = create_agent(
|
|
model=model,
|
|
tools=[active_tool, mcp_calc, mcp_other, setup.tool_search_tool],
|
|
middleware=[DeferredToolFilterMiddleware(setup.deferred_names, setup.catalog_hash)],
|
|
state_schema=ThreadState,
|
|
)
|
|
|
|
result = asyncio.run(graph.ainvoke({"messages": [HumanMessage(content="use the deferred calculator")]}))
|
|
|
|
assert len(bound) >= 2, f"expected >=2 model binds, got {bound}"
|
|
# Turn 1: both deferred MCP tools hidden.
|
|
assert "mcp_calc" not in bound[0] and "mcp_other" not in bound[0]
|
|
# Turn 2: the searched tool is promoted (visible); the un-searched one stays hidden.
|
|
assert "mcp_calc" in bound[1]
|
|
assert "mcp_other" not in bound[1]
|
|
# Promotion recorded in graph state, scoped by catalog hash.
|
|
assert result["promoted"] == {"catalog_hash": setup.catalog_hash, "names": ["mcp_calc"]}
|