Aloong

class MCPClient:
    def __init__(self):
        self.max_loop = 10  # 防止无限循环的最大对话轮次限制
        self.messages = []  # 存储完整的对话历史，用于保持上下文
        self.exit_stack = AsyncExitStack()  # 异步资源管理器，确保所有连接正确关闭
        self.sessions: List[ClientSession] = []  # 存储所有MCP服务器会话
        self.available_tools: List[dict] = []  # 存储转换为OpenAI格式的工具列表
        self.tool_session_mapping: Dict[str, ClientSession] = {}  # 工具名称到会话的映射，用于路由工具调用
        # 初始化DeepSeek客户端，从环境变量获取API密钥
        self.client = AsyncOpenAI(api_key=os.getenv("DEEPSEEK_API_KEY"), base_url="https://api.deepseek.com")

async def connect_to_server(self, server_name: str, server_config: dict):
    """连接单个MCP服务器并获取其工具列表"""
    try:
        # 根据配置创建服务器参数
        server_params = StdioServerParameters(
            command=server_config["command"],  # 服务器启动命令
            args=server_config["args"],        # 命令行参数
            env=server_config.get("env"),      # 环境变量（可选）
        )

        # 通过stdio协议建立与MCP服务器的双向通信
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        stdio, write = stdio_transport
        # 创建客户端会话
        session = await self.exit_stack.enter_async_context(ClientSession(stdio, write))
        self.sessions.append(session)
        logger.debug(f"Connected to server: {server_name}")

        # 初始化会话，建立协议握手
        await session.initialize()

        # 获取服务器提供的工具列表
        response = await session.list_tools()
        logger.debug(f"Tool Response: {response}")
        
        # 将MCP工具转换为OpenAI函数调用格式
        for tool in response.tools:
            self.tool_session_mapping[tool.name] = session  # 建立工具到会话的映射
            self.available_tools.append({
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.inputSchema  # MCP的inputSchema直接兼容OpenAI的parameters格式
                }
            })

    except Exception as e:
        logger.error(f"Error connecting to server {server_name}: {str(e)}")
        raise

async def connect_to_servers(self):
    """从servers.json读取配置并连接所有MCP服务器"""
    try:
        # 读取服务器配置文件
        with open("servers.json", "r") as f:
            data = json.load(f)
        servers = data.get("mcpServers", {})

        # 遍历配置文件中的所有服务器，逐一建立连接
        for server_name, server_config in servers.items():
            await self.connect_to_server(server_name, server_config)

    except FileNotFoundError:
        logger.error("No servers.json file found")
        raise
    except Exception as e:
        logger.error(f"Error connecting to servers: {str(e)}")
        raise

async def process_query(self, query: str, system_prompt: Optional[str] = None) -> None:
    """处理用户查询，实现LLM与工具调用的完整循环"""
    
    # 如果有系统提示词，先添加到消息历史
    if system_prompt:
        self.messages.append({"role": "system", "content": system_prompt})

    # 添加用户查询到消息历史
    self.messages.append({"role": "user", "content": query})

    # 开始LLM与工具调用的循环，最多执行max_loop次防止无限循环
    for i in range(self.max_loop):
        try:
            # 调用DeepSeek API，传入对话历史和可用工具列表
            response = await self.client.chat.completions.create(
                model="deepseek-chat",
                messages=self.messages,
                tools=self.available_tools,  # 传入所有可用的MCP工具
                tool_choice="auto"          # 让AI自动判断是否需要调用工具
            )
            logger.debug(f"Chat Completion Response: {response}")
        except Exception as e:
            logger.error(f"Error processing query: {str(e)}")
            raise

        # 获取AI的响应消息
        response_message = response.choices[0].message
        # 将AI响应添加到消息历史
        self.messages.append({
            "role": "assistant",
            "content": response_message.content,
            "tool_calls": response_message.tool_calls
        })

        # 如果AI不需要调用工具，直接输出结果并结束循环
        if not response_message.tool_calls:
            print("AI: ", response_message.content)
            break

        # AI决定调用工具，遍历所有工具调用请求
        for tool_call in response_message.tool_calls:
            tool_name = tool_call.function.name
            # 解析工具调用参数（OpenAI返回JSON字符串，需要转换为字典）
            try:
                arguments = json.loads(tool_call.function.arguments)
            except json.JSONDecodeError:
                arguments = {}
            print(f"Calling tool: [{tool_name}] with arguments: {arguments}")

            try:
                # 根据工具名称找到对应的MCP会话
                session = self.tool_session_mapping[tool_name]
                # 调用MCP工具
                mcp_result = await session.call_tool(
                    name=tool_name,
                    arguments=arguments
                )
                logger.debug(f"MCP Tool call Result: {mcp_result}")
                
                # 处理工具调用结果，转换为字符串格式
                if hasattr(mcp_result, 'content'):
                    if isinstance(mcp_result.content, list):
                        result_content = "\n".join([str(item) for item in mcp_result.content])
                    else:
                        result_content = str(mcp_result.content)
                else:
                    result_content = str(mcp_result)

                # 将工具调用结果添加到消息历史，供AI下一轮使用
                self.messages.append({
                    "role": "tool",
                    "name": tool_name,
                    "content": result_content,
                    "tool_call_id": tool_call.id
                })
            except Exception as e:
                # 工具调用失败时，将错误信息添加到消息历史
                error_msg = f"Error: {str(e)}"
                logger.error(error_msg)
                self.messages.append({
                    "role": "tool",
                    "name": tool_name,
                    "content": error_msg,
                    "tool_call_id": tool_call.id
                })

async def chat_loop(self):
    """运行交互式聊天循环，处理用户输入"""
    print("\nMCP Client Started!")
    print("Type 'exit' or 'quit' to quit.\n")

    while True:
        try:
            # 获取用户输入
            query = input("\nYou: ").strip()
            if query.lower() == "exit" or query.lower() == "quit":
                break  # 用户选择退出
            if not query:
                continue  # 忽略空输入
            
            # 调试模式下打印消息历史
            if self.messages:
                logger.debug(f"Messages[{len(self.messages)}]:")
                for message in self.messages:
                    logger.debug(message)
            
            # 处理用户查询
            await self.process_query(query)
        except Exception as e:
            logger.error(f"Error: {str(e)}")

async def cleanup(self):
    """清理所有异步资源，确保连接正确关闭"""
    await self.exit_stack.aclose()

async def main():
    """程序主入口，管理整个生命周期"""
    client = MCPClient()
    try:
        # 初始化阶段：连接所有配置的MCP服务器
        await client.connect_to_servers()
        # 交互阶段：启动聊天循环
        await client.chat_loop()
    except KeyboardInterrupt:
        print("Goodbye!")
    finally:
        # 确保资源正确释放
        await client.cleanup()

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

import asyncio
import json
import os
import logging
from contextlib import AsyncExitStack
from typing import Optional, List, Dict

import nest_asyncio
from dotenv import load_dotenv
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from openai import AsyncOpenAI

nest_asyncio.apply()
load_dotenv()

# Configure logging with basicConfig
logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "WARNING").upper(),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class MCPClient:
    def __init__(self):
        self.max_loop = 10
        self.messages = []
        self.exit_stack = AsyncExitStack()
        self.sessions: List[ClientSession] = []
        self.available_tools: List[dict] = []
        self.tool_session_mapping: Dict[str, ClientSession] = {}
        self.client = AsyncOpenAI(api_key=os.getenv("DEEPSEEK_API_KEY"), base_url="https://api.deepseek.com")

    async def connect_to_server(self, server_name: str, server_config: dict):
        """
        Connect to an MCP server

        Args:
            server_name: Name of the server
            server_config: Configuration for the server

        :return:
        """
        try:
            server_params = StdioServerParameters(
                command=server_config["command"],
                args=server_config["args"],  # Command line arguments
                env=server_config.get("env"),  # Optional environment variables
            )

            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
            stdio, write = stdio_transport
            session = await self.exit_stack.enter_async_context(ClientSession(stdio, write))
            self.sessions.append(session)
            logger.debug(f"Connected to server: {server_name}")

            await session.initialize()

            # list available tools
            response = await session.list_tools()
            logger.debug(f"Tool Response: {response}")
            for tool in response.tools:
                self.tool_session_mapping[tool.name] = session
                self.available_tools.append({
                    "type": "function",
                    "function": {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.inputSchema
                    }
                })

        except Exception as e:
            logger.error(f"Error connecting to server {server_name}: {str(e)}")
            raise

    async def connect_to_servers(self):
        """
        Connect to multiple MCP servers
        """
        try:
            with open("servers.json", "r") as f:
                data = json.load(f)
            servers = data.get("mcpServers", {})

            for server_name, server_config in servers.items():
                await self.connect_to_server(server_name, server_config)

        except FileNotFoundError:
            logger.error("No servers.json file found")
            raise
        except Exception as e:
            logger.error(f"Error connecting to servers: {str(e)}")
            raise

    async def process_query(self, query: str, system_prompt: Optional[str] = None) -> None:
        """
        Process a query using DeepSeek and available MCP tools

        :param query: The query user wants to process
        :param system_prompt: Optional system prompt
        :return:
        """

        if system_prompt:
            self.messages.append({"role": "system", "content": system_prompt})

        self.messages.append({"role": "user", "content": query})

        for i in range(self.max_loop):
            try:
                response = await self.client.chat.completions.create(
                    model="deepseek-chat",
                    messages=self.messages,
                    tools=self.available_tools,
                    tool_choice="auto"
                )
                logger.debug(f"Chat Completion Response: {response}")
            except Exception as e:
                logger.error(f"Error processing query: {str(e)}")
                raise

            response_message = response.choices[0].message
            self.messages.append({
                "role": "assistant",
                "content": response_message.content,
                "tool_calls": response_message.tool_calls
            })

            if not response_message.tool_calls:
                print("AI: ", response_message.content)
                break

            # If deepseek determines that it needs to call a tool
            for tool_call in response_message.tool_calls:
                tool_name = tool_call.function.name
                # for openai tools, arguments is a json string, mcp expects a dict
                try:
                    arguments = json.loads(tool_call.function.arguments)
                except json.JSONDecodeError:
                    arguments = {}
                print(f"Calling tool: [{tool_name}] with arguments: {arguments}")

                try:
                    session = self.tool_session_mapping[tool_name]
                    mcp_result = await session.call_tool(
                        name=tool_name,
                        arguments=arguments
                    )
                    logger.debug(f"MCP Tool call Result: {mcp_result}")
                    if hasattr(mcp_result, 'content'):
                        if isinstance(mcp_result.content, list):
                            result_content = "\n".join([str(item) for item in mcp_result.content])
                        else:
                            result_content = str(mcp_result.content)
                    else:
                        result_content = str(mcp_result)

                    self.messages.append({
                        "role": "tool",
                        "name": tool_name,
                        "content": result_content,
                        "tool_call_id": tool_call.id
                    })
                except Exception as e:
                    error_msg = f"Error: {str(e)}"
                    logger.error(error_msg)
                    self.messages.append({
                        "role": "tool",
                        "name": tool_name,
                        "content": error_msg,
                        "tool_call_id": tool_call.id
                    })


    async def chat_loop(self):
        """
        Run an interactive chat loop
        """
        print("\nMCP Client Started!")
        print("Type 'exit' or 'quit' to quit.\n")

        while True:
            try:
                query = input("\nYou: ").strip()
                if query.lower() == "exit" or query.lower() == "quit":
                    break
                if not query:
                    continue
                if self.messages:
                    logger.debug(f"Messages[{len(self.messages)}]:")
                    for message in self.messages:
                        logger.debug(message)
                await self.process_query(query)
            except Exception as e:
                logger.error(f"Error: {str(e)}")

    async def cleanup(self):
        await self.exit_stack.aclose()


async def main():
    client = MCPClient()
    try:
        await client.connect_to_servers()
        await client.chat_loop()
    except KeyboardInterrupt:
        print("Goodbye!")
    finally:
        await client.cleanup()


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