mcp = FastMCP(
    name="myMCP",
    host="0.0.0.0",
    port=8888,
    description="雜七雜八mcp",
    sse_path="/sse",
)

if __name__ == "__main__":
    # 初始化并運(yùn)行服務(wù)器
    try:
        print("Starting server...")
        mcp.run(transport="sse")
    except Exception as e:
        print(f"Error: {e}")

import json
import asyncio
import sys
from typing import Optional
from contextlib import AsyncExitStack

from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client

from dotenv import load_dotenv
from openai import AsyncOpenAI, OpenAI
import os

load_dotenv()
api_key = os.getenv("api_key")
base_url = os.getenv("base_url")


class Client:
    def __init__(self):
        self._exit_stack: Optional[AsyncExitStack] = None
        self.session: Optional[ClientSession] = None
        self._lock = asyncio.Lock()
        self.is_connected = False
        self.client = AsyncOpenAI(
            base_url=base_url,
            api_key=api_key,
        )
        self.model = "qwen-plus-2025-04-28"
        self.messages = []

    async def connect_server(self, server_config):
        async with self._lock:
            url = server_config["mcpServers"]["amap-amap-sse"]["url"]
            print(f"嘗試連接到: {url}")
            self._exit_stack = AsyncExitStack()
            sse_cm = sse_client(url)
            streams = await self._exit_stack.enter_async_context(sse_cm)

            session_cm = ClientSession(streams[0], streams[1])
            self.session = await self._exit_stack.enter_async_context(session_cm)

            await self.session.initialize()

            response = await self.session.list_tools()
            self.tools = {tool.name: tool for tool in response.tools}
            print(f"成功獲取 {len(self.tools)} 個(gè)工具:")

            # 將工具轉(zhuǎn)換為 OpenAI 格式
            self.openai_tools = [
                self.convert_mcp_tool_to_openai_tool(tool) for tool in response.tools
            ]

            print("連接成功并準(zhǔn)備就緒。")

    def convert_mcp_tool_to_openai_tool(self, mcp_tool):
        """將 MCP 工具轉(zhuǎn)換為 OpenAI 的 function call 格式"""
        return {
            "type": "function",
            "function": {
                "name": mcp_tool.name,
                "description": mcp_tool.description,
                "parameters": mcp_tool.inputSchema,
            },
        }

    async def chat(self, prompt, role="user"):
        self.messages.append({"role": role, "content": prompt})

        # 初始化 LLM API 調(diào)用，傳入 tools 參數(shù)
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=self.messages,
            tools=self.openai_tools,
            tool_choice="auto",  # 或者指定具體工具名
        )

        if response.choices[0].finish_reason == "tool_calls":
            # 獲取第一個(gè) tool_call
            tool_call = response.choices[0].message.tool_calls[0]

            print(f"[提示]：正在調(diào)用工具 {tool_call.function.name}")
            result = await self.session.call_tool(
                tool_call.function.name, json.loads(tool_call.function.arguments)
            )

            # 添加 tool_call 和 tool_response 到 messages
            self.messages.append(
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [tool_call.model_dump()],
                }
            )

            # 添加 tool_response 到 messages
            self.messages.append(
                {
                    "role": "tool",
                    "name": tool_call.function.name,
                    "content": str(result),
                }
            )

            # 繼續(xù)生成最終回答
            followup_response = await self.client.chat.completions.create(
                model=self.model,
                messages=self.messages,
            )
            content = followup_response.choices[0].message.content
            self.messages.append({"role": "assistant", "content": content})
            return content

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

    async def chat_loop(self):
        print("MCP 客戶端啟動(dòng)")
        print("輸入 /bye 退出")

        while True:
            prompt = input(">>> ").strip()
            if "/bye" in prompt.lower():
                break

            response = await self.chat(prompt)
            print(response)

    async def disconnect(self):
        """關(guān)閉 Session 和連接"""
        if self._exit_stack is not None:
            await self._exit_stack.aclose()
        self.is_connected = False
        print("客戶端已關(guān)閉")


def load_server_config(config_file):
    with open(config_file) as f:
        return json.load(f)


async def main():
    try:
        server_config = load_server_config("servers_config.json")
        client = Client()
        await client.connect_server(server_config)
        await client.chat_loop()
    except Exception as e:
        print(f"主程序發(fā)生錯(cuò)誤: {type(e).__name__}: {e}")
    finally:
        print("\n正在關(guān)閉客戶端...")
        await client.disconnect()
        print("客戶端已關(guān)閉。")


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

import json
import asyncio
import re
import sys
import traceback
from typing import Optional
from contextlib import AsyncExitStack

from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client

from dotenv import load_dotenv
from openai import AsyncOpenAI, OpenAI

import os

load_dotenv()
api_key = os.getenv("api_key")
base_url = os.getenv("base_url")


def format_tools_for_llm(tool) -> str:
    """對(duì)tool進(jìn)行格式化
    Returns:
        格式化之后的tool描述
    """
    args_desc = []
    if "properties" in tool.inputSchema:
        for param_name, param_info in tool.inputSchema["properties"].items():
            arg_desc = (
                f"- {param_name}: {param_info.get('description', 'No description')}"
            )
            if param_name in tool.inputSchema.get("required", []):
                arg_desc += " (required)"
            args_desc.append(arg_desc)

    return f"Tool: {tool.name}\nDescription: {tool.description}\nArguments:\n{chr(10).join(args_desc)}"


class Client:
    def __init__(self, url: str):
        self._exit_stack: Optional[AsyncExitStack] = None
        self.session: Optional[ClientSession] = None
        self._lock = asyncio.Lock()  # 防止并發(fā)連接/斷開(kāi)問(wèn)題
        self.is_connected = False
        self.server_url = url

    async def connect_server(self):
        async with self._lock:  # 防止并發(fā)調(diào)用 connect
            url = self.server_url
            print(f"嘗試連接到: {url}")
            self._exit_stack = AsyncExitStack()
            # 1. 進(jìn)入 SSE 上下文，但不退出
            sse_cm = sse_client(url)
            # 手動(dòng)調(diào)用 __aenter__ 獲取流，并存儲(chǔ)上下文管理器以便后續(xù)退出
            streams = await self._exit_stack.enter_async_context(sse_cm)
            print("SSE 流已獲取。")

            # 2. 進(jìn)入 Session 上下文，但不退出
            session_cm = ClientSession(streams[0], streams[1])
            # 手動(dòng)調(diào)用 __aenter__ 獲取 session
            self.session = await self._exit_stack.enter_async_context(session_cm)
            print("ClientSession 已創(chuàng)建。")

            # 3. 初始化 Session
            await self.session.initialize()
            print("Session 已初始化。")

            # 4. 獲取并存儲(chǔ)工具列表
            response = await self.session.list_tools()
            self.tools = {tool.name: tool for tool in response.tools}
            print(f"成功獲取 {len(self.tools)} 個(gè)工具:")
            for name, tool in self.tools.items():
                print(
                    f"  - {name}: {tool.description[:50]}...:{tool.annotations}"
                )  # 打印部分描述

            print("連接成功并準(zhǔn)備就緒。")

        # # 列出可用工具
        # response = await self.session.list_tools()
        # tools = response.tools

    async def disconnect(self):
        """關(guān)閉 Session 和連接。"""
        async with self._lock:
            await self._exit_stack.aclose()


class Host:
    def __init__(self):

        self.client = AsyncOpenAI(
            base_url=base_url,
            api_key=api_key,
        )
        self.model = "qwen-plus-2025-04-28"
        self.messages = []
        self.all_tools = []
        self.mcp_clients = []

    async def connect_mcp_servers(self, server_config):
        servers = server_config["mcpServers"]
        for name, server_info in servers.items():
            try:

                server_url = server_info["url"]
                client = Client(server_url)
                print(f"正在連接MCP服務(wù)器: {name} {server_url}")
                await client.connect_server()

                self.mcp_clients.append(client)

                response = await client.session.list_tools()
                tools = response.tools
                self.all_tools.extend(tools)

            except Exception as e:
                print(f"連接MCP服務(wù)器: {name} 失敗: {type(e).__name__}: {e}")

        tools_description = "\n".join(
            [format_tools_for_llm(tool) for tool in self.all_tools]
        )
        # 修改系統(tǒng)提示
        system_prompt = (
            "You are a helpful assistant with access to these tools:\n\n"
            f"{tools_description}\n"
            "Choose the appropriate tool based on the user's question. "
            "If no tool is needed, reply directly.\n\n"
            "IMPORTANT: When you need to use a tool, you must ONLY respond with "
            "the exact JSON object format below, nothing else:\n"
            "{\n"
            '    "tool": "tool-name",\n'
            '    "arguments": {\n'
            '        "argument-name": "value"\n'
            "    }\n"
            "}\n\n"
            '"```json" is not allowed'
            "After receiving a tool's response:\n"
            "1. Transform the raw data into a natural, conversational response\n"
            "2. Keep responses concise but informative\n"
            "3. Focus on the most relevant information\n"
            "4. Use appropriate context from the user's question\n"
            "5. Avoid simply repeating the raw data\n\n"
            "Please use only the tools that are explicitly defined above."
        )
        self.messages.append({"role": "system", "content": system_prompt})

    async def disconnect_mcp_servers(self):
        for mcp_server in self.mcp_clients:
            await mcp_server.disconnect()

    async def chat(self, prompt, role="user"):
        """與LLM進(jìn)行交互"""
        self.messages.append({"role": role, "content": prompt})

        # 初始化 LLM API 調(diào)用
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=self.messages,
        )
        llm_response = response.choices[0].message.content
        return llm_response

    async def find_client_tool(self, tool_name: str):
        """查找客戶端工具"""
        target_client = None
        target_tool = None
        for client in self.mcp_clients:
            response = await client.session.list_tools()
            tools = response.tools
            for tool in tools:
                if tool.name == tool_name:
                    target_client = client
                    target_tool = tool
                    break

            if target_tool:
                break

        return target_client, target_tool

    async def execute_tool(self, llm_response: str):
        """Process the LLM response and execute tools if needed.
        Args:
            llm_response: The response from the LLM.
        Returns:
            The result of tool execution or the original response.
        """
        import json

        print(f"LLM Response: {llm_response}")
        try:
            pattern = r"```json\n(.*?)\n?```"
            match = re.search(pattern, llm_response, re.DOTALL)
            if match:
                llm_response = match.group(1)
            tool_call = json.loads(llm_response)

            if "tool" in tool_call and "arguments" in tool_call:
                # result = await self.session.call_tool(tool_name, tool_args)
                # 如果有多個(gè)client，需要先找到clinet，再調(diào)用tool
                # response = await self.session.list_tools()
                # tools = response.tools

                target_client, target_tool = await self.find_client_tool(
                    tool_call["tool"]
                )

                if target_client != None and target_tool != None:
                    try:
                        print(f"[提示]：正在調(diào)用工具 {tool_call['tool']}")
                        result = await target_client.session.call_tool(
                            tool_call["tool"], tool_call["arguments"]
                        )

                        if isinstance(result, dict) and "progress" in result:
                            progress = result["progress"]
                            total = result["total"]
                            percentage = (progress / total) * 100
                            print(f"Progress: {progress}/{total} ({percentage:.1f}%)")
                        print(f"[執(zhí)行結(jié)果]: {result}")
                        return f"Tool execution result: {result}"
                    except Exception as e:
                        error_msg = f"Error executing tool: {str(e)}"
                        print(error_msg)
                        return error_msg

                return f"No server found with tool: {tool_call['tool']}"
            return llm_response
        except json.JSONDecodeError:
            return llm_response

    async def chat_loop(self):
        """運(yùn)行交互式聊天循環(huán)"""
        print("MCP 客戶端啟動(dòng)")
        print("輸入 /bye 退出")

        while True:
            prompt = input(">>> ").strip()
            if "/bye" in prompt.lower():
                break

            response = await self.chat(prompt)
            self.messages.append({"role": "assistant", "content": response})

            result = await self.execute_tool(response)
            while result != response:
                response = await self.chat(result, "system")
                self.messages.append({"role": "assistant", "content": response})
                result = await self.execute_tool(response)
            print(response)


def load_server_config(config_file):
    with open(config_file) as f:
        return json.load(f)


async def main():
    host = None
    try:
        server_config = load_server_config("servers_config.json")
        host = Host()
        await host.connect_mcp_servers(server_config)
        await host.chat_loop()
    except Exception as e:
        print(f"主程序發(fā)生錯(cuò)誤: {type(e).__name__}: {e}")
        # 打印完整的調(diào)用堆棧
        traceback.print_exc()
    finally:
        # 無(wú)論如何，最后都要嘗試斷開(kāi)連接并清理資源
        print("\n正在關(guān)閉客戶端...")
        await host.disconnect_mcp_servers()
        print("客戶端已關(guān)閉。")


if __name__ == "__main__":
    # 我要去濟(jì)南奧體中心出差，請(qǐng)你查詢附近5km的酒店，為我安排行程
    asyncio.run(main())