""" 工具存储管理器 负责存储和管理外部数据工具的配置和代码 符合 MCP-Server 调用规范 """ import os import json import logging from typing import Dict, List, Optional, Any from datetime import datetime from pathlib import Path logger = logging.getLogger(__name__) class ToolStorage: """工具存储管理器 - 管理外部数据工具的配置和代码文件""" def __init__(self, storage_path: str = None): """ 初始化工具存储 Args: storage_path: 工具存储路径，默认使用 ./tool_storage """ self.storage_path = Path(storage_path or os.getenv("TOOL_STORAGE_PATH", "./tool_storage")) self.storage_path.mkdir(parents=True, exist_ok=True) # 内存缓存 self._cache: Dict[str, Dict] = {} # 加载已有工具 self._load_existing_tools() def _load_existing_tools(self): """从文件系统加载已有的工具配置""" try: for tool_dir in self.storage_path.iterdir(): if tool_dir.is_dir(): config_file = tool_dir / "config.json" if config_file.exists(): with open(config_file, "r", encoding="utf-8") as f: config = json.load(f) self._cache[config["tool_ref_id"]] = config logger.info(f"✅ 已加载 {len(self._cache)} 个工具配置") except Exception as e: logger.error(f"加载工具配置失败: {e}") def _get_tool_dir(self, tool_ref_id: str) -> Path: """获取工具的存储目录""" return self.storage_path / tool_ref_id def save_tool( self, tool_ref_id: str, name: str, description: str, config: Dict, code: str, user_id: str, tenant_id: Optional[str] = None ) -> Dict[str, Any]: """ 保存工具配置和代码 Args: tool_ref_id: 工具唯一标识 name: 工具名称 description: 工具描述 config: 工具配置（含认证、参数等） code: 生成的 Python 代码 user_id: 用户 ID tenant_id: 租户 ID Returns: 保存结果 """ tool_dir = self._get_tool_dir(tool_ref_id) tool_dir.mkdir(parents=True, exist_ok=True) # 保存配置 tool_data = { "tool_ref_id": tool_ref_id, "name": name, "description": description, "config": config, "user_id": user_id, "tenant_id": tenant_id, "created_at": datetime.utcnow().isoformat(), "updated_at": datetime.utcnow().isoformat(), "status": "created" } config_file = tool_dir / "config.json" with open(config_file, "w", encoding="utf-8") as f: json.dump(tool_data, f, ensure_ascii=False, indent=2) # 保存代码文件 code_file = tool_dir / f"{self._to_python_name(name)}.py" with open(code_file, "w", encoding="utf-8") as f: f.write(code) # 更新缓存 tool_data["code_file"] = str(code_file) self._cache[tool_ref_id] = tool_data logger.info(f"✅ 工具已保存: {tool_ref_id}") return {"success": True, "tool_ref_id": tool_ref_id, "code_file": str(code_file)} def get_tool(self, tool_ref_id: str) -> Optional[Dict]: """ 获取工具配置 Args: tool_ref_id: 工具唯一标识 Returns: 工具配置，不存在返回 None """ # 先从缓存获取 if tool_ref_id in self._cache: return self._cache[tool_ref_id] # 从文件系统加载 tool_dir = self._get_tool_dir(tool_ref_id) config_file = tool_dir / "config.json" if config_file.exists(): with open(config_file, "r", encoding="utf-8") as f: tool_data = json.load(f) self._cache[tool_ref_id] = tool_data return tool_data return None def get_tool_code(self, tool_ref_id: str) -> Optional[str]: """ 获取工具代码 Args: tool_ref_id: 工具唯一标识 Returns: 工具代码，不存在返回 None """ tool = self.get_tool(tool_ref_id) if not tool: return None name = tool.get("name", "tool") tool_dir = self._get_tool_dir(tool_ref_id) code_file = tool_dir / f"{self._to_python_name(name)}.py" if code_file.exists(): with open(code_file, "r", encoding="utf-8") as f: return f.read() return None def update_tool(self, tool_ref_id: str, updates: Dict) -> Optional[Dict]: """ 更新工具配置 Args: tool_ref_id: 工具唯一标识 updates: 更新内容 Returns: 更新后的配置 """ tool = self.get_tool(tool_ref_id) if not tool: return None # 更新配置 if "description" in updates: tool["description"] = updates["description"] tool["config"]["description"] = updates["description"] for key in ["url", "method", "headers", "auth", "request_params", "request_body", "timeout"]: if key in updates and updates[key] is not None: tool["config"][key] = updates[key] tool["updated_at"] = datetime.utcnow().isoformat() # 保存到文件 tool_dir = self._get_tool_dir(tool_ref_id) config_file = tool_dir / "config.json" save_data = {k: v for k, v in tool.items() if k != "code_file"} with open(config_file, "w", encoding="utf-8") as f: json.dump(save_data, f, ensure_ascii=False, indent=2) # 更新缓存 self._cache[tool_ref_id] = tool return tool def delete_tool(self, tool_ref_id: str) -> bool: """ 删除工具 Args: tool_ref_id: 工具唯一标识 Returns: 是否删除成功 """ tool_dir = self._get_tool_dir(tool_ref_id) if tool_dir.exists(): import shutil shutil.rmtree(tool_dir) if tool_ref_id in self._cache: del self._cache[tool_ref_id] logger.info(f"✅ 工具已删除: {tool_ref_id}") return True def list_tools(self, user_id: Optional[str] = None, tenant_id: Optional[str] = None) -> List[Dict]: """ 列出工具 Args: user_id: 用户 ID 过滤 tenant_id: 租户 ID 过滤 Returns: 工具列表 """ tools = list(self._cache.values()) if user_id: tools = [t for t in tools if t.get("user_id") == user_id] if tenant_id: tools = [t for t in tools if t.get("tenant_id") == tenant_id] return tools def get_tools_by_refs(self, tool_refs: List[str]) -> List[Dict]: """ 根据 tool_ref_id 列表获取工具 Args: tool_refs: tool_ref_id 列表 Returns: 工具列表（仅返回存在的工具） """ tools = [] for ref in tool_refs: tool = self.get_tool(ref) if tool: tools.append(tool) return tools def check_tool_in_use(self, tool_ref_id: str) -> bool: """ 检查工具是否被 Agent 使用 Args: tool_ref_id: 工具唯一标识 Returns: 是否被使用 """ tool = self.get_tool(tool_ref_id) if not tool: return False return tool.get("status") == "in_use" or bool(tool.get("used_by_agents")) def mark_tool_in_use(self, tool_ref_id: str, agent_name: str) -> bool: """ 标记工具被 Agent 使用 Args: tool_ref_id: 工具唯一标识 agent_name: Agent 名称 Returns: 是否成功 """ tool = self.get_tool(tool_ref_id) if not tool: return False if "used_by_agents" not in tool: tool["used_by_agents"] = [] if agent_name not in tool["used_by_agents"]: tool["used_by_agents"].append(agent_name) tool["status"] = "in_use" # 保存更新 self.update_tool(tool_ref_id, {"status": "in_use"}) return True def _to_python_name(self, name: str) -> str: """将名称转换为 Python 函数名格式""" import re name = re.sub(r'[^a-zA-Z0-9_]', '_', name) if name and name[0].isdigit(): name = '_' + name return name.lower() # 全局实例 tool_storage = ToolStorage()