chendandan_homework/app/mcp_client.py

from __future__ import annotations

"""MCP工具加载器：连接自建旅行MCP与高德MCP两个服务器，校验必需工具存在后打包为MCPToolBundle供Agent使用。"""

import sys
from dataclasses import dataclass

from langchain_core.tools import BaseTool
from langchain_mcp_adapters.client import (
    MultiServerMCPClient,
)

from app.config import get_settings


TRAVEL_TOOL_NAMES = {
    "search_airports",
    "search_flights",
    "search_return_flights",
    "search_hotels",
}


@dataclass
class MCPToolBundle:
    """项目使用的全部MCP工具。"""

    client: MultiServerMCPClient
    all_tools: list[BaseTool]
    travel_tools: list[BaseTool]
    amap_tools: list[BaseTool]


async def load_mcp_tools() -> MCPToolBundle:
    """连接自建旅行MCP和高德MCP。"""

    settings = get_settings()

    client = MultiServerMCPClient(
        {
            "travel_search": {
                # 使用当前虚拟环境的Python启动本地Server。
                "transport": "stdio",
                "command": sys.executable,
                "args": [
                    "-m",
                    "mcp_servers.travel_search_server",
                ],
            },
            "amap": {
                # LangChain配置中写http，
                # 实际协议为Streamable HTTP。
                "transport": "http",
                "url": (
                    "https://mcp.amap.com/mcp"
                    f"?key={settings.require('amap_api_key')}"
                ),
            },
        }
    )

    all_tools = await client.get_tools()

    travel_tools = [
        tool
        for tool in all_tools
        if tool.name in TRAVEL_TOOL_NAMES
    ]

    amap_tools = [
        tool
        for tool in all_tools
        if tool.name.startswith("maps_")
    ]

    loaded_travel_names = {
        tool.name
        for tool in travel_tools
    }

    missing_tools = (
        TRAVEL_TOOL_NAMES
        - loaded_travel_names
    )

    if missing_tools:
        raise RuntimeError(
            "旅行MCP缺少工具："
            + "、".join(sorted(missing_tools))
        )

    if not amap_tools:
        raise RuntimeError(
            "没有从高德MCP加载到任何地图工具。"
        )

    return MCPToolBundle(
        client=client,
        all_tools=all_tools,
        travel_tools=travel_tools,
        amap_tools=amap_tools,
    )