# file: math_server.py

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("Math")

@mcp.tool()
def add(a: int, b: int) -> int:
    """Add two numbers."""
    return a + b

@mcp.tool()
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b

if __name__ == "__main__":
    # 使用 stdio 作为传输机制，方便与本地客户端集成
    # 客户端只要能启动这个进程并接管 stdin/stdout 就能使用这些工具
    mcp.run(transport="stdio")

# file: ch11_agent_single_mcp.py

import asyncio

from langchain.agents import create_agent
from langchain_deepseek import ChatDeepSeek
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.tools import load_mcp_tools

async def main():
    # 1. 初始化 DeepSeek 模型（LangChain v1，开启 v1 输出）
    model = ChatDeepSeek(
        model="deepseek-reasoner",
        api_key="sk-your-api-key",  # 替换为你的 DeepSeek API 密钥
        temperature=0,
        output_version="v1",
    )

    # 2. 创建 MCP 客户端，只配置一个本地 math 服务器
    client = MultiServerMCPClient(
        {
            "math": {
                "command": "python",          # 启动本地 Python 进程
                "args": ["./math_server.py"], # 与 math_server.py 路径保持一致
                "transport": "stdio",         # 本地开发推荐 stdio，零网络依赖
            }
        }
    )

    # 3. 建立与 math 服务器的会话，并加载工具
    async with client.session("math") as session:
        # 自动把 MCP 工具描述转换为 LangChain v1 的 Tool 对象
        tools = await load_mcp_tools(session)

        # 4. 用 create_agent 组装 LangChain v1 Agent
        agent = create_agent(
            model=model,
            tools=tools,
            system_prompt=(
                "你是一个数学助理，可以使用 MCP 提供的数学工具进行计算，"
                "在必要时给出一步一步的推导过程。"
            ),
        )

        # 5. 测试一次调用
        result = await agent.ainvoke(
            {"messages": [{"role": "user", "content": "计算 (3 + 5) × 12 并解释步骤"}]}
        )

        for msg in result["messages"]:
            print(msg)

if __name__ == "__main__":
    asyncio.run(main())

# file: weather_server.py

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("Weather")

@mcp.tool()
async def get_weather(location: str) -> str:
    """Get weather for location."""
    # 教学示例，真实场景可以接第三方天气 API
    return f"The weather in {location} is always sunny in this demo."

if __name__ == "__main__":
    # 这里可以用 SSE 或 stdio，示例中用 stdio 简化本地开发
    mcp.run(transport="stdio")

# file: ch11_agent_multi_mcp.py

import asyncio

from langchain.agents import create_agent
from langchain_deepseek import ChatDeepSeek
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.tools import load_mcp_tools

async def main():
    model = ChatDeepSeek(
        model="deepseek-chat",
        api_key="sk-your-api-key",
        temperature=0.3,
        output_version="v1",
    )

    client = MultiServerMCPClient(
        {
            "math": {
                "command": "python",            # 使用同一套二进制启动不同 MCP 服务
                "args": ["./math_server.py"],   # 数学工具服务器
                "transport": "stdio",
            },
            "weather": {
                "command": "python",
                "args": ["./weather_server.py"], # 天气查询服务器
                "transport": "stdio",
            },
        }
    )

    # 演示：从多个服务器加载工具并合并
    all_tools = []
    async with client.session("math") as math_session:
        all_tools.extend(await load_mcp_tools(math_session))

    async with client.session("weather") as weather_session:
        all_tools.extend(await load_mcp_tools(weather_session))

    agent = create_agent(
        model=model,
        tools=all_tools,
        system_prompt=(
            "你可以使用 MCP 工具进行数学计算和查询天气。"
            "在回答问题时，优先利用可用工具给出精确结果。"
        ),
    )

    # 测试数学问题
    math_result = await agent.ainvoke(
        {"messages": [{"role": "user", "content": "what's (3 + 5) x 12?"}]}
    )
    print("数学问题回答：")
    print(math_result["messages"][-1].content)

    # 测试天气问题
    weather_result = await agent.ainvoke(
        {"messages": [{"role": "user", "content": "what is the weather in New York?"}]}
    )
    print("天气问题回答：")
    print(weather_result["messages"][-1].content)

if __name__ == "__main__":
    asyncio.run(main())

# file: ch11_agent_mcp_hitl.py

import asyncio

from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langchain_deepseek import ChatDeepSeek
from langchain_core.runnables import RunnableConfig
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.tools import load_mcp_tools
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command

async def main():
    model = ChatDeepSeek(
        model="deepseek-reasoner",
        api_key="sk-your-api-key",
        temperature=0,
        output_version="v1",
    )

    client = MultiServerMCPClient(
        {
            "math": {
                "command": "python",
                "args": ["./math_server.py"],
                "transport": "stdio",
            },
            "shell": {
                "command": "python",
                "args": ["./shell_server.py"],  # 假设有一个执行 shell 命令的 MCP 服务器
                "transport": "stdio",
            },
        }
    )

    tools = []
    async with client.session("math") as math_session:
        tools.extend(await load_mcp_tools(math_session))
    async with client.session("shell") as shell_session:
        tools.extend(await load_mcp_tools(shell_session))

    # LangGraph 内存 checkpointer，支持线程级状态持久化
    checkpointer = InMemorySaver()

    agent = create_agent(
        model=model,
        tools=tools,
        system_prompt=(
            "你可以使用 MCP 工具执行数学运算和系统命令。"
            "对于可能修改系统状态的命令，必须在获得人类批准后才能执行。"
        ),
        middleware=[
            HumanInTheLoopMiddleware(
                interrupt_on={
                    # 数学工具完全自动化
                    "add": False,
                    "multiply": False,
                    # shell 相关工具一律需要审批
                    "run_shell_command": True,
                }
            )
        ],
        checkpointer=checkpointer,
    )

    thread_id = "mcp-hitl-demo-001"  # 同一 thread_id 下的多次调用共享一份 Agent 状态
    config: RunnableConfig = {"configurable": {"thread_id": thread_id}}

    # 第一次调用：Agent 可能在准备执行 run_shell_command 时中断
    first_result = await agent.ainvoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "请帮我查看当前目录下有哪些文件，如果必要可以执行系统命令。",
                }
            ]
        },
        config=config,
    )

    print("第一次调用最后一条消息：")
    print(first_result["messages"][-1])

    if not first_result.get("interrupts"):
        print("未触发任何中断。")
        return

    interrupt_info = first_result["interrupts"][0]
    print("需要审批的工具调用：", interrupt_info)

    # 在真实场景中，这里应当由人类审阅工具名和参数后再决定
    decisions = [{"type": "approve"}]

    resume_cmd = Command(resume={"decisions": decisions})

    resumed_result = await agent.ainvoke(resume_cmd, config=config)

    print("恢复执行后的最终回答：")
    print(resumed_result["messages"][-1])

if __name__ == "__main__":
    asyncio.run(main())