Coverage for aipyapp/cli/command/markdown_command.py: 0%
193 statements
« prev ^ index » next coverage.py v7.10.3, created at 2025-08-11 12:45 +0200
1import argparse
2from pathlib import Path
3from typing import Any, Dict, List, Optional, NamedTuple
4from collections import OrderedDict
5import re
6import io
7import sys
8from contextlib import redirect_stdout, redirect_stderr
10from rich.markdown import Markdown
11from jinja2 import Template, Environment, BaseLoader
13from .base import Completable, CommandMode
14from .base_parser import ParserCommand
15from .custom_command_manager import CustomCommandConfig
16from .result import TaskModeResult
18class CodeBlock(NamedTuple):
19 """Represents a code block with its metadata"""
20 language: str
21 code: str
22 start_pos: int
23 end_pos: int
26class 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
32class StringTemplateLoader(BaseLoader):
33 """Simple template loader for string templates"""
35 def __init__(self, template_string: str):
36 self.template_string = template_string
38 def get_source(self, environment, template):
39 return self.template_string, None, lambda: True
42class CodeExecutor:
43 """Unified code execution engine for both TASK and MAIN modes"""
45 def __init__(self, ctx):
46 self.ctx = ctx
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
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()
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 }
69 with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):
70 exec(code, exec_globals)
72 stdout_content = stdout_buffer.getvalue()
73 stderr_content = stderr_buffer.getvalue()
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()}")
81 return "\n".join(output) if output else None
83 except Exception as e:
84 return f"Python 执行错误: {e}"
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)
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()}")
98 return "\n".join(output) if output else None
100 except Exception as e:
101 return f"Shell 执行错误: {e}"
104class ContentParser:
105 """Unified content parser for extracting code blocks"""
107 EXECUTABLE_PATTERN = re.compile(r'````(python|bash|shell|exec)\n(.*?)````', re.DOTALL | re.IGNORECASE)
109 def parse_content(self, content: str) -> ParsedContent:
110 """Parse content into markdown parts and code blocks"""
111 parts = []
112 last_end = 0
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))
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()
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))
140 return ParsedContent(parts=parts, num_code_blocks=num_code_blocks)
142class MarkdownCommand(ParserCommand):
143 """Custom command loaded from markdown file"""
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
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
157 # Initialize manager to None, will be set by CommandManager
158 self.manager = None
160 # Template environment
161 self.template_env = Environment(loader=StringTemplateLoader(content))
162 self.template = self.template_env.from_string(content)
164 # Unified content parser
165 self.content_parser = ContentParser()
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)
172 # Add universal --test argument for all custom commands
173 parser.add_argument(
174 '--test',
175 action='store_true',
176 help='测试模式：预览命令输出，不发送给LLM'
177 )
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 )
187 # Add arguments for subcommand
188 for arg_config in subcmd_config.get('arguments', []):
189 self._add_argument_from_config(subcmd_parser, arg_config)
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')
200 kwargs = {'help': help_text}
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
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'] = '?'
223 parser.add_argument(name, **kwargs)
226 def cmd(self, args, ctx):
227 """Execute the main command"""
228 return self._execute_with_content(args, ctx)
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)
236 # Render template with arguments
237 rendered_content = self._render_template(args, subcommand, ctx)
239 # Parse content once
240 parsed_content = self.content_parser.parse_content(rendered_content)
242 # 渲染代码块
243 final_content = self._render_code_block(parsed_content, ctx)
245 # 检查是否是测试模式
246 is_test_mode = getattr(args, 'test', False)
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
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
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)
266 ctx.console.print(Markdown(final_content))
267 return True
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]
274 result_parts = []
275 executor = CodeExecutor(ctx)
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("```")
287 return "\n".join(result_parts)
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 = {}
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
299 # Add config template variables
300 template_vars.update(self.config.template_vars)
302 # Add subcommand info
303 if subcommand:
304 template_vars['subcommand'] = subcommand
306 # Add context object for main mode commands
307 if ctx:
308 template_vars['ctx'] = ctx
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
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