Coverage for aipyapp/cli/command/custom_command_manager.py: 0%

103 statements  

« prev     ^ index     » next       coverage.py v7.10.3, created at 2025-08-11 12:02 +0200

1import os 

2import re 

3import yaml 

4from pathlib import Path 

5from typing import Dict, List, Optional, Any, Tuple 

6from dataclasses import dataclass, field 

7from loguru import logger 

8 

9try: 

10 import jinja2 

11except ImportError: 

12 raise ImportError("Jinja2 is required for custom commands. Install with: pip install jinja2") 

13 

14from .base import CommandMode 

15 

16 

17@dataclass 

18class CustomCommandConfig: 

19 """Configuration for a custom command""" 

20 name: str 

21 description: str = "" 

22 modes: List[CommandMode] = field(default_factory=lambda: [CommandMode.TASK]) 

23 arguments: List[Dict[str, Any]] = field(default_factory=list) 

24 subcommands: Dict[str, Dict[str, Any]] = field(default_factory=dict) 

25 template_vars: Dict[str, Any] = field(default_factory=dict) 

26 task: bool|None = None # 是否在MAIN模式下创建新任务 

27 

28 

29class CustomCommandManager: 

30 """Manager for custom markdown-based commands""" 

31 

32 def __init__(self): 

33 self.command_dirs: set[Path] = set() 

34 self.commands: Dict[str, 'MarkdownCommand'] = {} 

35 self.log = logger.bind(src="CustomCommandManager") 

36 

37 def add_command_dir(self, command_dir: str | Path): 

38 """Add a custom command directory""" 

39 self.command_dirs.add(Path(command_dir)) 

40 self.log.info(f"Added custom command directory: {command_dir}") 

41 

42 def scan_commands(self) -> List['MarkdownCommand']: 

43 """Scan the command directories for markdown commands""" 

44 commands = [] 

45 

46 for command_dir in self.command_dirs: 

47 if not command_dir.exists(): 

48 # Create default directory if it doesn't exist 

49 if command_dir.name == "custom_commands": 

50 self._ensure_default_command_dir(command_dir) 

51 else: 

52 self.log.warning(f"Command directory does not exist: {command_dir}") 

53 continue 

54 

55 # Scan for .md files 

56 for md_file in command_dir.rglob("*.md"): 

57 try: 

58 command = self._load_command_from_file(md_file) 

59 if command: 

60 commands.append(command) 

61 self.commands[command.name] = command 

62 self.log.info(f"Loaded custom command: {command.name} from {md_file.relative_to(command_dir)}") 

63 except Exception as e: 

64 self.log.error(f"Failed to load command from {md_file}: {e}") 

65 

66 if commands: 

67 self.log.info(f"Loaded {len(commands)} custom commands") 

68 return commands 

69 

70 def _load_command_from_file(self, md_file: Path) -> Optional['MarkdownCommand']: 

71 """Load a command from a markdown file""" 

72 try: 

73 content = md_file.read_text(encoding='utf-8') 

74 frontmatter, body = self._parse_frontmatter(content) 

75 

76 if frontmatter: 

77 # File has frontmatter, parse configuration from it 

78 config = self._parse_command_config(frontmatter, md_file.stem) 

79 else: 

80 # No frontmatter, use default configuration 

81 self.log.info(f"No frontmatter found in {md_file}, using default configuration") 

82 config = self._create_default_config(md_file.stem, content) 

83 body = content # Use entire content as body 

84 

85 # Import here to avoid circular imports 

86 from .markdown_command import MarkdownCommand 

87 return MarkdownCommand(config, body, md_file) 

88 

89 except Exception as e: 

90 self.log.error(f"Error loading command from {md_file}: {e}") 

91 return None 

92 

93 def _parse_frontmatter(self, content: str) -> Tuple[Optional[Dict[str, Any]], str]: 

94 """Parse YAML frontmatter from markdown content""" 

95 # Use regex to match YAML frontmatter pattern 

96 pattern = r'^\s*---\n(.*?)\n---\n?(.*)' 

97 match = re.match(pattern, content, re.DOTALL) 

98 

99 if not match: 

100 return None, content 

101 

102 yaml_content = match.group(1) 

103 body_content = match.group(2) 

104 

105 try: 

106 frontmatter = yaml.safe_load(yaml_content) if yaml_content.strip() else {} 

107 return frontmatter, body_content 

108 except yaml.YAMLError as e: 

109 self.log.error(f"Invalid YAML frontmatter: {e}") 

110 return None, content 

111 

112 def _create_default_config(self, default_name: str, content: str = "") -> CustomCommandConfig: 

113 """Create default configuration for pure markdown files""" 

114 mode = CommandMode.TASK 

115 

116 return CustomCommandConfig( 

117 name=default_name, 

118 description=f"Custom command: {default_name}", 

119 modes=[mode], 

120 arguments=[], 

121 subcommands={}, 

122 template_vars={}, 

123 task=True 

124 ) 

125 

126 def _parse_command_config(self, frontmatter: Dict[str, Any], default_name: str) -> CustomCommandConfig: 

127 """Parse command configuration from frontmatter""" 

128 config = CustomCommandConfig( 

129 name=frontmatter.get('name', default_name), 

130 description=frontmatter.get('description', ''), 

131 arguments=frontmatter.get('arguments', []), 

132 subcommands=frontmatter.get('subcommands', {}), 

133 template_vars=frontmatter.get('template_vars', {}), 

134 task=frontmatter.get('task') 

135 ) 

136 

137 # Parse modes 

138 mode_strings = frontmatter.get('modes', ['task']) 

139 config.modes = [] 

140 for mode_str in mode_strings: 

141 try: 

142 mode = CommandMode[mode_str.upper()] 

143 config.modes.append(mode) 

144 except KeyError: 

145 self.log.warning(f"Invalid command mode: {mode_str}, using TASK as default") 

146 config.modes.append(CommandMode.TASK) 

147 

148 return config 

149 

150 def get_command(self, name: str) -> Optional['MarkdownCommand']: 

151 """Get a custom command by name""" 

152 return self.commands.get(name) 

153 

154 def get_all_commands(self) -> List['MarkdownCommand']: 

155 """Get all loaded custom commands""" 

156 return list(self.commands.values()) 

157 

158 def reload_commands(self) -> List['MarkdownCommand']: 

159 """Reload all custom commands""" 

160 self.commands.clear() 

161 return self.scan_commands() 

162 

163 def validate_command_name(self, name: str, existing_commands: List[str]) -> bool: 

164 """Validate that a command name doesn't conflict with existing commands""" 

165 if name in existing_commands: 

166 self.log.warning(f"Custom command '{name}' conflicts with existing command") 

167 return False 

168 return True