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

1 import argparse

2 from pathlib import Path

3 from typing import Any, Dict, List, Optional, NamedTuple

4 from collections import OrderedDict

5 import re

6 import io

7 import sys

8 from contextlib import redirect_stdout, redirect_stderr

9

10 from rich.markdown import Markdown

11 from jinja2 import Template, Environment, BaseLoader

12

13 from .base import Completable, CommandMode

14 from .base_parser import ParserCommand

15 from .custom_command_manager import CustomCommandConfig

16 from .result import TaskModeResult

17

18 class CodeBlock(NamedTuple):

19 """Represents a code block with its metadata"""

20 language: str

21 code: str

22 start_pos: int

23 end_pos: int

24

25

26 class ParsedContent(NamedTuple):

27 """Represents parsed content with separated markdown and code blocks"""

28 parts: List[tuple] # List of ('markdown', content) or ('code', CodeBlock) tuples

29 num_code_blocks: int

30

31

32 class StringTemplateLoader(BaseLoader):

33 """Simple template loader for string templates"""

34

35 def __init__(self, template_string: str):

36 self.template_string = template_string

37

38 def get_source(self, environment, template):

39 return self.template_string, None, lambda: True

40

41

42 class CodeExecutor:

43 """Unified code execution engine for both TASK and MAIN modes"""

44

45 def __init__(self, ctx):

46 self.ctx = ctx

47

48 def execute_code_block(self, code_block: CodeBlock) -> Optional[str]:

49 """Execute a code block and return output"""

50 if code_block.language == 'python':

51 return self._execute_python(code_block.code)

52 elif code_block.language in ['bash', 'shell', 'exec']:

53 return self._execute_shell(code_block.code)

54 return None

55

56 def _execute_python(self, code: str) -> Optional[str]:

57 """Execute Python code and capture output"""

58 stdout_buffer = io.StringIO()

59 stderr_buffer = io.StringIO()

60

61 try:

62 exec_globals = {

63 'ctx': self.ctx,

64 'tm': getattr(self.ctx, 'tm', None),

65 'console': self.ctx.console,

66 'print': print

67 }

68

69 with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):

70 exec(code, exec_globals)

71

72 stdout_content = stdout_buffer.getvalue()

73 stderr_content = stderr_buffer.getvalue()

74

75 output = []

76 if stdout_content.strip():

77 output.append(stdout_content.strip())

78 if stderr_content.strip():

79 output.append(f"错误: {stderr_content.strip()}")

80

81 return "\n".join(output) if output else None

82

83 except Exception as e:

84 return f"Python 执行错误: {e}"

85

86 def _execute_shell(self, code: str) -> Optional[str]:

87 """Execute shell code and capture output"""

88 try:

89 import subprocess

90 result = subprocess.run(code, shell=True, capture_output=True, text=True)

91

92 output = []

93 if result.stdout.strip():

94 output.append(result.stdout.strip())

95 if result.stderr.strip():

96 output.append(f"错误: {result.stderr.strip()}")

97

98 return "\n".join(output) if output else None

99

100 except Exception as e:

101 return f"Shell 执行错误: {e}"

102

103

104 class ContentParser:

105 """Unified content parser for extracting code blocks"""

106

107 EXECUTABLE_PATTERN = re.compile(r'````(python|bash|shell|exec)\n(.*?)````', re.DOTALL | re.IGNORECASE)

108

109 def parse_content(self, content: str) -> ParsedContent:

110 """Parse content into markdown parts and code blocks"""

111 parts = []

112 last_end = 0

113

114 num_code_blocks = 0

115 for match in self.EXECUTABLE_PATTERN.finditer(content):

116 # Add markdown before code block

117 if match.start() > last_end:

118 markdown_content = content[last_end: match.start()].strip()

119 if markdown_content:

120 parts.append(('markdown', markdown_content))

121

122 # Add code block

123 code_block = CodeBlock(

124 language=match.group(1).lower(),

125 code=match.group(2).strip(),

126 start_pos=match.start(),

127 end_pos=match.end()

128 )

129 parts.append(('code', code_block))

130 num_code_blocks += 1

131 last_end = match.end()

132

133 # Add remaining markdown after last code block

134 if last_end < len(content):

135 remaining_content = content[last_end:] if last_end else content

136 remaining_content = remaining_content.strip()

137 if remaining_content:

138 parts.append(('markdown', remaining_content))

139

140 return ParsedContent(parts=parts, num_code_blocks=num_code_blocks)

141

142 class MarkdownCommand(ParserCommand):

143 """Custom command loaded from markdown file"""

144

145 def __init__(self, config: CustomCommandConfig, content: str, file_path: Path):

146 super().__init__()

147 self.config = config

148 self.content = content

149 self.file_path = file_path

150

151 # Set command properties from config

152 self.name = config.name

153 self.desc = config.description

154 self.description = config.description #TODO: fix this

155 self.modes = config.modes

156

157 # Initialize manager to None, will be set by CommandManager

158 self.manager = None

159

160 # Template environment

161 self.template_env = Environment(loader=StringTemplateLoader(content))

162 self.template = self.template_env.from_string(content)

163

164 # Unified content parser

165 self.content_parser = ContentParser()

166

167 def add_arguments(self, parser):

168 """Add arguments defined in the command configuration"""

169 for arg_config in self.config.arguments:

170 self._add_argument_from_config(parser, arg_config)

171

172 # Add universal --test argument for all custom commands

173 parser.add_argument(

174 '--test',

175 action='store_true',

176 help='测试模式：预览命令输出，不发送给LLM'

177 )

178

179 def add_subcommands(self, subparsers):

180 """Add subcommands defined in the configuration"""

181 for subcmd_name, subcmd_config in self.config.subcommands.items():

182 subcmd_parser = subparsers.add_parser(

183 subcmd_name,

184 help=subcmd_config.get('description', '')

185 )

186

187 # Add arguments for subcommand

188 for arg_config in subcmd_config.get('arguments', []):

189 self._add_argument_from_config(subcmd_parser, arg_config)

190

191 def _add_argument_from_config(self, parser, arg_config: Dict[str, Any]):

192 """Add a single argument from configuration"""

193 name = arg_config['name']

194 arg_type = arg_config.get('type', 'str')

195 required = arg_config.get('required', False)

196 default = arg_config.get('default')

197 help_text = arg_config.get('help', '')

198 choices = arg_config.get('choices')

199

200 kwargs = {'help': help_text}

201

202 # Handle different argument types

203 if arg_type == 'flag':

204 kwargs['action'] = 'store_true'

205 elif arg_type in ('str', 'int', 'float'):

206 kwargs['type'] = eval(arg_type)

207 if default is not None:

208 kwargs['default'] = default

209 elif arg_type == 'choice' and choices:

210 kwargs['choices'] = choices

211 if default is not None:

212 kwargs['default'] = default

213

214 # Add required flag for positional arguments

215 if not name.startswith('-') and required and default is None:

216 # For positional arguments, required is implicit

217 pass

218 elif name.startswith('-') and required:

219 kwargs['required'] = True

220 elif not name.startswith('-') and not required:

221 kwargs['nargs'] = '?'

222

223 parser.add_argument(name, **kwargs)

224

225

226 def cmd(self, args, ctx):

227 """Execute the main command"""

228 return self._execute_with_content(args, ctx)

229

230 def _execute_with_content(self, args, ctx, subcommand: Optional[str] = None):

231 """Execute command by rendering template and processing content"""

232 # Get subcommand from args if available

233 if not subcommand:

234 subcommand = getattr(args, 'subcommand', None)

235

236 # Render template with arguments

237 rendered_content = self._render_template(args, subcommand, ctx)

238

239 # Parse content once

240 parsed_content = self.content_parser.parse_content(rendered_content)

241

242 # 渲染代码块

243 final_content = self._render_code_block(parsed_content, ctx)

244

245 # 检查是否是测试模式

246 is_test_mode = getattr(args, 'test', False)

247

248 if is_test_mode:

249 # 测试模式：始终显示输出，不发送给LLM

250 ctx.console.print("[yellow]🧪 测试模式 - 以下是命令输出预览：[/yellow]")

251 ctx.console.print(Markdown(final_content))

252 ctx.console.print("[yellow]💡 移除 --test 参数即可正常执行命令[/yellow]")

253 return True

254

255 # 判断是否发送给LLM

256 should_send_to_llm = self.config.task

257 if should_send_to_llm is None:

258 should_send_to_llm = True if ctx.task else False

259

260 if should_send_to_llm:

261 if ctx.task:

262 return ctx.task.run(final_content, title=self.desc)

263 else:

264 return TaskModeResult(instruction=final_content, title=self.desc)

265

266 ctx.console.print(Markdown(final_content))

267 return True

268

269 def _render_code_block(self, parsed_content: ParsedContent, ctx) -> str:

270 """Render code block"""

271 if parsed_content.num_code_blocks == 0:

272 return parsed_content.parts[0][1]

273

274 result_parts = []

275 executor = CodeExecutor(ctx)

276

277 for part_type, content in parsed_content.parts:

278 if part_type == 'markdown':

279 result_parts.append(content)

280 elif part_type == 'code':

281 output = executor.execute_code_block(content)

282 if output:

283 result_parts.append("```")

284 result_parts.append(output)

285 result_parts.append("```")

286

287 return "\n".join(result_parts)

288

289 def _render_template(self, args, subcommand: Optional[str] = None, ctx=None) -> str:

290 """Render the command template with arguments"""

291 # Build template variables

292 template_vars = {}

293

294 # Add argument values

295 for key, value in vars(args).items():

296 if key not in ('subcommand', 'raw_args'):

297 template_vars[key] = value

298

299 # Add config template variables

300 template_vars.update(self.config.template_vars)

301

302 # Add subcommand info

303 if subcommand:

304 template_vars['subcommand'] = subcommand

305

306 # Add context object for main mode commands

307 if ctx:

308 template_vars['ctx'] = ctx

309

310 try:

311 return self.template.render(**template_vars)

312 except Exception as e:

313 if self.log:

314 self.log.error(f"Template rendering error: {e}")

315 return self.content

316

317 def get_arg_values(self, arg, subcommand=None, partial_value=''):

318 """Get argument values for auto-completion"""

319 choices = arg.get('choices')

320 if choices:

321 return list(choices.values())

322 return None