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 python_a2a import MCPClient, to_langchain_tool

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)


# 创建一个异步函数，来实现客户端的创建及使用
async def main():
    # 创建MCP客户端
    url = "http://localhost:8010"
    mcp_client = MCPClient(server_url=url)

    try:
        # 获取可用的工具列表
        tools = await mcp_client.get_tools()
        print(f'MCP tools-->{tools}')

        # 将 MCP tool 转成 LangChain 的工具
        get_weather_tool = to_langchain_tool(url, "get_weather")
        query_high_frequency_question = to_langchain_tool(url, "query_high_frequency_question")
        tools = [get_weather_tool, query_high_frequency_question]

        # 创建prompt模板
        # agent_scratchpad 这个参数是agent在进行推理的时候，自动生成的，用于记录agent的推理过程。
        prompt_template = ChatPromptTemplate.from_messages([
            ("system", "你是一个乐于助人的助手，能够调用工具回答用户问题。"),
            ("human", "{user_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({"user_input": query})
                print(f"response-->{response}")
            except Exception:
                print("解析有问题")
    except Exception as e:
        print(f"Error: {e}")


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