pip install mcp openai

pip install uv

import asyncio
import os
import json
import sys
from typing import List, Dict, Any, Optional

# 導(dǎo)入 MCP 相關(guān)庫
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import mcp.types as types

# 導(dǎo)入 OpenAI SDK (DeepSeek 兼容 OpenAI 格式)
from openai import AsyncOpenAI

# ================= 配置部分 =================

# DeepSeek API 配置
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "<YOUR_DEEPSEEK_API_KEY>")
DEEPSEEK_BASE_URL = "https://api.deepseek.com"
MODEL_NAME = "deepseek-chat"  # 或者 "deepseek-reasoner" (R1)

# MCP ClickHouse Server 配置 (來自你的 JSON)
MCP_CLICKHOUSE_CONFIG = {
    "command": "uv",
    "args": [
        "run",
        "--with",
        "mcp-clickhouse",
        "--python",
        "3.10",
        "mcp-clickhouse"
    ],
    "env": {
        "CLICKHOUSE_HOST": "<clickhouse-host>",
        "CLICKHOUSE_PORT": "<clickhouse-port>",
        "CLICKHOUSE_USER": "<clickhouse-user>",
        "CLICKHOUSE_PASSWORD": "<clickhouse-password>",
        "CLICKHOUSE_ROLE": "<clickhouse-role>",
        "CLICKHOUSE_SECURE": "true",
        "CLICKHOUSE_VERIFY": "true",
        "CLICKHOUSE_CONNECT_TIMEOUT": "30",
        "CLICKHOUSE_SEND_RECEIVE_TIMEOUT": "30"
    }
}

# ================= 工具函數(shù) =================

def convert_mcp_tool_to_openai(mcp_tool: types.Tool) -> Dict[str, Any]:
    """
    將 MCP 工具定義轉(zhuǎn)換為 OpenAI/DeepSeek 支持的 Function Calling 格式。
    """
    return {
        "type": "function",
        "function": {
            "name": mcp_tool.name,
            "description": mcp_tool.description,
            "parameters": mcp_tool.inputSchema
        }
    }

# ================= 主邏輯 =================

async def run_chat_loop():
    # 1. 初始化 DeepSeek 客戶端
    client = AsyncOpenAI(api_key=DEEPSEEK_API_KEY, base_url=DEEPSEEK_BASE_URL)
    
    # 2. 準(zhǔn)備 MCP Server 參數(shù)
    server_params = StdioServerParameters(
        command=MCP_CLICKHOUSE_CONFIG["command"],
        args=MCP_CLICKHOUSE_CONFIG["args"],
        env={**os.environ, **MCP_CLICKHOUSE_CONFIG["env"]} # 合并當(dāng)前環(huán)境變量
    )

    print(f"正在啟動 MCP Server: {MCP_CLICKHOUSE_CONFIG['command']}...")
    
    # 3. 連接 MCP Server 并開始對話循環(huán)
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            # 初始化連接
            await session.initialize()
            
            # 獲取 MCP Server 提供的工具列表
            mcp_tools_list = await session.list_tools()
            openai_tools = [convert_mcp_tool_to_openai(tool) for tool in mcp_tools_list.tools]
            
            print(f"\n已連接 MCP Server，加載了 {len(openai_tools)} 個工具: {[t['function']['name'] for t in openai_tools]}")
            print("-" * 50)
            print("你可以開始詢問關(guān)于 ClickHouse 的問題了 (輸入 'quit' 退出)。")

            messages = [
                {"role": "system", "content": "你是一個智能助手，可以利用工具查詢 ClickHouse 數(shù)據(jù)庫來回答用戶問題。"}
            ]

            while True:
                user_input = input("\n用戶: ")
                if user_input.lower() in ["quit", "exit"]:
                    break

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

                # 第一輪調(diào)用：發(fā)送用戶問題 + 工具定義給 DeepSeek
                try:
                    response = await client.chat.completions.create(
                        model=MODEL_NAME,
                        messages=messages,
                        tools=openai_tools,
                    )
                except Exception as e:
                    print(f"API 調(diào)用錯誤: {e}")
                    continue

                assistant_msg = response.choices[0].message
                messages.append(assistant_msg)

                # 檢查 DeepSeek 是否想調(diào)用工具
                if assistant_msg.tool_calls:
                    print(f"\n[思考] 模型決定調(diào)用工具...")
                    
                    for tool_call in assistant_msg.tool_calls:
                        tool_name = tool_call.function.name
                        tool_args = json.loads(tool_call.function.arguments)
                        
                        print(f"  -> 調(diào)用工具: {tool_name}, 參數(shù): {tool_args}")

                        # 執(zhí)行 MCP 工具
                        try:
                            mcp_result = await session.call_tool(tool_name, arguments=tool_args)
                            
                            # 獲取文本結(jié)果 (ClickHouse MCP 通常返回文本或 JSON 文本)
                            tool_output = ""
                            if mcp_result.content:
                                for content in mcp_result.content:
                                    if content.type == "text":
                                        tool_output += content.text
                            
                            print(f"  <- 工具返回結(jié)果 (前100字符): {tool_output[:100]}...")

                            # 將工具結(jié)果添加回對話歷史
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": tool_output
                            })

                        except Exception as e:
                            error_msg = f"工具執(zhí)行出錯: {str(e)}"
                            print(f"  ! {error_msg}")
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": error_msg
                            })

                    # 第二輪調(diào)用：將工具結(jié)果發(fā)送回 DeepSeek 獲取最終回答
                    final_response = await client.chat.completions.create(
                        model=MODEL_NAME,
                        messages=messages,
                        # 這里依然傳入 tools，以防模型需要繼續(xù)多步調(diào)用
                        tools=openai_tools, 
                    )
                    
                    final_content = final_response.choices[0].message.content
                    print(f"\nDeepSeek: {final_content}")
                    messages.append(final_response.choices[0].message)
                
                else:
                    # 如果不需要調(diào)用工具，直接輸出結(jié)果
                    print(f"\nDeepSeek: {assistant_msg.content}")

if __name__ == "__main__":
    # Windows 下通常需要這個策略
    if sys.platform.startswith('win'):
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
        
    asyncio.run(run_chat_loop())