1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
| import asyncio
import json
import os
import shutil
from contextlib import AsyncExitStack
from typing import Dict, Any, List, Optional
import httpx
from loguru import logger
from openai import AsyncOpenAI, OpenAI
from agents import (
Agent,
Runner,
gen_trace_id,
trace,
set_default_openai_api,
set_default_openai_client,
set_tracing_disabled,
)
from agents.mcp import MCPServer, MCPServerStdio
class OpenAIMCPClient:
def __init__(self, servers_config_path) :
self.exit_stack = AsyncExitStack()
self.mcp_servers: List[MCPServerStdio] = []
self.system_prompt = """
你是一个有用的助手，可以帮助用户解决问题。
你能够通过工具获取当前准确时间。
"""
self.servers_config = self.load_server_config(servers_config_path)
self._initialized = False
self.model = "gpt-4o"
async def initialize(self) -> None:
if self._initialized:
return
self.mcp_servers = await self._init_mcp_servers(self.servers_config)
self.agent = self._init_openai_agent(self.system_prompt, self.mcp_servers, self.model)
self._initialized = True
logger.info("所有MCP服务器初始化完成")
def load_server_config(self, file_path: str) -> Dict[str, Any]:
"""加载服务器配置"""
with open(file_path, "r") as f:
return json.load(f)
def _init_openai_agent(self, system_prompt, mcp_servers, model: str = "gpt-4o"):
"""初始化OpenAI Agent"""
client = AsyncOpenAI(
api_key=os.getenv("OPENAI_API_KEY"),
base_url=os.getenv("OPENAI_BASE_URL"),
http_client=httpx.AsyncClient(verify=False),
)
set_default_openai_client(client=client, use_for_tracing=False)
set_default_openai_api("chat_completions")
set_tracing_disabled(disabled=True)
agent = Agent(
name="Assistant",
instructions=system_prompt,
mcp_servers=mcp_servers,
model=model,
)
return agent
async def _init_mcp_servers(self, servers_config: Dict[str, Any]) -> List[MCPServerStdio]:
"""异步初始化MCP服务器并返回服务器列表"""
mcp_servers = []
for name, config in servers_config["mcpServers"].items():
command = (
shutil.which("npx")
if config["command"] == "npx"
else config["command"]
)
if command is None:
raise ValueError("命令必须是有效的字符串，不能为None。")
server = MCPServerStdio(
name=name,
params={
"command": command,
"args": config["args"],
"env": {**os.environ, **config["env"]} if config.get("env") else None,
}
)
await self.exit_stack.enter_async_context(server)
mcp_servers.append(server)
return mcp_servers
async def process_message(self, user_message: str) -> str:
"""处理用户消息并返回响应"""
if not self._initialized:
raise RuntimeError("MCP服务器未初始化，请先调用initialize方法")
try:
trace_id = gen_trace_id()
with trace(workflow_name="User Query", trace_id=trace_id):
logger.info(f"处理用户消息: {user_message}")
result = await Runner.run(starting_agent=self.agent, input=user_message)
logger.info(f"处理完成，获得响应")
return result.final_output
except Exception as e:
error_msg = f"处理消息时出错: {str(e)}"
logger.error(error_msg)
return error_msg
async def chat_loop(self) -> None:
if not self._initialized:
await self.initialize()
try:
print("欢迎使用MCP聊天助手！输入 'quit' 或 'exit' 退出。")
while True:
try:
user_input = input("\n您: ").strip()
if user_input.lower() in ["quit", "exit"]:
print("\n退出...")
break
print("\n处理中...")
response = await self.process_message(user_input)
print(f"\n助手: {response}")
except KeyboardInterrupt:
print("\n退出...")
break
except Exception as e:
print(f"\n出错: {str(e)}")
except Exception as e:
logger.error(f"聊天循环出错: {str(e)}")
finally:
await self.cleanup()
async def cleanup(self) -> None:
"""清理服务器资源"""
if not self._initialized:
return
try:
print("正在清理资源...")
logger.info("开始清理所有服务器...")
await self.exit_stack.aclose()
logger.info("所有服务器清理完成")
except Exception as e:
logger.error(f"清理服务器资源时出错: {e}")
finally:
self._initialized = False
self.mcp_servers = []
self.agent = None
async def main() -> None:
"""初始化并运行聊天会话"""
try:
mcp_client = OpenAIMCPClient("servers_config.json")
await mcp_client.initialize()
await mcp_client.chat_loop()
except asyncio.CancelledError:
logger.info("程序被取消")
except Exception as e:
logger.error(f"运行时出错: {str(e)}")
if __name__ == "__main__":
asyncio.run(main())
|
