Zhiwen Assistant MCP

""" 模板管理工具 提供模板管理、自定义模板创建等功能。 """ from typing import Dict, Any, List, Optional from datetime import datetime from services import TemplateService class TemplateFilter: """模板过滤器类""" @staticmethod def filter_by_type(templates: List[str], template_type: str) -> List[str]: """ 根据类型过滤模板 Args: templates: 模板名称列表 template_type: 模板类型 (all, readme, mindmap, custom) Returns: 过滤后的模板列表 """ template_type = template_type.lower() if template_type == "readme": return [t for t in templates if 'readme' in t.lower()] elif template_type == "mindmap": return [t for t in templates if 'mindmap' in t.lower() or 'mermaid' in t.lower()] elif template_type == "custom": prefixes = ['readme', 'mindmap', 'mermaid'] return [t for t in templates if not any(prefix in t.lower() for prefix in prefixes)] else: return templates class TemplateDetailsCollector: """模板详细信息收集器""" def __init__(self, template_service: TemplateService): self.template_service = template_service def collect_details(self, template_names: List[str]) -> List[Dict[str, Any]]: """ 收集模板详细信息 Args: template_names: 模板名称列表 Returns: 模板详细信息列表 """ template_details = [] for template_name in template_names: try: info = self.template_service.get_template_info(template_name) template_details.append(info) except (RuntimeError, ValueError) as e: # 如果获取详细信息失败，至少包含基本名称 template_details.append({ 'name': template_name, 'error': str(e) }) return template_details class TemplateToolRegistrar: """模板工具注册器""" def __init__(self, mcp_server, template_service: TemplateService): self.mcp_server = mcp_server self.template_service = template_service self.filter = TemplateFilter() self.details_collector = TemplateDetailsCollector(template_service) def register_list_templates_tool(self): """注册列出模板工具""" @self.mcp_server.tool() async def folder_docs_list_templates( template_type: str = "all" ) -> Dict[str, Any]: """ 列出可用的模板 列出系统中所有可用的文档生成模板。 Args: template_type: 模板类型过滤 (all, readme, mindmap, custom) Returns: 可用模板列表，包含模板名称、类型、大小等信息 """ try: # 获取所有模板 all_templates = self.template_service.list_available_templates() # 按类型过滤 filtered_templates = self.filter.filter_by_type(all_templates, template_type) # 获取模板详细信息 template_details = self.details_collector.collect_details(filtered_templates) return { 'templates': template_details, 'total_count': len(template_details), 'filter_type': template_type } except (RuntimeError, ValueError) as e: return { 'error': f"列出模板失败: {str(e)}", 'templates': [], 'total_count': 0 } def register_get_template_tool(self): """注册获取模板内容工具""" @self.mcp_server.tool() async def folder_docs_get_template( template_name: str ) -> Dict[str, Any]: """ 获取模板内容 获取指定模板的详细内容和信息。 Args: template_name: 模板名称 Returns: 模板内容和详细信息 """ try: content = self.template_service.get_template_content(template_name) info = self.template_service.get_template_info(template_name) return { 'template_name': template_name, 'content': content, 'info': info } except (RuntimeError, ValueError) as e: return { 'error': f"获取模板失败: {str(e)}", 'template_name': template_name, 'content': None, 'info': None } def register_create_template_tool(self): """注册创建模板工具""" @self.mcp_server.tool() async def folder_docs_create_template( template_name: str, template_content: str, template_type: str = "custom", description: str = "" ) -> Dict[str, Any]: """ 创建自定义模板 创建新的自定义模板，用于后续文档生成。 Args: template_name: 模板名称 template_content: 模板内容 template_type: 模板类型 description: 模板描述 Returns: 创建结果 """ try: result = self.template_service.create_custom_template( template_name=template_name, content=template_content, template_type=template_type, description=description ) return { 'success': True, 'template_name': template_name, 'message': '模板创建成功', 'result': result } except (RuntimeError, ValueError) as e: return { 'success': False, 'template_name': template_name, 'error': f"创建模板失败: {str(e)}" } def register_delete_template_tool(self): """注册删除模板工具""" @self.mcp_server.tool() async def folder_docs_delete_template( template_name: str ) -> Dict[str, Any]: """ 删除模板 删除指定的自定义模板。 Args: template_name: 模板名称 Returns: 删除结果 """ try: result = self.template_service.delete_template(template_name) return { 'success': True, 'template_name': template_name, 'message': '模板删除成功', 'result': result } except (RuntimeError, ValueError) as e: return { 'success': False, 'template_name': template_name, 'error': f"删除模板失败: {str(e)}" } def register_all_tools(self): """注册所有模板工具""" self.register_list_templates_tool() self.register_get_template_tool() self.register_create_template_tool() self.register_delete_template_tool() def register_template_tools(mcp_server, template_service: TemplateService): """ 注册模板管理相关工具 Args: mcp_server: FastMCP服务器实例 template_service: 模板服务实例 """ registrar = TemplateToolRegistrar(mcp_server, template_service) registrar.register_all_tools()