mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 02:05:45 +00:00
Ryker_FengandGitHub
5ba25b06ec
* feat: add MCP routing hints * test: isolate mcp routing prompt config * fix: address mcp routing review feedback
137 lines
4.6 KiB
Python
137 lines
4.6 KiB
Python
"""Tests for MCP routing hint prompt rendering."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
from langchain_core.tools import StructuredTool
|
|
from langchain_core.utils.function_calling import convert_to_openai_function
|
|
from pydantic import BaseModel, Field
|
|
|
|
from deerflow.agents.lead_agent.prompt import apply_prompt_template
|
|
from deerflow.tools.builtins.tool_search import assemble_deferred_tools, get_mcp_routing_hints_prompt_section
|
|
from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool
|
|
|
|
|
|
class _Args(BaseModel):
|
|
query: str = Field(..., description="query")
|
|
|
|
|
|
def _tool(name: str, description: str = "Query internal data") -> StructuredTool:
|
|
async def _call(query: str) -> str:
|
|
return query
|
|
|
|
return StructuredTool(
|
|
name=name,
|
|
description=description,
|
|
args_schema=_Args,
|
|
coroutine=_call,
|
|
)
|
|
|
|
|
|
def _routed_tool(name: str, *, priority: int, keywords: list[str], mode: str = "prefer") -> StructuredTool:
|
|
tool = tag_mcp_tool(_tool(name))
|
|
tag_mcp_routing(
|
|
tool,
|
|
{
|
|
"mode": mode,
|
|
"priority": priority,
|
|
"keywords": keywords,
|
|
},
|
|
)
|
|
return tool
|
|
|
|
|
|
def _minimal_prompt_app_config() -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
sandbox=SimpleNamespace(mounts=[]),
|
|
skills=SimpleNamespace(container_path="/mnt/skills", get_skills_path=lambda: Path("/tmp/skills")),
|
|
skill_evolution=SimpleNamespace(enabled=False),
|
|
acp_agents={},
|
|
)
|
|
|
|
|
|
def test_zero_mcp_routing_tools_render_empty_section():
|
|
assert get_mcp_routing_hints_prompt_section([]) == ""
|
|
|
|
|
|
def test_off_mode_and_empty_keywords_are_excluded():
|
|
section = get_mcp_routing_hints_prompt_section(
|
|
[
|
|
_routed_tool("postgres_query", priority=100, keywords=["订单"], mode="off"),
|
|
_routed_tool("metrics_query", priority=90, keywords=[]),
|
|
]
|
|
)
|
|
|
|
assert section == ""
|
|
|
|
|
|
def test_routing_hints_are_ordered_by_priority_then_name():
|
|
section = get_mcp_routing_hints_prompt_section(
|
|
[
|
|
_routed_tool("z_tool", priority=50, keywords=["z"]),
|
|
_routed_tool("a_tool", priority=50, keywords=["a"]),
|
|
_routed_tool("top_tool", priority=90, keywords=["top", "SQL"]),
|
|
]
|
|
)
|
|
|
|
assert section.startswith("<mcp_routing_hints>")
|
|
top_index = section.index("`top_tool`")
|
|
a_index = section.index("`a_tool`")
|
|
z_index = section.index("`z_tool`")
|
|
assert top_index < a_index < z_index
|
|
assert "When the user's request involves top, or SQL:" in section
|
|
assert "prefer the `top_tool` tool." in section
|
|
assert "priority" not in section
|
|
|
|
|
|
def test_deferred_routing_hints_use_tool_search_promotion():
|
|
routed = _routed_tool("postgres_query", priority=100, keywords=["订单"])
|
|
_, deferred_setup = assemble_deferred_tools([routed], enabled=True)
|
|
|
|
section = get_mcp_routing_hints_prompt_section([routed], deferred_names=deferred_setup.deferred_names)
|
|
|
|
assert "When the user's request involves 订单:" in section
|
|
assert "use `tool_search` to fetch `postgres_query`, then prefer that MCP tool." in section
|
|
assert "prefer the `postgres_query` tool." not in section
|
|
|
|
|
|
def test_apply_prompt_template_places_routing_hints_after_deferred_tools(monkeypatch):
|
|
section = get_mcp_routing_hints_prompt_section(
|
|
[
|
|
_routed_tool("postgres_query", priority=100, keywords=["订单"]),
|
|
]
|
|
)
|
|
empty_storage = SimpleNamespace(load_skills=lambda *, enabled_only: [])
|
|
monkeypatch.setattr("deerflow.agents.lead_agent.prompt.get_or_new_skill_storage", lambda **kwargs: empty_storage)
|
|
monkeypatch.setattr("deerflow.agents.lead_agent.prompt.get_or_new_user_skill_storage", lambda *args, **kwargs: empty_storage)
|
|
monkeypatch.setattr("deerflow.agents.lead_agent.prompt.get_agent_soul", lambda agent_name=None: "")
|
|
|
|
prompt = apply_prompt_template(
|
|
app_config=_minimal_prompt_app_config(),
|
|
deferred_names=frozenset({"postgres_query"}),
|
|
mcp_routing_hints_section=section,
|
|
)
|
|
|
|
assert "<available-deferred-tools>" in prompt
|
|
assert "<mcp_routing_hints>" in prompt
|
|
assert prompt.index("<available-deferred-tools>") < prompt.index("<mcp_routing_hints>")
|
|
|
|
|
|
def test_routing_metadata_does_not_change_openai_function_schema():
|
|
tool = tag_mcp_tool(_tool("postgres_query"))
|
|
before = convert_to_openai_function(tool)
|
|
|
|
tag_mcp_routing(
|
|
tool,
|
|
{
|
|
"mode": "prefer",
|
|
"priority": 100,
|
|
"keywords": ["订单"],
|
|
},
|
|
)
|
|
after = convert_to_openai_function(tool)
|
|
|
|
assert after == before
|