""" 工具注册中心 集中管理所有 LLM 工具的注册、查找和执行 - O(1) 工具查找（替代 O(n) 插件遍历） - 统一的超时保护 - 工具元信息管理 使用示例: from utils.tool_registry import get_tool_registry registry = get_tool_registry() # 注册工具 registry.register( name="generate_image", plugin_name="AIChat", schema={...}, executor=some_async_func, timeout=120 ) # 执行工具 result = await registry.execute("generate_image", arguments, bot, from_wxid) """ import asyncio from dataclasses import dataclass, field from threading import Lock from typing import Any, Callable, Dict, List, Optional, Awaitable from loguru import logger @dataclass class ToolDefinition: """工具定义""" name: str plugin_name: str schema: Dict[str, Any] # OpenAI-compatible tool schema executor: Callable[..., Awaitable[Dict[str, Any]]] timeout: float = 60.0 priority: int = 50 # 同名工具时优先级高的生效 description: str = "" def __post_init__(self): # 从 schema 提取描述 if not self.description and self.schema: func_def = self.schema.get("function", {}) self.description = func_def.get("description", "") class ToolRegistry: """ 工具注册中心（线程安全单例） 功能: - 工具注册与注销 - O(1) 工具查找 - 统一超时保护执行 - 工具列表导出（供 LLM 使用） """ _instance: Optional["ToolRegistry"] = None _lock = Lock() def __new__(cls): if cls._instance is None: with cls._lock: if cls._instance is None: instance = super().__new__(cls) instance._initialized = False cls._instance = instance return cls._instance def __init__(self): if self._initialized: return self._tools: Dict[str, ToolDefinition] = {} self._tools_by_plugin: Dict[str, List[str]] = {} # plugin_name -> [tool_names] self._registry_lock = Lock() self._initialized = True logger.debug("ToolRegistry 初始化完成") def register( self, name: str, plugin_name: str, schema: Dict[str, Any], executor: Callable[..., Awaitable[Dict[str, Any]]], timeout: float = 60.0, priority: int = 50, ) -> bool: """ 注册工具 Args: name: 工具名称（唯一标识） plugin_name: 所属插件名 schema: OpenAI-compatible tool schema executor: 异步执行函数，签名: async (tool_name, arguments, bot, from_wxid) -> dict timeout: 执行超时（秒） priority: 优先级（同名工具时高优先级覆盖低优先级） Returns: 是否注册成功 """ with self._registry_lock: # 检查是否已存在同名工具 existing = self._tools.get(name) if existing: if existing.priority >= priority: logger.debug(f"工具 {name} 已存在且优先级更高，跳过注册") return False logger.info(f"工具 {name} 被 {plugin_name} 覆盖（优先级 {priority} > {existing.priority}）") # 从旧插件的工具列表中移除 old_plugin_tools = self._tools_by_plugin.get(existing.plugin_name, []) if name in old_plugin_tools: old_plugin_tools.remove(name) # 注册新工具 tool_def = ToolDefinition( name=name, plugin_name=plugin_name, schema=schema, executor=executor, timeout=timeout, priority=priority, ) self._tools[name] = tool_def # 更新插件工具映射 if plugin_name not in self._tools_by_plugin: self._tools_by_plugin[plugin_name] = [] if name not in self._tools_by_plugin[plugin_name]: self._tools_by_plugin[plugin_name].append(name) logger.debug(f"注册工具: {name} (插件: {plugin_name}, 超时: {timeout}s)") return True def unregister(self, name: str) -> bool: """注销工具""" with self._registry_lock: tool_def = self._tools.pop(name, None) if tool_def: plugin_tools = self._tools_by_plugin.get(tool_def.plugin_name, []) if name in plugin_tools: plugin_tools.remove(name) logger.debug(f"注销工具: {name}") return True return False def unregister_plugin(self, plugin_name: str) -> int: """ 注销插件的所有工具 Args: plugin_name: 插件名 Returns: 注销的工具数量 """ with self._registry_lock: tool_names = self._tools_by_plugin.pop(plugin_name, []) count = 0 for name in tool_names: if self._tools.pop(name, None): count += 1 if count > 0: logger.info(f"注销插件 {plugin_name} 的 {count} 个工具") return count def get(self, name: str) -> Optional[ToolDefinition]: """获取工具定义（O(1) 查找）""" return self._tools.get(name) def get_all_schemas(self) -> List[Dict[str, Any]]: """获取所有工具的 schema 列表（供 LLM 使用）""" return [tool.schema for tool in self._tools.values()] def get_plugin_schemas(self, plugin_name: str) -> List[Dict[str, Any]]: """获取指定插件的工具 schema 列表""" tool_names = self._tools_by_plugin.get(plugin_name, []) return [self._tools[name].schema for name in tool_names if name in self._tools] def list_tools(self) -> List[str]: """列出所有工具名""" return list(self._tools.keys()) def list_plugin_tools(self, plugin_name: str) -> List[str]: """列出插件的所有工具名""" return self._tools_by_plugin.get(plugin_name, []).copy() async def execute( self, name: str, arguments: Dict[str, Any], bot, from_wxid: str, timeout_override: float = None, ) -> Dict[str, Any]: """ 执行工具（带超时保护和统一错误处理） Args: name: 工具名 arguments: 工具参数 bot: WechatHookClient 实例 from_wxid: 消息来源 wxid timeout_override: 覆盖默认超时时间 Returns: 工具执行结果字典 """ from utils.errors import ( ToolNotFoundError, ToolTimeoutError, ToolExecutionError, handle_error ) tool_def = self._tools.get(name) if not tool_def: err = ToolNotFoundError(f"工具 {name} 不存在") return err.to_dict() timeout = timeout_override if timeout_override is not None else tool_def.timeout try: result = await asyncio.wait_for( tool_def.executor(name, arguments, bot, from_wxid), timeout=timeout ) return result except asyncio.TimeoutError: err = ToolTimeoutError( message=f"工具 {name} 执行超时 ({timeout}s)", user_message=f"工具执行超时 ({timeout}s)", context={"tool_name": name, "timeout": timeout} ) logger.warning(err.message) result = err.to_dict() result["timeout"] = True return result except Exception as e: error_result = handle_error( e, context=f"执行工具 {name}", log=True, ) return error_result.to_dict() def get_stats(self) -> Dict[str, Any]: """获取注册统计信息""" return { "total_tools": len(self._tools), "plugins": len(self._tools_by_plugin), "tools_by_plugin": { plugin: len(tools) for plugin, tools in self._tools_by_plugin.items() } } def clear(self): """清空所有注册（用于测试或重置）""" with self._registry_lock: self._tools.clear() self._tools_by_plugin.clear() logger.info("ToolRegistry 已清空") # ==================== 便捷函数 ==================== def get_tool_registry() -> ToolRegistry: """获取工具注册中心实例""" return ToolRegistry() # ==================== 导出列表 ==================== __all__ = [ 'ToolDefinition', 'ToolRegistry', 'get_tool_registry', ]