#!python3 # coding=utf-8 from openai import OpenAI import fire import httpx import subprocess import sys import signal import time from rich.console import Console from rich.panel import Panel from rich.table import Table from rich.progress import Progress, SpinnerColumn, TextColumn from rich.syntax import Syntax from rich import box _VERSION_ = "1.0.0" _MODEL_ = "doubao-seed-2-0-code-preview-260215" _BASE_URL_ = "https://ark.cn-beijing.volces.com/api/v3" _API_KEY_ = "0fb5c683-9002-4bca-9811-b81f46d07e9c" console = Console() def signal_handler(sig, frame): console.print("\n[yellow]程序已中断[/yellow]") sys.exit(0) signal.signal(signal.SIGINT, signal_handler) client = OpenAI( base_url=_BASE_URL_, api_key=_API_KEY_, http_client=httpx.Client(verify=False, timeout=60.0), ) prompt_zh = """你是一位专业的命令行命令生成助手。请根据用户需求生成可直接执行的一条 shell 命令。 1. 优先生成 macOS/zsh 可用命令 2. 只返回一条命令，不要解释，不要代码块，不要序号 3. 除非用户明确要求，不要生成破坏性命令（如 rm -rf、dd、mkfs） 4. 命令应尽量简洁、可靠，可直接复制执行 5. 如果用户给出修改建议，必须基于建议调整命令 需求: tmux创建窗口 输出: tmux new-window 需求: 查找当前目录下所有 md 文件并统计数量 输出: find . -type f -name '*.md' | wc -l """ prompt_en = """You are a professional command generator. Generate one executable shell command based on the user's requirement. 1. Prefer commands that work on macOS/zsh 2. Return exactly one command, no explanation, no code block, no numbering 3. Unless explicitly requested, avoid destructive commands (e.g. rm -rf, dd, mkfs) 4. Keep the command concise and directly executable 5. If user provides revision feedback, you MUST revise accordingly """ def generate_command(requirement, language="zh", feedback_history=None): """ 使用 AI 生成命令 :param requirement: 用户初始需求 :param language: 语言 zh/en :param feedback_history: 修改建议历史 :return: 生成的命令字符串 """ try: if feedback_history is None: feedback_history = [] system_prompt = prompt_zh if language == "zh" else prompt_en user_message = f""" {requirement} """ if feedback_history: history_text = "\n".join([f"- {item}" for item in feedback_history]) if language == "zh": user_message += f""" 以下是用户对上一个命令的修改建议，请严格参考： {history_text} """ else: user_message += f""" The following feedback must be applied to revise the previous command: {history_text} """ completion = client.chat.completions.create( model=_MODEL_, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_message}, ], ) command = completion.choices[0].message.content.strip() # 清理潜在代码块包裹 if command.startswith("```"): command = command.strip("`").strip() if "\n" in command: lines = command.split("\n") if lines and lines[0].lower() in ["bash", "sh", "zsh", "shell"]: lines = lines[1:] command = "\n".join(lines).strip() # 只取第一行，确保是一条命令 if "\n" in command: command = command.split("\n")[0].strip() return command except Exception as e: console.print(f"[red]✗ 生成命令失败: {str(e)}[/red]") return None def run_command(command): """ 执行命令 :param command: 要执行的命令 :return: True 成功，False 失败 """ start_time = time.perf_counter() try: console.print(Panel.fit("[bold cyan]正在执行命令[/bold cyan]", border_style="cyan")) console.print(Syntax(command, "bash", theme="monokai", line_numbers=False)) subprocess.run(command, shell=True, check=True) elapsed = time.perf_counter() - start_time console.print("[bold green]✓ 执行成功！[/bold green]") console.print(f"[dim]执行用时: {elapsed:.2f}s[/dim]") return True except subprocess.CalledProcessError as e: elapsed = time.perf_counter() - start_time console.print(f"[red]✗ 执行失败: {e}[/red]") console.print(f"[dim]执行用时: {elapsed:.2f}s[/dim]") return False def show_documentation(): """显示使用文档""" console.print(Panel.fit( "[bold cyan]命令生成助手 - 使用文档[/bold cyan]", border_style="cyan" )) console.print("\n[bold]📖 简介[/bold]") console.print("输入一句简短需求，生成一条可执行命令；回车直接执行，输入文本可继续迭代优化。\n") console.print("[bold]🚀 基础用法[/bold]") usage_table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2)) usage_table.add_column(style="cyan") usage_table.add_column(style="white") usage_table.add_row("python cmd-generator.py", "交互输入需求并生成命令") usage_table.add_row("python cmd-generator.py --req='tmux创建窗口'", "直接传入需求") usage_table.add_row("python cmd-generator.py --lang=en --req='create tmux window'", "英文生成") console.print(usage_table) console.print() console.print("[bold]⚙️ 参数说明[/bold]") param_table = Table(box=box.ROUNDED, show_header=True, header_style="bold magenta") param_table.add_column("参数", style="cyan", width=20) param_table.add_column("类型", style="yellow", width=10) param_table.add_column("默认值", style="green", width=10) param_table.add_column("说明", style="white") param_table.add_row("--req", "string", "''", "用户需求，例如 'tmux创建窗口'") param_table.add_row("--lang", "string", "'zh'", "生成语言：'zh' 或 'en'") param_table.add_row("--help", "bool", "False", "显示帮助文档") console.print(param_table) console.print() console.print("[bold]💡 交互规则[/bold]") rule_table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2)) rule_table.add_column(style="cyan bold", width=14) rule_table.add_column(style="white") rule_table.add_row("直接回车", "立即执行当前命令") rule_table.add_row("输入文本", "作为修改建议，重新生成命令") rule_table.add_row("输入 n/q", "取消并退出") console.print(rule_table) console.print() console.print(Panel.fit( f"[dim]Version {_VERSION_} | Author: @louisyoungx | ❤️[/dim]", border_style="dim" )) def main(req="", lang="zh", help=False): """ 主函数：根据需求生成并执行命令 :param req: 用户需求 :param lang: 生成语言，zh 或 en :param help: 显示帮助文档 """ if help: show_documentation() return console.print(Panel.fit( "[bold cyan]命令生成助手[/bold cyan]", border_style="cyan" )) requirement = req.strip() if not requirement: requirement = console.input("请输入你的需求: ").strip() if not requirement: console.print("[yellow]未提供需求，已退出[/yellow]") return feedback_history = [] while True: with Progress( SpinnerColumn(), TextColumn("[bold blue]正在生成命令...[/bold blue]"), console=console, transient=True ) as progress: progress.add_task("", total=None) command = generate_command(requirement, lang, feedback_history) if not command: return panel = Panel( Syntax(command, "bash", theme="monokai", line_numbers=False), title="[bold cyan]生成的命令[/bold cyan]", border_style="cyan", box=box.ROUNDED ) console.print(panel) action = console.input("[cyan]回车执行 / 输入修改建议继续生成 / n或q取消[/cyan]: ").strip() if action == "": run_command(command) return if action.lower() in ["n", "q", "quit", "exit"]: console.print("[yellow]已取消[/yellow]") return feedback_history.append(action) if __name__ == "__main__": fire.Fire(main)