Zhiwen Assistant MCP

""" 思维导图生成工具 提供文件夹结构的思维导图生成功能。 """ from typing import Dict, Any, List from datetime import datetime from data_access import FileSystemInterface from services import DocumentService def register_mindmap_tools(mcp_server, document_service: DocumentService, file_system: FileSystemInterface): """ 注册思维导图相关工具 Args: mcp_server: FastMCP服务器实例 document_service: 文档服务实例 file_system: 文件系统接口实例 """ @mcp_server.tool() async def folder_docs_generate_mindmap( folder_path: str, format: str = "mermaid", max_depth: int = 5, include_files: bool = True, output_file: str = "" ) -> Dict[str, Any]: """ 生成文件夹结构的思维导图 生成指定文件夹结构的思维导图，支持多种格式输出。 Args: folder_path: 要分析的文件夹路径 format: 思维导图格式 (mermaid, markdown, plantuml) max_depth: 最大显示深度 include_files: 是否包含文件节点 output_file: 输出文件路径（可选） Returns: 思维导图生成结果，包含导图内容、统计信息等 """ try: # 安全验证 folder_path = file_system.security_validator.validate_path(folder_path) # 分析文件夹结构 from services import AnalysisService from services import CacheService cache_service = CacheService(file_system.cache) analysis_service = AnalysisService(file_system, cache_service) structure = await analysis_service.analyze_folder_structure(folder_path, max_depth) # 生成思维导图 if format.lower() == "mermaid": mindmap_content = await generate_mermaid_mindmap(structure, include_files) elif format.lower() == "markdown": mindmap_content = await generate_markdown_mindmap(structure, include_files) elif format.lower() == "plantuml": mindmap_content = await generate_plantuml_mindmap(structure, include_files) else: raise ValueError(f"不支持的格式: {format}") # 保存到文件 if output_file: output_path = file_system.security_validator.validate_path(output_file) await file_system.write_file(output_path, mindmap_content) saved_to = output_path else: saved_to = "" return { "success": True, "format": format, "mindmap_content": mindmap_content, "saved_to": saved_to, "statistics": structure.get("statistics", {}), "generation_time": datetime.now().isoformat(), "total_items": structure.get("statistics", {}).get("total_files", 0) + structure.get("statistics", {}).get("total_folders", 0) } except (RuntimeError, ValueError) as e: return { "success": False, "error": str(e), "error_type": type(e).__name__, "folder_path": folder_path } async def generate_mermaid_mindmap(structure: Dict[str, Any], include_files: bool = True) -> str: """生成Mermaid格式思维导图""" lines = ["mindmap"] lines.append(f" root({structure['name']})") def add_folder_node(folder: Dict[str, Any], prefix: str = " "): lines.append(f"{prefix}{folder['name']}") # 添加子文件夹 for subfolder in folder.get('folders', []): add_folder_node(subfolder, prefix + " ") # 添加文件（如果启用） if include_files: for file_info in folder.get('files', []): file_name = file_info['name'] file_ext = file_info.get('extension', '') file_type = file_info.get('file_type', 'other') # 根据文件类型添加图标 icons = { 'code': '📝', 'documentation': '📚', 'image': '🖼️', 'config': '⚙️', 'other': '📄' } icon = icons.get(file_type, '📄') lines.append(f"{prefix} {icon} {file_name}") # 递归添加文件夹节点 add_folder_node(structure['structure']) return "\n".join(lines) async def generate_markdown_mindmap(structure: Dict[str, Any], include_files: bool = True) -> str: """生成Markdown格式思维导图""" lines = [f"# {structure['name']} 项目结构思维导图"] lines.append("") lines.append("```") def add_node(item: Dict[str, Any], level: int = 0): indent = " " * level if item.get('type') == 'folder': lines.append(f"{indent}📁 {item['name']}/") # 递归处理子项 for sub_item in item.get('items', []): add_node(sub_item, level + 1) elif item.get('type') == 'file' and include_files: file_type = item.get('file_type', 'other') icons = { 'code': '📝', 'documentation': '📚', 'image': '🖼️', 'config': '⚙️', 'other': '📄' } icon = icons.get(file_type, '📄') file_size = item.get('size_human', '') size_info = f" ({file_size})" if file_size else "" lines.append(f"{indent}{icon} {item['name']}{size_info}") # 从根目录开始 add_node(structure['structure']) lines.append("```") lines.append("") # 添加统计信息 stats = structure.get('statistics', {}) lines.append("## 📊 统计信息") lines.append(f"- **总文件数**: {stats.get('total_files', 0)}") lines.append(f"- **总文件夹数**: {stats.get('total_folders', 0)}") lines.append(f"- **总大小**: {stats.get('total_size_human', '0 B')}") lines.append(f"- **代码文件**: {stats.get('code_files', 0)}") lines.append(f"- **文档文件**: {stats.get('doc_files', 0)}") lines.append(f"- **图片文件**: {stats.get('image_files', 0)}") return "\n".join(lines) async def generate_plantuml_mindmap(structure: Dict[str, Any], include_files: bool = True) -> str: """生成PlantUML格式思维导图""" lines = ["@startmindmap"] lines.append(f"* {structure['name']}") def add_folder_node(folder: Dict[str, Any], prefix: str = ""): current_prefix = prefix + "**" lines.append(f"{current_prefix} {folder['name']}") # 添加子文件夹 for subfolder in folder.get('folders', []): add_folder_node(subfolder, prefix + "***") # 添加文件（如果启用） if include_files: file_prefix = prefix + "***" for file_info in folder.get('files', []): file_name = file_info['name'] file_type = file_info.get('file_type', 'other') # 根据文件类型添加样式 styles = { 'code': '<color:blue>', 'documentation': '<color:green>', 'image': '<color:orange>', 'config': '<color:red>', 'other': '' } style = styles.get(file_type, '') end_style = '</color>' if style else '' lines.append(f"{file_prefix} {style}{file_name}{end_style}") # 递归添加文件夹节点 add_folder_node(structure['structure']) lines.append("@endmindmap") return "\n".join(lines)