feat:mcp
This commit is contained in:
269
plugins/MCPManager/main.py
Normal file
269
plugins/MCPManager/main.py
Normal file
@@ -0,0 +1,269 @@
|
||||
"""
|
||||
MCP 管理插件
|
||||
|
||||
管理 MCP (Model Context Protocol) 服务器，将 MCP 工具自动注册到 ToolRegistry，
|
||||
使 AI 可以调用各种 MCP 服务器提供的工具。
|
||||
|
||||
功能：
|
||||
- 自动连接配置的 MCP 服务器
|
||||
- 将 MCP 工具转换为 OpenAI 格式并注册
|
||||
- 支持热重载（禁用/启用插件时自动管理连接）
|
||||
- 提供管理命令（查看状态、重连等）
|
||||
"""
|
||||
|
||||
import tomllib
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
from loguru import logger
|
||||
|
||||
from utils.plugin_base import PluginBase
|
||||
from utils.decorators import on_text_message
|
||||
from utils.tool_registry import get_tool_registry
|
||||
|
||||
from .mcp_client import MCPManager, MCPServerConfig
|
||||
|
||||
|
||||
class MCPManagerPlugin(PluginBase):
|
||||
"""MCP 管理插件"""
|
||||
|
||||
description = "MCP 服务器管理，自动注册 MCP 工具到 AI"
|
||||
author = "ShiHao"
|
||||
version = "1.0.0"
|
||||
|
||||
# 高优先级加载，确保在其他插件之前注册工具
|
||||
load_priority = 90
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.config = None
|
||||
self.mcp_manager: MCPManager = None
|
||||
self._registered_tools: List[str] = [] # 已注册的工具名列表
|
||||
|
||||
async def async_init(self):
|
||||
"""插件异步初始化"""
|
||||
# 读取配置
|
||||
config_path = Path(__file__).parent / "config.toml"
|
||||
with open(config_path, "rb") as f:
|
||||
self.config = tomllib.load(f)
|
||||
|
||||
mcp_config = self.config.get("mcp", {})
|
||||
|
||||
if not mcp_config.get("enabled", True):
|
||||
logger.info("MCPManager: MCP 功能已禁用")
|
||||
return
|
||||
|
||||
# 初始化 MCP 管理器
|
||||
self.mcp_manager = MCPManager(
|
||||
tool_timeout=mcp_config.get("tool_timeout", 60),
|
||||
server_start_timeout=mcp_config.get("server_start_timeout", 30)
|
||||
)
|
||||
|
||||
# 连接所有配置的服务器
|
||||
servers = mcp_config.get("servers", [])
|
||||
if not servers:
|
||||
logger.info("MCPManager: 未配置任何 MCP 服务器")
|
||||
return
|
||||
|
||||
if mcp_config.get("auto_connect", True):
|
||||
await self._connect_all_servers(servers)
|
||||
|
||||
logger.success(f"MCPManager 插件已加载，已连接 {len(self.mcp_manager.clients)} 个 MCP 服务器")
|
||||
|
||||
async def _connect_all_servers(self, servers: List[Dict]):
|
||||
"""连接所有配置的服务器"""
|
||||
for server_data in servers:
|
||||
try:
|
||||
config = MCPServerConfig.from_dict(server_data)
|
||||
success = await self.mcp_manager.add_server(config)
|
||||
|
||||
if success:
|
||||
# 注册工具到 ToolRegistry
|
||||
await self._register_server_tools(config.name)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"MCPManager: 连接服务器 {server_data.get('name', '未知')} 失败: {e}")
|
||||
|
||||
async def _register_server_tools(self, server_name: str):
|
||||
"""将服务器的工具注册到 ToolRegistry"""
|
||||
client = self.mcp_manager.clients.get(server_name)
|
||||
if not client:
|
||||
return
|
||||
|
||||
registry = get_tool_registry()
|
||||
prefix = client.config.tool_prefix
|
||||
|
||||
for tool in client.tools.values():
|
||||
schema = tool.to_openai_schema(prefix)
|
||||
tool_name = schema["function"]["name"]
|
||||
|
||||
# 创建工具执行器
|
||||
async def executor(name: str, arguments: Dict, bot, from_wxid: str, _tn=tool_name) -> Dict:
|
||||
return await self.mcp_manager.call_tool(_tn, arguments)
|
||||
|
||||
# 注册到 ToolRegistry
|
||||
success = registry.register(
|
||||
name=tool_name,
|
||||
plugin_name=f"MCP:{server_name}",
|
||||
schema=schema,
|
||||
executor=executor,
|
||||
timeout=self.mcp_manager.tool_timeout,
|
||||
priority=40 # 比普通插件优先级稍低
|
||||
)
|
||||
|
||||
if success:
|
||||
self._registered_tools.append(tool_name)
|
||||
logger.debug(f"MCPManager: 注册工具 {tool_name}")
|
||||
|
||||
logger.info(f"MCPManager: 从 {server_name} 注册了 {len(client.tools)} 个工具")
|
||||
|
||||
async def on_disable(self):
|
||||
"""插件禁用时清理"""
|
||||
await super().on_disable()
|
||||
|
||||
# 注销所有已注册的工具
|
||||
registry = get_tool_registry()
|
||||
for tool_name in self._registered_tools:
|
||||
registry.unregister(tool_name)
|
||||
self._registered_tools.clear()
|
||||
|
||||
# 关闭所有 MCP 服务器连接
|
||||
if self.mcp_manager:
|
||||
await self.mcp_manager.shutdown()
|
||||
self.mcp_manager = None
|
||||
|
||||
logger.info("MCPManager: 已清理所有 MCP 连接和工具注册")
|
||||
|
||||
# ==================== 管理命令 ====================
|
||||
|
||||
@on_text_message(priority=80)
|
||||
async def handle_admin_commands(self, bot, message: dict):
|
||||
"""处理管理命令"""
|
||||
content = message.get("Content", "").strip()
|
||||
|
||||
# 只响应管理员的命令
|
||||
# TODO: 从主配置读取管理员列表
|
||||
if not content.startswith("/mcp"):
|
||||
return True
|
||||
|
||||
parts = content.split()
|
||||
if len(parts) < 2:
|
||||
return True
|
||||
|
||||
cmd = parts[1].lower()
|
||||
reply_to = message.get("FromWxid", "")
|
||||
|
||||
if cmd == "status":
|
||||
await self._cmd_status(bot, reply_to)
|
||||
return False
|
||||
|
||||
elif cmd == "list":
|
||||
await self._cmd_list_tools(bot, reply_to)
|
||||
return False
|
||||
|
||||
elif cmd == "reload":
|
||||
await self._cmd_reload(bot, reply_to)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def _cmd_status(self, bot, reply_to: str):
|
||||
"""查看 MCP 服务器状态"""
|
||||
if not self.mcp_manager:
|
||||
await bot.send_text(reply_to, "MCP 功能未启用")
|
||||
return
|
||||
|
||||
servers = self.mcp_manager.list_servers()
|
||||
if not servers:
|
||||
await bot.send_text(reply_to, "没有已连接的 MCP 服务器")
|
||||
return
|
||||
|
||||
lines = ["📡 MCP 服务器状态："]
|
||||
for s in servers:
|
||||
status = "✅" if s["connected"] else "❌"
|
||||
lines.append(f"{status} {s['name']}: {s['tools_count']} 个工具")
|
||||
|
||||
await bot.send_text(reply_to, "\n".join(lines))
|
||||
|
||||
async def _cmd_list_tools(self, bot, reply_to: str):
|
||||
"""列出所有 MCP 工具"""
|
||||
if not self.mcp_manager:
|
||||
await bot.send_text(reply_to, "MCP 功能未启用")
|
||||
return
|
||||
|
||||
servers = self.mcp_manager.list_servers()
|
||||
if not servers:
|
||||
await bot.send_text(reply_to, "没有已连接的 MCP 服务器")
|
||||
return
|
||||
|
||||
lines = ["🔧 MCP 工具列表："]
|
||||
for s in servers:
|
||||
if s["tools"]:
|
||||
lines.append(f"\n【{s['name']}】")
|
||||
for tool in s["tools"][:10]: # 限制显示数量
|
||||
lines.append(f" • {tool}")
|
||||
if len(s["tools"]) > 10:
|
||||
lines.append(f" ... 还有 {len(s['tools']) - 10} 个")
|
||||
|
||||
await bot.send_text(reply_to, "\n".join(lines))
|
||||
|
||||
async def _cmd_reload(self, bot, reply_to: str):
|
||||
"""重新加载 MCP 服务器"""
|
||||
await bot.send_text(reply_to, "正在重新加载 MCP 服务器...")
|
||||
|
||||
# 清理现有连接
|
||||
if self.mcp_manager:
|
||||
registry = get_tool_registry()
|
||||
for tool_name in self._registered_tools:
|
||||
registry.unregister(tool_name)
|
||||
self._registered_tools.clear()
|
||||
await self.mcp_manager.shutdown()
|
||||
|
||||
# 重新读取配置
|
||||
config_path = Path(__file__).parent / "config.toml"
|
||||
with open(config_path, "rb") as f:
|
||||
self.config = tomllib.load(f)
|
||||
|
||||
mcp_config = self.config.get("mcp", {})
|
||||
|
||||
# 重新初始化
|
||||
self.mcp_manager = MCPManager(
|
||||
tool_timeout=mcp_config.get("tool_timeout", 60),
|
||||
server_start_timeout=mcp_config.get("server_start_timeout", 30)
|
||||
)
|
||||
|
||||
servers = mcp_config.get("servers", [])
|
||||
await self._connect_all_servers(servers)
|
||||
|
||||
await bot.send_text(
|
||||
reply_to,
|
||||
f"MCP 重新加载完成，已连接 {len(self.mcp_manager.clients)} 个服务器"
|
||||
)
|
||||
|
||||
# ==================== LLM 工具接口（备用） ====================
|
||||
|
||||
def get_llm_tools(self) -> List[Dict]:
|
||||
"""
|
||||
返回 MCP 工具列表（备用接口）
|
||||
|
||||
注意：工具已通过 ToolRegistry 注册，此方法仅供参考
|
||||
"""
|
||||
if not self.mcp_manager:
|
||||
return []
|
||||
return self.mcp_manager.get_all_tools()
|
||||
|
||||
async def execute_llm_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
arguments: Dict[str, Any],
|
||||
bot,
|
||||
from_wxid: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
执行 MCP 工具（备用接口）
|
||||
|
||||
注意：工具已通过 ToolRegistry 注册，此方法仅供备用
|
||||
"""
|
||||
if not self.mcp_manager:
|
||||
return {"success": False, "error": "MCP 功能未启用"}
|
||||
|
||||
return await self.mcp_manager.call_tool(tool_name, arguments)
|
||||
Reference in New Issue
Block a user