gemini_jsonl_batch.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. #!/usr/bin/env python3
  2. """
  3. 使用 Gemini CLI 按固定系统提示词,将指定目录下的 .md 提示词批量转换为 JSONL。
  4. 特点:
  5. - 内置系统提示词,与《Gemini 无头模式 JSONL 规范化指引》一致
  6. - 禁用工具调用 (--allowed-tools ''), 输出纯文本,每个文件生成一行 JSON
  7. - 默认输入目录为仓库根下的 `2/`,输出为 `2/prompts.jsonl`
  8. 用法示例:
  9. python3 gemini_jsonl_batch.py
  10. python3 gemini_jsonl_batch.py --input 2 --output 2/prompts.jsonl --model gemini-2.5-flash
  11. """
  12. from __future__ import annotations
  13. import argparse
  14. import subprocess
  15. import sys
  16. from pathlib import Path
  17. # ==================== 固定系统提示词 ====================
  18. SYS_PROMPT = """{"category_id": 1, "category": "JSONL规范化", "row": 2, "col": 1, "title": "# JSONL 提示词转换器 - 系统提示词", "content": "# JSONL 提示词转换器 - 系统提示词\n\n你是一个专业 的提示词格式转换器。将用户提供的提示词内容转换为标准 JSONL 格式。\n\n## 输出格式\n\n```json\n{\\"title\\\": \\"<标题>\\", \\"content\\\": \\"<完整内容>\\"}\n```\n\n### 字段说明\n\n| 字段 | 类型 | 说明 |\n|------|------|------|\n| `title` | string | 提示词标题,取内容的第一行或前 50 字符 |\n| `content` | string | 完整的提示词内容 |\n\n## 转换规则\n\n1. **标题提取**:\n - 若内容以 `#` 开头,取第一个标题作为 title\n - 否则取前 50 字符(去除换行)\n2. **内容转义**:\n - 换行符 转为 `\\\\n`\n - 双引号转为 `\\\\\"`\n - 反斜杠转为 `\\\\\\\\`\n\n## 输出要求\n\n- 每行一个完整的 JSON 对象\n- 不要添加任何解释、注释或额外文字\n- 不要用 ```json 代码块包裹\n- 直接输出纯 JSONL 内容\n\n## 示例\n\n### 输入\n```\n# Role:智能文档助手\n\n## Background\n用户需要一个能够处理文档的 AI 助手。\n\n## Skills\n- 文档解析\n- 格式转换\n```\n\n### 输出\n```\n{\\"title\\\": \\"# Role:智能文档助手\\", \\"content\\\": \\"# Role:智能文档助手\\\\n\\\\n## Background\\\\n用户需要一个能够处 理文档的 AI 助手。\\\\n\\\\n## Skills\\\\n- 文档解析\\\\n- 格式转换\\"}\n```\n\n---\n\n现在,请将用户提供的内容转换为标准 JSONL 格式。"}"""
  19. def parse_args() -> argparse.Namespace:
  20. parser = argparse.ArgumentParser(
  21. description="使用 Gemini CLI 批量将 .md 提示词转换为 JSONL(固定系统提示词)。",
  22. )
  23. parser.add_argument(
  24. "-i",
  25. "--input",
  26. type=Path,
  27. default=Path("2"),
  28. help="输入目录,遍历其中的 .md 文件(默认:仓库根目录下的 2/)",
  29. )
  30. parser.add_argument(
  31. "-o",
  32. "--output",
  33. type=Path,
  34. default=None,
  35. help="输出 JSONL 文件路径,默认写入 <input>/prompts.jsonl",
  36. )
  37. parser.add_argument(
  38. "-m",
  39. "--model",
  40. default="gemini-2.5-flash",
  41. help="Gemini 模型名称(默认:gemini-2.5-flash)",
  42. )
  43. parser.add_argument(
  44. "--gemini-cmd",
  45. default="gemini",
  46. help="Gemini CLI 可执行文件名或路径(默认:gemini)",
  47. )
  48. parser.add_argument(
  49. "-v",
  50. "--verbose",
  51. action="store_true",
  52. help="输出处理中的详细信息",
  53. )
  54. return parser.parse_args()
  55. def run_gemini(content: str, model: str, cmd: str) -> str:
  56. """调用 Gemini CLI,将单个文本转换为一行 JSON。"""
  57. proc = subprocess.run(
  58. [
  59. cmd,
  60. "-m",
  61. model,
  62. "--output-format",
  63. "text",
  64. "--allowed-tools",
  65. "",
  66. SYS_PROMPT,
  67. ],
  68. input=content.encode("utf-8"),
  69. stdout=subprocess.PIPE,
  70. stderr=subprocess.PIPE,
  71. check=False,
  72. )
  73. stdout = proc.stdout.decode("utf-8", errors="replace").strip()
  74. stderr = proc.stderr.decode("utf-8", errors="replace").strip()
  75. if proc.returncode != 0:
  76. raise RuntimeError(f"Gemini 调用失败 (code={proc.returncode}): {stderr or '无错误输出'}")
  77. if stderr:
  78. # 某些 CLI 可能在 stderr 打印警告,保留但不中断
  79. print(f"⚠️ Gemini 警告: {stderr}", file=sys.stderr)
  80. if not stdout:
  81. raise RuntimeError("Gemini 未返回内容")
  82. # 去除多余行,只保留非空行并合并
  83. lines = [ln for ln in stdout.splitlines() if ln.strip()]
  84. return " ".join(lines).strip()
  85. def main() -> None:
  86. args = parse_args()
  87. input_dir = args.input.resolve()
  88. if not input_dir.exists() or not input_dir.is_dir():
  89. print(f"❌ 输入目录不存在: {input_dir}")
  90. sys.exit(1)
  91. output_path = args.output or (input_dir / "prompts.jsonl")
  92. output_path = output_path.resolve()
  93. output_path.parent.mkdir(parents=True, exist_ok=True)
  94. md_files = sorted(f for f in input_dir.iterdir() if f.suffix == ".md")
  95. if not md_files:
  96. print(f"⚠️ 未找到任何 .md 文件: {input_dir}")
  97. sys.exit(0)
  98. results = []
  99. for md in md_files:
  100. content = md.read_text(encoding="utf-8")
  101. if args.verbose:
  102. print(f"→ 处理 {md.name}")
  103. try:
  104. json_line = run_gemini(content, args.model, args.gemini_cmd)
  105. results.append(json_line)
  106. except Exception as exc: # noqa: BLE001
  107. print(f"❌ 处理失败 {md.name}: {exc}", file=sys.stderr)
  108. with output_path.open("w", encoding="utf-8") as f:
  109. for line in results:
  110. f.write(line + "\n")
  111. print(f"✅ 完成:{len(results)} 条 → {output_path}")
  112. if __name__ == "__main__":
  113. main()