tool-templates.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271
  1. """
  2. 工具模板 - 复制后可按你的 Agent 定制。
  3. 每个工具都需要:
  4. 1. 定义(给模型看的 JSON schema)
  5. 2. 实现(Python 函数)
  6. """
  7. from pathlib import Path
  8. import subprocess
  9. WORKDIR = Path.cwd()
  10. # =============================================================================
  11. # 工具定义(放入 TOOLS 列表)
  12. # =============================================================================
  13. BASH_TOOL = {
  14. "name": "bash",
  15. "description": "运行 shell 命令。适用于 ls、find、grep、git、npm、python 等。",
  16. "input_schema": {
  17. "type": "object",
  18. "properties": {
  19. "command": {
  20. "type": "string",
  21. "description": "要执行的 shell 命令"
  22. }
  23. },
  24. "required": ["command"],
  25. },
  26. }
  27. READ_FILE_TOOL = {
  28. "name": "read_file",
  29. "description": "读取文件内容,返回 UTF-8 文本。",
  30. "input_schema": {
  31. "type": "object",
  32. "properties": {
  33. "path": {
  34. "type": "string",
  35. "description": "文件的相对路径"
  36. },
  37. "limit": {
  38. "type": "integer",
  39. "description": "最多读取多少行(默认全部读取)"
  40. },
  41. },
  42. "required": ["path"],
  43. },
  44. }
  45. WRITE_FILE_TOOL = {
  46. "name": "write_file",
  47. "description": "向文件写入内容。需要时会创建父目录。",
  48. "input_schema": {
  49. "type": "object",
  50. "properties": {
  51. "path": {
  52. "type": "string",
  53. "description": "文件的相对路径"
  54. },
  55. "content": {
  56. "type": "string",
  57. "description": "要写入的内容"
  58. },
  59. },
  60. "required": ["path", "content"],
  61. },
  62. }
  63. EDIT_FILE_TOOL = {
  64. "name": "edit_file",
  65. "description": "在文件中替换一次完全匹配的文本,适合小范围精确修改。",
  66. "input_schema": {
  67. "type": "object",
  68. "properties": {
  69. "path": {
  70. "type": "string",
  71. "description": "文件的相对路径"
  72. },
  73. "old_text": {
  74. "type": "string",
  75. "description": "要查找的精确文本(必须完全匹配)"
  76. },
  77. "new_text": {
  78. "type": "string",
  79. "description": "替换后的文本"
  80. },
  81. },
  82. "required": ["path", "old_text", "new_text"],
  83. },
  84. }
  85. TODO_WRITE_TOOL = {
  86. "name": "TodoWrite",
  87. "description": "更新任务清单,用于规划并跟踪进度。",
  88. "input_schema": {
  89. "type": "object",
  90. "properties": {
  91. "items": {
  92. "type": "array",
  93. "description": "完整任务列表",
  94. "items": {
  95. "type": "object",
  96. "properties": {
  97. "content": {"type": "string", "description": "任务描述"},
  98. "status": {"type": "string", "enum": ["pending", "in_progress", "completed"]},
  99. "activeForm": {"type": "string", "description": "当前进行态描述,例如“正在读取文件”"},
  100. },
  101. "required": ["content", "status", "activeForm"],
  102. },
  103. }
  104. },
  105. "required": ["items"],
  106. },
  107. }
  108. TASK_TOOL_TEMPLATE = """
  109. # 根据 Agent 类型动态生成
  110. TASK_TOOL = {
  111. "name": "Task",
  112. "description": f"为聚焦子任务启动一个子 Agent。\\n\\nAgent 类型:\\n{get_agent_descriptions()}",
  113. "input_schema": {
  114. "type": "object",
  115. "properties": {
  116. "description": {"type": "string", "description": "简短任务名(3-5 个词)"},
  117. "prompt": {"type": "string", "description": "详细指令"},
  118. "agent_type": {"type": "string", "enum": list(AGENT_TYPES.keys())},
  119. },
  120. "required": ["description", "prompt", "agent_type"],
  121. },
  122. }
  123. """
  124. # =============================================================================
  125. # 工具实现
  126. # =============================================================================
  127. def safe_path(p: str) -> Path:
  128. """
  129. 安全检查:确保路径留在工作区内。
  130. 防止 ../../../etc/passwd 这类路径逃逸攻击。
  131. """
  132. path = (WORKDIR / p).resolve()
  133. if not path.is_relative_to(WORKDIR):
  134. raise ValueError(f"路径逃逸出工作区:{p}")
  135. return path
  136. def run_bash(command: str) -> str:
  137. """
  138. 带安全检查地执行 shell 命令。
  139. 安全特性:
  140. - 阻止明显危险的命令
  141. - 60 秒超时
  142. - 输出截断到 50KB
  143. """
  144. dangerous = ["rm -rf /", "sudo", "shutdown", "reboot", "> /dev/"]
  145. if any(d in command for d in dangerous):
  146. return "错误:已阻止危险命令"
  147. try:
  148. result = subprocess.run(
  149. command,
  150. shell=True,
  151. cwd=WORKDIR,
  152. capture_output=True,
  153. text=True,
  154. timeout=60
  155. )
  156. output = (result.stdout + result.stderr).strip()
  157. return output[:50000] if output else "(无输出)"
  158. except subprocess.TimeoutExpired:
  159. return "错误:命令执行超时(60 秒)"
  160. except Exception as e:
  161. return f"错误:{e}"
  162. def run_read_file(path: str, limit: int = None) -> str:
  163. """
  164. 读取文件内容,可选行数限制。
  165. 特性:
  166. - 安全路径解析
  167. - 大文件可限制读取行数
  168. - 输出截断到 50KB
  169. """
  170. try:
  171. text = safe_path(path).read_text()
  172. lines = text.splitlines()
  173. if limit and limit < len(lines):
  174. lines = lines[:limit]
  175. lines.append(f"...(还有 {len(text.splitlines()) - limit} 行)")
  176. return "\n".join(lines)[:50000]
  177. except Exception as e:
  178. return f"错误:{e}"
  179. def run_write_file(path: str, content: str) -> str:
  180. """
  181. 向文件写入内容,需要时创建父目录。
  182. 特性:
  183. - 安全路径解析
  184. - 自动创建父目录
  185. - 返回写入字节数用于确认
  186. """
  187. try:
  188. fp = safe_path(path)
  189. fp.parent.mkdir(parents=True, exist_ok=True)
  190. fp.write_text(content)
  191. return f"已写入 {len(content)} 字节到 {path}"
  192. except Exception as e:
  193. return f"错误:{e}"
  194. def run_edit_file(path: str, old_text: str, new_text: str) -> str:
  195. """
  196. 在文件中替换一次完全匹配的文本(精确修改)。
  197. 特性:
  198. - 精确字符串匹配(不是正则)
  199. - 只替换首次出现的位置(更安全)
  200. - 找不到文本时返回清晰错误
  201. """
  202. try:
  203. fp = safe_path(path)
  204. content = fp.read_text()
  205. if old_text not in content:
  206. return f"错误:在 {path} 中未找到目标文本"
  207. new_content = content.replace(old_text, new_text, 1)
  208. fp.write_text(new_content)
  209. return f"已编辑 {path}"
  210. except Exception as e:
  211. return f"错误:{e}"
  212. # =============================================================================
  213. # 分发器模式
  214. # =============================================================================
  215. def execute_tool(name: str, args: dict) -> str:
  216. """
  217. 将工具调用分发到对应实现。
  218. 这种模式便于添加新工具:
  219. 1. 将工具定义加入 TOOLS 列表
  220. 2. 添加实现函数
  221. 3. 在这个分发器中增加分支
  222. """
  223. if name == "bash":
  224. return run_bash(args["command"])
  225. if name == "read_file":
  226. return run_read_file(args["path"], args.get("limit"))
  227. if name == "write_file":
  228. return run_write_file(args["path"], args["content"])
  229. if name == "edit_file":
  230. return run_edit_file(args["path"], args["old_text"], args["new_text"])
  231. # 在这里继续添加更多工具...
  232. return f"未知工具:{name}"