mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
261 lines
10 KiB
Python
261 lines
10 KiB
Python
import os
|
||
import json
|
||
import traceback
|
||
from .route import Route, Response, RouteContext
|
||
from quart import request
|
||
from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
|
||
from astrbot.core import logger
|
||
|
||
DEFAULT_MCP_CONFIG = {"mcpServers": {}}
|
||
|
||
|
||
class ToolsRoute(Route):
|
||
def __init__(
|
||
self, context: RouteContext, core_lifecycle: AstrBotCoreLifecycle
|
||
) -> None:
|
||
super().__init__(context)
|
||
self.core_lifecycle = core_lifecycle
|
||
self.routes = {
|
||
"/tools/mcp/servers": ("GET", self.get_mcp_servers),
|
||
"/tools/mcp/add": ("POST", self.add_mcp_server),
|
||
"/tools/mcp/update": ("POST", self.update_mcp_server),
|
||
"/tools/mcp/delete": ("POST", self.delete_mcp_server),
|
||
}
|
||
self.register_routes()
|
||
self.tool_mgr = self.core_lifecycle.provider_manager.llm_tools
|
||
|
||
@property
|
||
def mcp_config_path(self):
|
||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||
data_dir = os.path.abspath(os.path.join(current_dir, "../../../data"))
|
||
return os.path.join(data_dir, "mcp_server.json")
|
||
|
||
def load_mcp_config(self):
|
||
if not os.path.exists(self.mcp_config_path):
|
||
# 配置文件不存在，创建默认配置
|
||
os.makedirs(os.path.dirname(self.mcp_config_path), exist_ok=True)
|
||
with open(self.mcp_config_path, "w", encoding="utf-8") as f:
|
||
json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4)
|
||
return DEFAULT_MCP_CONFIG
|
||
|
||
try:
|
||
with open(self.mcp_config_path, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
except Exception as e:
|
||
logger.error(f"加载 MCP 配置失败: {e}")
|
||
return DEFAULT_MCP_CONFIG
|
||
|
||
def save_mcp_config(self, config):
|
||
try:
|
||
with open(self.mcp_config_path, "w", encoding="utf-8") as f:
|
||
json.dump(config, f, ensure_ascii=False, indent=4)
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"保存 MCP 配置失败: {e}")
|
||
return False
|
||
|
||
async def get_mcp_servers(self):
|
||
try:
|
||
config = self.load_mcp_config()
|
||
servers = []
|
||
|
||
# 获取所有服务器并添加它们的工具列表
|
||
for name, server_config in config["mcpServers"].items():
|
||
server_info = {
|
||
"name": name,
|
||
"active": server_config.get("active", True),
|
||
}
|
||
|
||
# 复制所有配置字段
|
||
for key, value in server_config.items():
|
||
if key != "active": # active 已经处理
|
||
server_info[key] = value
|
||
|
||
# 如果MCP客户端已初始化，从客户端获取工具名称
|
||
for (
|
||
name_key,
|
||
mcp_client,
|
||
) in self.tool_mgr.mcp_client_dict.items():
|
||
if name_key == name:
|
||
server_info["tools"] = [tool.name for tool in mcp_client.tools]
|
||
break
|
||
else:
|
||
server_info["tools"] = []
|
||
|
||
servers.append(server_info)
|
||
|
||
return Response().ok(servers).__dict__
|
||
except Exception as e:
|
||
logger.error(traceback.format_exc())
|
||
return Response().error(f"获取 MCP 服务器列表失败: {str(e)}").__dict__
|
||
|
||
async def add_mcp_server(self):
|
||
try:
|
||
server_data = await request.json
|
||
|
||
name = server_data.get("name", "")
|
||
|
||
# 检查必填字段
|
||
if not name:
|
||
return Response().error("服务器名称不能为空").__dict__
|
||
|
||
# 移除特殊字段并检查配置是否有效
|
||
has_valid_config = False
|
||
server_config = {"active": server_data.get("active", True)}
|
||
|
||
# 复制所有配置字段
|
||
for key, value in server_data.items():
|
||
if key not in ["name", "active", "tools"]: # 排除特殊字段
|
||
if key == "mcpServers":
|
||
key_0 = list(server_data["mcpServers"].keys())[0] # 不考虑为空的情况
|
||
server_config = server_data["mcpServers"][key_0]
|
||
else:
|
||
server_config[key] = value
|
||
has_valid_config = True
|
||
|
||
if not has_valid_config:
|
||
return Response().error("必须提供有效的服务器配置").__dict__
|
||
|
||
config = self.load_mcp_config()
|
||
|
||
if name in config["mcpServers"]:
|
||
return Response().error(f"服务器 {name} 已存在").__dict__
|
||
|
||
config["mcpServers"][name] = server_config
|
||
|
||
if self.save_mcp_config(config):
|
||
# 动态初始化新MCP客户端
|
||
self.tool_mgr.mcp_service_queue.put_nowait(
|
||
{
|
||
"type": "init",
|
||
"name": name,
|
||
"cfg": config["mcpServers"][name],
|
||
}
|
||
)
|
||
return Response().ok(None, f"成功添加 MCP 服务器 {name}").__dict__
|
||
else:
|
||
return Response().error("保存配置失败").__dict__
|
||
except Exception as e:
|
||
logger.error(traceback.format_exc())
|
||
return Response().error(f"添加 MCP 服务器失败: {str(e)}").__dict__
|
||
|
||
async def update_mcp_server(self):
|
||
try:
|
||
server_data = await request.json
|
||
|
||
name = server_data.get("name", "")
|
||
|
||
if not name:
|
||
return Response().error("服务器名称不能为空").__dict__
|
||
|
||
config = self.load_mcp_config()
|
||
|
||
if name not in config["mcpServers"]:
|
||
return Response().error(f"服务器 {name} 不存在").__dict__
|
||
|
||
# 获取活动状态
|
||
active = server_data.get(
|
||
"active", config["mcpServers"][name].get("active", True)
|
||
)
|
||
|
||
# 创建新的配置对象
|
||
server_config = {"active": active}
|
||
|
||
# 仅更新活动状态的特殊处理
|
||
only_update_active = True
|
||
|
||
# 复制所有配置字段
|
||
for key, value in server_data.items():
|
||
if key not in ["name", "active", "tools"]: # 排除特殊字段
|
||
if key == "mcpServers":
|
||
key_0 = list(server_data["mcpServers"].keys())[0] # 不考虑为空的情况
|
||
server_config = server_data["mcpServers"][key_0]
|
||
else:
|
||
server_config[key] = value
|
||
only_update_active = False
|
||
|
||
# 如果只更新活动状态，保留原始配置
|
||
if only_update_active:
|
||
for key, value in config["mcpServers"][name].items():
|
||
if key != "active": # 除了active之外的所有字段都保留
|
||
server_config[key] = value
|
||
|
||
config["mcpServers"][name] = server_config
|
||
|
||
if self.save_mcp_config(config):
|
||
# 处理MCP客户端状态变化
|
||
if active:
|
||
# 如果要激活服务器或者配置已更改
|
||
if name in self.tool_mgr.mcp_client_dict or not only_update_active:
|
||
await self.tool_mgr.mcp_service_queue.put(
|
||
{
|
||
"type": "terminate",
|
||
"name": name,
|
||
}
|
||
)
|
||
await self.tool_mgr.mcp_service_queue.put(
|
||
{
|
||
"type": "init",
|
||
"name": name,
|
||
"cfg": config["mcpServers"][name],
|
||
}
|
||
)
|
||
else:
|
||
# 客户端不存在，初始化
|
||
self.tool_mgr.mcp_service_queue.put_nowait(
|
||
{
|
||
"type": "init",
|
||
"name": name,
|
||
"cfg": config["mcpServers"][name],
|
||
}
|
||
)
|
||
else:
|
||
# 如果要停用服务器
|
||
if name in self.tool_mgr.mcp_client_dict:
|
||
self.tool_mgr.mcp_service_queue.put_nowait(
|
||
{
|
||
"type": "terminate",
|
||
"name": name,
|
||
}
|
||
)
|
||
|
||
return Response().ok(None, f"成功更新 MCP 服务器 {name}").__dict__
|
||
else:
|
||
return Response().error("保存配置失败").__dict__
|
||
except Exception as e:
|
||
logger.error(traceback.format_exc())
|
||
return Response().error(f"更新 MCP 服务器失败: {str(e)}").__dict__
|
||
|
||
async def delete_mcp_server(self):
|
||
try:
|
||
server_data = await request.json
|
||
name = server_data.get("name", "")
|
||
|
||
if not name:
|
||
return Response().error("服务器名称不能为空").__dict__
|
||
|
||
config = self.load_mcp_config()
|
||
|
||
if name not in config["mcpServers"]:
|
||
return Response().error(f"服务器 {name} 不存在").__dict__
|
||
|
||
# 删除服务器配置
|
||
del config["mcpServers"][name]
|
||
|
||
if self.save_mcp_config(config):
|
||
# 关闭并删除MCP客户端
|
||
if name in self.tool_mgr.mcp_client_dict:
|
||
self.tool_mgr.mcp_service_queue.put_nowait(
|
||
{
|
||
"type": "terminate",
|
||
"name": name,
|
||
}
|
||
)
|
||
|
||
return Response().ok(None, f"成功删除 MCP 服务器 {name}").__dict__
|
||
else:
|
||
return Response().error("保存配置失败").__dict__
|
||
except Exception as e:
|
||
logger.error(traceback.format_exc())
|
||
return Response().error(f"删除 MCP 服务器失败: {str(e)}").__dict__
|