import asyncio
import logging

# from langchain.agents import create_react_agent
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_mcp_adapters.client import MultiServerMCPClient
import os
import dotenv

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

dotenv.load_dotenv()

## 设置环境变量
os.environ['OPENAI_API_KEY'] = os.getenv("OLLAMA_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("OLLAMA_BASE_URL")

# 默认的 'model_name': 'deepseek-ai/DeepSeek-V3.1',
llm = ChatOpenAI(model="gemma4:26b")


async def execute():
    # 1. 创建langchain中的mcp客户端
    # uv add langchain_mcp_adapters
    client = MultiServerMCPClient(
        {
            # 这里是定义服务端信息的，可以有多个服务端
            "weather": {
                "url": "http://localhost:9000/mcp",
                "transport": "streamable_http",
            }
        }
    )

    try:
        # 2. 通过客户端获取工具列表
        ## 这里是通过服务端获取，所以可能会有异常（比如服务端没有启动，或者网络连接有问题
        tools = await client.get_tools()

        # 3. 创建一个智能代理，能够完成 思考--> 行动 --> 观察 --> 思考 --> 行动 --> ... --> 最终答案
        agent = create_react_agent(model=llm, tools=tools)
        agent_response = await agent.ainvoke({"messages":"北京的天气怎么样？适合出门吗？"})
        print(agent_response)

    finally:
        # 资源回收
        if hasattr(client, 'close'):
            client.close()


if __name__ == '__main__':
    asyncio.run(execute())