agent_management/tool_storage.py

"""
工具存储管理器
负责存储和管理外部数据工具的配置和代码
符合 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()