from mcp.client.streamable_http import streamablehttp_client
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_openai import ChatOpenAI
from mcp import ClientSession
import asyncio

from agent_learn.config import Config

conf = Config()

# 创建模型
llm = ChatOpenAI(base_url=conf.base_url,
                 api_key=conf.api_key,
                 model=conf.model_name,
                 temperature=0.1)

# MCP Server URL
server_url = "http://localhost:8001/mcp"

# 定义mcp客户端
mcp_client = None

# 创建一个异步函数，来实现客户端的创建及使用
async def main():
    global mcp_client
    # 启动 MCP server，并通过标准输入输出建立异步连接。
    async with streamablehttp_client(url=server_url) as (read, write, _):
        # 使用读写流创建MCP会话对象
        async with ClientSession(read, write) as session:
            # 初始化session
            await session.initialize()

            # 动态创建一个临时类 MCPClientHolder，把 session 放进去。这样就可以在函数外部通过 mcp_client.session 调用 MCP 工具
            mcp_client = type("MCPClientHolder", (), {"session": session})()

            # 从 session 自动获取 MCP server 提供的工具列表。
            tools = await load_mcp_tools(session)
            print(f"tools-->{tools}")

            # 创建prompt模板
            # agent_scratchpad 这个参数是agent在进行推理的时候，自动生成的，用于记录agent的推理过程。
            prompt_template = ChatPromptTemplate.from_messages([
                ("system", "你是一个乐于助人的助手，能够调用工具回答用户问题。"),
                ("human", "{input}"),
                ("placeholder", "{agent_scratchpad}"),
            ])

            # 创建agent对象
            agent = create_tool_calling_agent(llm, tools, prompt_template)
            # 使用agent对象创建agent执行器
            agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

            # agent调用
            print("MCP客户端启动，输入'quit'退出")
            while True:
                # 接收用户查询
                query = input("\nQuery: ").strip()
                if query.lower() == "quit":
                    break
                try:
                    # 发送用户查询给代理，并打印  ainvoke 指的是异步调用
                    response = await agent_executor.ainvoke({"input": query})
                    print(f"response-->{response}")
                except Exception:
                    print("解析有问题")

    return # 终止进程

# 启动运行
if __name__ == '__main__':
    asyncio.run(main())