mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-09 16:08:31 +00:00
5ba25b06ec
* feat: add MCP routing hints * test: isolate mcp routing prompt config * fix: address mcp routing review feedback
122 lines
3.6 KiB
Python
122 lines
3.6 KiB
Python
"""Tests for MCP routing metadata tags."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from langchain_core.tools import StructuredTool
|
|
from pydantic import BaseModel, Field
|
|
|
|
from deerflow.config.extensions_config import ExtensionsConfig
|
|
from deerflow.tools.mcp_metadata import MCP_TOOL_METADATA_KEY, MCP_TOOL_ROUTING_METADATA_KEY, get_mcp_routing, tag_mcp_routing, tag_mcp_tool
|
|
|
|
|
|
class _Args(BaseModel):
|
|
query: str = Field(..., description="query")
|
|
|
|
|
|
def _tool(name: str = "postgres_query") -> StructuredTool:
|
|
async def _call(query: str) -> str:
|
|
return query
|
|
|
|
return StructuredTool(
|
|
name=name,
|
|
description="Query internal data",
|
|
args_schema=_Args,
|
|
coroutine=_call,
|
|
)
|
|
|
|
|
|
def test_tag_mcp_routing_preserves_existing_mcp_flag():
|
|
tool = tag_mcp_tool(_tool())
|
|
|
|
tagged = tag_mcp_routing(
|
|
tool,
|
|
{
|
|
"mode": "prefer",
|
|
"priority": 80,
|
|
"keywords": ["订单"],
|
|
},
|
|
)
|
|
|
|
assert tagged.metadata[MCP_TOOL_METADATA_KEY] is True
|
|
assert tagged.metadata[MCP_TOOL_ROUTING_METADATA_KEY]["priority"] == 80
|
|
assert get_mcp_routing(tagged)["keywords"] == ["订单"]
|
|
|
|
|
|
def test_get_mcp_routing_returns_none_for_non_mcp_tools():
|
|
tool = tag_mcp_routing(
|
|
_tool(),
|
|
{
|
|
"mode": "prefer",
|
|
"priority": 80,
|
|
"keywords": ["订单"],
|
|
},
|
|
)
|
|
|
|
assert get_mcp_routing(tool) is None
|
|
|
|
|
|
def test_get_mcp_routing_returns_none_for_off_mode():
|
|
tool = tag_mcp_tool(_tool())
|
|
tag_mcp_routing(
|
|
tool,
|
|
{
|
|
"mode": "off",
|
|
"priority": 80,
|
|
"keywords": ["订单"],
|
|
},
|
|
)
|
|
|
|
assert get_mcp_routing(tool) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("transport", ["http", "stdio"])
|
|
async def test_get_mcp_tools_tags_effective_routing_metadata(transport: str):
|
|
from deerflow.mcp.tools import get_mcp_tools
|
|
|
|
tool = _tool("postgres_query")
|
|
extensions_config = ExtensionsConfig.model_validate(
|
|
{
|
|
"mcpServers": {
|
|
"postgres": {
|
|
"type": transport,
|
|
"url": "http://localhost:8000/mcp",
|
|
"command": "npx",
|
|
"routing": {
|
|
"mode": "prefer",
|
|
"priority": 50,
|
|
"keywords": ["database"],
|
|
},
|
|
"tools": {
|
|
"query": {
|
|
"routing": {
|
|
"priority": 100,
|
|
"keywords": ["查库"],
|
|
}
|
|
}
|
|
},
|
|
}
|
|
}
|
|
}
|
|
)
|
|
|
|
with (
|
|
patch("deerflow.mcp.tools.ExtensionsConfig.from_file", return_value=extensions_config),
|
|
patch(
|
|
"deerflow.mcp.tools.build_servers_config",
|
|
return_value={"postgres": {"transport": transport, "url": "http://localhost:8000/mcp", "command": "npx"}},
|
|
),
|
|
patch("deerflow.mcp.tools.get_initial_oauth_headers", return_value={}),
|
|
patch("deerflow.mcp.tools.build_oauth_tool_interceptor", return_value=None),
|
|
patch("langchain_mcp_adapters.client.MultiServerMCPClient") as MockClient,
|
|
):
|
|
MockClient.return_value.get_tools = AsyncMock(return_value=[tool])
|
|
tools = await get_mcp_tools()
|
|
|
|
routing = get_mcp_routing(tools[0])
|
|
assert routing is not None
|
|
assert routing["priority"] == 100
|
|
assert routing["keywords"] == ["查库"]
|