update:合并最新代码

This commit is contained in:
hrz
2025-04-04 00:27:04 +08:00
parent c8a3d378b7
commit 0b4a4df1af
42 changed files with 4559 additions and 532 deletions
+86
View File
@@ -0,0 +1,86 @@
from datetime import timedelta
from typing import Optional
from contextlib import AsyncExitStack
import os, shutil
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from config.logger import setup_logging
TAG = __name__
class MCPClient:
def __init__(self, config):
# Initialize session and client objects
self.session: Optional[ClientSession] = None
self.exit_stack = AsyncExitStack()
self.logger = setup_logging()
self.config = config
self.tolls = []
async def initialize(self):
args = self.config.get("args", [])
command = (
shutil.which("npx")
if self.config["command"] == "npx"
else self.config["command"]
)
env={**os.environ}
if self.config.get("env"):
env.update(self.config["env"])
server_params = StdioServerParameters(
command=command,
args=args,
env=env
)
stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
self.stdio, self.write = stdio_transport
time_out_delta = timedelta(seconds=15)
self.session = await self.exit_stack.enter_async_context(ClientSession(read_stream=self.stdio, write_stream=self.write, read_timeout_seconds=time_out_delta))
await self.session.initialize()
# List available tools
response = await self.session.list_tools()
tools = response.tools
self.tools = tools
self.logger.bind(tag=TAG).info(f"Connected to server with tools:{[tool.name for tool in tools]}")
def has_tool(self, tool_name):
return any(tool.name == tool_name for tool in self.tools)
def get_available_tools(self):
available_tools = [{"type": "function", "function":{
"name": tool.name,
"description": tool.description,
"parameters": tool.inputSchema
} } for tool in self.tools]
return available_tools
async def call_tool(self, tool_name: str, tool_args: dict):
self.logger.bind(tag=TAG).info(f"MCPClient Calling tool {tool_name} with args: {tool_args}")
try:
response = await self.session.call_tool(tool_name, tool_args)
except Exception as e:
self.logger.bind(tag=TAG).error(f"Error calling tool {tool_name}: {e}")
from types import SimpleNamespace
error_content = SimpleNamespace(
type='text',
text=f"Error calling tool {tool_name}: {e}"
)
error_response = SimpleNamespace(
content=[error_content],
isError=True
)
return error_response
self.logger.bind(tag=TAG).info(f"MCPClient Response from tool {tool_name}: {response}")
return response
async def cleanup(self):
"""Clean up resources"""
await self.exit_stack.aclose()
+110
View File
@@ -0,0 +1,110 @@
"""MCP服务管理器"""
import os, json
from typing import Dict, Any, List
from .MCPClient import MCPClient
from config.logger import setup_logging
from core.utils.util import get_project_dir
from plugins_func.register import register_function, ActionResponse, Action, ToolType
TAG = __name__
class MCPManager:
"""管理多个MCP服务的集中管理器"""
def __init__(self,conn) -> None:
"""
初始化MCP管理器
"""
self.conn = conn
self.logger = setup_logging()
self.config_path = get_project_dir() + 'data/.mcp_server_settings.json'
if os.path.exists(self.config_path) == False:
self.config_path = ""
self.logger.bind(tag=TAG).warning(f"请检查mcp服务配置文件:data/.mcp_server_settings.json")
self.client: Dict[str, MCPClient] = {}
self.tools = []
def load_config(self) -> Dict[str, Any]:
"""加载MCP服务配置
Returns:
Dict[str, Any]: 服务配置字典
"""
if len(self.config_path) == 0:
return {}
try:
with open(self.config_path, 'r', encoding='utf-8') as f:
config = json.load(f)
return config.get('mcpServers', {})
except Exception as e:
self.logger.bind(tag=TAG).error(f"Error loading MCP config from {self.config_path}: {e}")
return {}
async def initialize_servers(self) -> None:
"""初始化所有MCP服务"""
config = self.load_config()
for name, srv_config in config.items():
if not srv_config.get("command"):
self.logger.bind(tag=TAG).warning(f"Skipping server {name}: command not specified")
continue
try:
client = MCPClient(srv_config)
await client.initialize()
self.client[name] = client
self.logger.bind(tag=TAG).info(f"Initialized MCP client: {name}")
client_tools = client.get_available_tools()
self.tools.extend(client_tools)
for tool in client_tools:
func_name = "mcp_"+tool["function"]["name"]
register_function(func_name, tool, ToolType.MCP_CLIENT)(self.execute_tool)
self.conn.func_handler.function_registry.register_function(func_name)
except Exception as e:
self.logger.bind(tag=TAG).error(f"Failed to initialize MCP server {name}: {e}")
self.conn.func_handler.upload_functions_desc()
def get_all_tools(self) -> List[Dict[str, Any]]:
"""获取所有服务的工具function定义
Returns:
List[Dict[str, Any]]: 所有工具的function定义列表
"""
return self.tools
def is_mcp_tool(self, tool_name: str) -> bool:
"""检查是否是MCP工具
Args:
tool_name: 工具名称
Returns:
bool: 是否是MCP工具
"""
for tool in self.tools:
if tool.get("function") != None and tool["function"].get("name") == tool_name:
return True
return False
async def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""执行工具调用
Args:
tool_name: 工具名称
arguments: 工具参数
Returns:
Any: 工具执行结果
Raises:
ValueError: 工具未找到时抛出
"""
self.logger.bind(tag=TAG).info(f"Executing tool {tool_name} with arguments: {arguments}")
for client in self.client.values():
if client.has_tool(tool_name):
return await client.call_tool(tool_name, arguments)
raise ValueError(f"Tool {tool_name} not found in any MCP server")
async def cleanup_all(self) -> None:
for name, client in self.client.items():
try:
await client.cleanup()
self.logger.bind(tag=TAG).info(f"Cleaned up MCP client: {name}")
except Exception as e:
self.logger.bind(tag=TAG).error(f"Error cleaning up MCP client {name}: {e}")
self.client.clear()