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

1 import os

2 import re

3 import yaml

4 from pathlib import Path

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

6 from dataclasses import dataclass, field

7 from loguru import logger

8

9 try:

10 import jinja2

11 except ImportError:

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

13

14 from .base import CommandMode

15

16

17@dataclass

18 class 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

29 class 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