# mcp_server.py
from fastmcp import FastMCP

mcp = FastMCP("demo-mcp-server")


# --- Tool Example ---
# 定义工具 1: 加法
@mcp.tool()
def add(a: int, b: int) -> int:
    """Add two integers together."""
    return a + b

# 定义工具 2: 乘法
@mcp.tool()
def multiply(a: int, b: int) -> int:
    """Multiply two integers together."""
    return a * b

# 定义工具 3: 减法
@mcp.tool()
def subtract(a: int, b: int) -> int:
    """Subtract b from a."""
    return a - b

# 定义工具 4: 除法
@mcp.tool()
def divide(a: float, b: float) -> float:
    """Divide a by b."""
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b


# --- Prompt Example ---

@mcp.prompt
def summarize(text: str) -> str:
    """Prompt: ask the LLM to summarize a given text."""
    return (
        f"请对以下文本进行简要总结，控制在 3 句话以内：\n\n{text}"
    )


if __name__ == "__main__":
    mcp.run(transport="streamable-http", host="127.0.0.1", port=8081)

# agent_demo.py
import asyncio

import httpx
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain.agents import create_agent
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field


class MathResult(BaseModel):
    """数学运算结果"""
    expression: str = Field(..., description="计算表达式")
    result: int = Field(..., description="计算结果")
    step: str = Field(..., description="计算步骤说明")


def _make_mcp_http_client(
    headers: dict | None  = None,
    timeout: httpx.Timeout | None = None,
    auth: httpx.Auth | None = None
):
    "避免使用代理和ssl验证客户端，必须有headers，timeout， auth这三个参数，否则会报错"
    return httpx.AsyncClient(trust_env=False, proxy=None, verify=False, timeout=100, headers=headers, auth=auth)


def mcp_client():
    # 配置 MCP 客户端 - 使用 streamable_http 传输
    # 注意：transport 参数可以设置为 "streamable_http" 或简写为 "http"
    # MultiServerMCPClient 支持四种传输协议：stdio、sse、streamable_http、websocket[reference:2]
    client = MultiServerMCPClient(
        {
            "math_server": {
                "url": "http://localhost:8081/mcp",  # MCP 服务器端点
                "transport": "streamable_http",  # 使用 streamable HTTP 传输
                "httpx_client_factory": _make_mcp_http_client
            }
        },
        tool_name_prefix=True,  # 为工具添加服务器前缀以区分来源
    )
    return client


async def creat_agent_():
    client = mcp_client()
    tools = await client.get_tools()  # 从 MCP 服务器加载工具
    print(f"Loaded {len(tools)} tools from MCP server:")
    for tool in tools:
        print(f"  - {tool.name}: {tool.description}")

    llm = ChatOpenAI(
        model="qwen2.5",
        base_url="http://xxxx/v1",
        default_headers={  # 必须设置的请求头参数
            "Authorization": "xxx",
            "Content-Type": "application/json"
        },
        http_async_client=httpx.AsyncClient(trust_env=False, proxy=None, verify=False, timeout=30),
        api_key="EMPTY",  # 必须得有，否则报错，因为vllm启动模型时，没有设置api—key，此处访问时必须设置为EMPTY
        temperature=0.7,
        extra_body={"chat_template_kwargs": {"enable_thinking": False}},  # 额外body体参数，此处表示不启用思考
    )

    # 4. 创建 ReAct Agent
    agent = create_agent(
        llm,
        tools=tools,
        response_format=MathResult  # 没起作用，不清楚为什么？
    )
    return agent


async def main():
    agent = await creat_agent_()

    query = "What is 15 plus 27?"

    # 1. 获取 Agent 响应
    response = await agent.ainvoke({"messages": [{"role": "user", "content": query}]})

    for msg in response['messages']:
        msg.pretty_print()


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