subagent-pattern.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243
  1. """
  2. 子 Agent 模式 - 如何实现用于上下文隔离的 Task 工具。
  3. 核心洞察:用隔离上下文启动子 Agent,防止探索细节填满主对话,
  4. 造成“上下文污染”。
  5. """
  6. import time
  7. import sys
  8. # 假设 client、MODEL、execute_tool 已在其他位置定义
  9. # =============================================================================
  10. # Agent 类型注册表
  11. # =============================================================================
  12. AGENT_TYPES = {
  13. # Explore:只读,用于搜索和分析
  14. "explore": {
  15. "description": "只读 Agent,用于探索代码、查找文件和搜索",
  16. "tools": ["bash", "read_file"], # 没有写入权限!
  17. "prompt": "你是一个探索型 Agent。请搜索并分析,但绝不要修改文件。返回你发现内容的简洁摘要。",
  18. },
  19. # Code:完整权限,用于实现
  20. "code": {
  21. "description": "完整权限 Agent,用于实现功能和修复 bug",
  22. "tools": "*", # 所有工具
  23. "prompt": "你是一个编码 Agent。高效实现请求的修改,并返回变更摘要。",
  24. },
  25. # Plan:只读,用于设计方案
  26. "plan": {
  27. "description": "规划型 Agent,用于设计实现策略",
  28. "tools": ["bash", "read_file"], # 只读
  29. "prompt": "你是一个规划型 Agent。分析代码库并输出编号实现计划。不要做任何修改。",
  30. },
  31. # 在这里添加你自己的类型...
  32. # "test": {
  33. # "description": "Testing agent for running and analyzing tests",
  34. # "tools": ["bash", "read_file"],
  35. # "prompt": "Run tests and report results. Don't modify code.",
  36. # },
  37. }
  38. def get_agent_descriptions() -> str:
  39. """为 Task 工具 schema 生成 Agent 类型描述。"""
  40. return "\n".join(
  41. f"- {name}: {cfg['description']}"
  42. for name, cfg in AGENT_TYPES.items()
  43. )
  44. def get_tools_for_agent(agent_type: str, base_tools: list) -> list:
  45. """
  46. 根据 Agent 类型过滤工具。
  47. '*' 表示所有基础工具。
  48. 否则只允许白名单中的指定工具名。
  49. 注意:子 Agent 不会拿到 Task 工具,以防无限递归。
  50. """
  51. allowed = AGENT_TYPES.get(agent_type, {}).get("tools", "*")
  52. if allowed == "*":
  53. return base_tools # 所有基础工具,但不包括 Task
  54. return [t for t in base_tools if t["name"] in allowed]
  55. # =============================================================================
  56. # Task 工具定义
  57. # =============================================================================
  58. TASK_TOOL = {
  59. "name": "Task",
  60. "description": f"""为聚焦子任务启动一个子 Agent。
  61. 子 Agent 运行在隔离上下文中,看不到父 Agent 的历史。
  62. 用它保持主对话清爽。
  63. Agent 类型:
  64. {get_agent_descriptions()}
  65. 使用示例:
  66. - Task(explore): "查找所有使用 auth 模块的文件"
  67. - Task(plan): "设计数据库迁移策略"
  68. - Task(code): "实现用户注册表单"
  69. """,
  70. "input_schema": {
  71. "type": "object",
  72. "properties": {
  73. "description": {
  74. "type": "string",
  75. "description": "用于进度显示的简短任务名(3-5 个词)"
  76. },
  77. "prompt": {
  78. "type": "string",
  79. "description": "给子 Agent 的详细指令"
  80. },
  81. "agent_type": {
  82. "type": "string",
  83. "enum": list(AGENT_TYPES.keys()),
  84. "description": "要启动的 Agent 类型"
  85. },
  86. },
  87. "required": ["description", "prompt", "agent_type"],
  88. },
  89. }
  90. # =============================================================================
  91. # 子 Agent 执行
  92. # =============================================================================
  93. def run_task(description: str, prompt: str, agent_type: str,
  94. client, model: str, workdir, base_tools: list, execute_tool) -> str:
  95. """
  96. 用隔离上下文执行一个子 Agent 任务。
  97. 关键概念:
  98. 1. 隔离历史:子 Agent 从空白历史开始,没有父 Agent 上下文
  99. 2. 过滤工具:根据 Agent 类型决定工具权限
  100. 3. 专属提示词:提供特化行为
  101. 4. 只返回摘要:父 Agent 只看到最终结果
  102. 参数:
  103. description: 用于进度显示的简短名称
  104. prompt: 给子 Agent 的详细指令
  105. agent_type: AGENT_TYPES 中的键
  106. client: Anthropic 客户端
  107. model: 要使用的模型
  108. workdir: 工作目录
  109. base_tools: 工具定义列表
  110. execute_tool: 执行工具的函数
  111. 返回:
  112. 子 Agent 的最终文本输出
  113. """
  114. if agent_type not in AGENT_TYPES:
  115. return f"错误:未知 Agent 类型 '{agent_type}'"
  116. config = AGENT_TYPES[agent_type]
  117. # Agent 专属系统提示词
  118. sub_system = f"""你是位于 {workdir} 的 {agent_type} 子 Agent。
  119. {config["prompt"]}
  120. 完成任务,并返回清晰简洁的摘要。"""
  121. # 针对此 Agent 类型过滤后的工具
  122. sub_tools = get_tools_for_agent(agent_type, base_tools)
  123. # 关键点:隔离的消息历史!
  124. # 子 Agent 从空白历史开始,看不到父 Agent 的对话
  125. sub_messages = [{"role": "user", "content": prompt}]
  126. # 进度显示
  127. print(f" [{agent_type}] {description}")
  128. start = time.time()
  129. tool_count = 0
  130. # 运行相同的 Agent 循环(但保持安静)
  131. while True:
  132. response = client.messages.create(
  133. model=model,
  134. system=sub_system,
  135. messages=sub_messages,
  136. tools=sub_tools,
  137. max_tokens=8000,
  138. )
  139. # 检查是否完成
  140. if response.stop_reason != "tool_use":
  141. break
  142. # 执行工具
  143. tool_calls = [b for b in response.content if b.type == "tool_use"]
  144. results = []
  145. for tc in tool_calls:
  146. tool_count += 1
  147. output = execute_tool(tc.name, tc.input)
  148. results.append({
  149. "type": "tool_result",
  150. "tool_use_id": tc.id,
  151. "content": output
  152. })
  153. # 更新进度(在同一行原位刷新)
  154. elapsed = time.time() - start
  155. sys.stdout.write(
  156. f"\r [{agent_type}] {description} ... {tool_count} 次工具调用,{elapsed:.1f}s"
  157. )
  158. sys.stdout.flush()
  159. sub_messages.append({"role": "assistant", "content": response.content})
  160. sub_messages.append({"role": "user", "content": results})
  161. # 最终进度更新
  162. elapsed = time.time() - start
  163. sys.stdout.write(
  164. f"\r [{agent_type}] {description} - 已完成({tool_count} 次工具调用,{elapsed:.1f}s)\n"
  165. )
  166. # 只提取并返回最终文本
  167. # This is what the parent agent sees - a clean summary
  168. for block in response.content:
  169. if hasattr(block, "text"):
  170. return block.text
  171. return "(子 Agent 未返回文本)"
  172. # =============================================================================
  173. # 使用示例
  174. # =============================================================================
  175. """
  176. # 在主 Agent 的 execute_tool 函数中:
  177. def execute_tool(name: str, args: dict) -> str:
  178. if name == "Task":
  179. return run_task(
  180. description=args["description"],
  181. prompt=args["prompt"],
  182. agent_type=args["agent_type"],
  183. client=client,
  184. model=MODEL,
  185. workdir=WORKDIR,
  186. base_tools=BASE_TOOLS,
  187. execute_tool=execute_tool # 传入自身用于递归
  188. )
  189. # ... 其他工具 ...
  190. # 在 TOOLS 列表中:
  191. TOOLS = BASE_TOOLS + [TASK_TOOL]
  192. """