code.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. #!/usr/bin/env python3
  2. """
  3. s01_agent_loop.py - Agent 循环
  4. AI 编码 Agent 的核心秘密可以浓缩成一个模式:
  5. while stop_reason == "tool_use":
  6. response = LLM(messages, tools)
  7. 执行工具
  8. 追加结果
  9. +----------+ +-------+ +---------+
  10. | User | ---> | LLM | ---> | Tool |
  11. | prompt | | | | execute |
  12. +----------+ +---+---+ +----+----+
  13. ^ |
  14. | tool_result |
  15. +---------------+
  16. (循环继续)
  17. 这就是核心循环:把工具结果喂回给模型,直到模型决定停止。
  18. 生产级 Agent 会在这个基础上叠加策略、Hooks 和生命周期控制。
  19. 用法:
  20. pip install anthropic python-dotenv
  21. ANTHROPIC_API_KEY=... python s01_agent_loop/code.py
  22. """
  23. import os
  24. import subprocess
  25. try:
  26. import readline
  27. # macOS 的 libedit 在处理中文输入时有退格问题,这四行修复它
  28. readline.parse_and_bind('set bind-tty-special-chars off')
  29. readline.parse_and_bind('set input-meta on')
  30. readline.parse_and_bind('set output-meta on')
  31. readline.parse_and_bind('set convert-meta off')
  32. except ImportError:
  33. pass
  34. from anthropic import Anthropic
  35. from dotenv import load_dotenv
  36. load_dotenv(override=True)
  37. if os.getenv("ANTHROPIC_BASE_URL"):
  38. os.environ.pop("ANTHROPIC_AUTH_TOKEN", None)
  39. client = Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
  40. MODEL = os.environ["MODEL_ID"]
  41. SYSTEM = f"你是位于 {os.getcwd()}. 使用 bash 解决任务。直接行动,不要只解释。"
  42. # ── 工具定义:只有 bash ────────────────────────────
  43. TOOLS = [{
  44. "name": "bash",
  45. "description": "运行一条 shell 命令。",
  46. "input_schema": {
  47. "type": "object",
  48. "properties": {"command": {"type": "string"}},
  49. "required": ["command"],
  50. },
  51. }]
  52. # ── 工具执行 ────────────────────────────────────────
  53. def run_bash(command: str) -> str:
  54. dangerous = ["rm -rf /", "sudo", "shutdown", "reboot", "> /dev/"]
  55. if any(d in command for d in dangerous):
  56. return "错误:危险命令已被拦截"
  57. try:
  58. r = subprocess.run(command, shell=True, cwd=os.getcwd(),
  59. capture_output=True, text=True, timeout=120)
  60. out = (r.stdout + r.stderr).strip()
  61. return out[:50000] if out else "(无输出)"
  62. except subprocess.TimeoutExpired:
  63. return "错误:执行超时(120 秒)"
  64. except (FileNotFoundError, OSError) as e:
  65. return f"错误:{e}"
  66. # ── 核心模式:while 循环持续调用工具,直到模型停止 ──
  67. def agent_loop(messages: list):
  68. while True:
  69. response = client.messages.create(
  70. model=MODEL, system=SYSTEM, messages=messages,
  71. tools=TOOLS, max_tokens=8000,
  72. )
  73. # 追加 assistant 轮次
  74. messages.append({"role": "assistant", "content": response.content})
  75. # 如果模型没有调用工具,就结束
  76. if response.stop_reason != "tool_use":
  77. return
  78. # 执行每个工具调用并收集结果
  79. results = []
  80. for block in response.content:
  81. if block.type == "tool_use":
  82. print(f"\033[33m$ {block.input['command']}\033[0m")
  83. output = run_bash(block.input["command"])
  84. print(output[:200])
  85. results.append({
  86. "type": "tool_result",
  87. "tool_use_id": block.id,
  88. "content": output,
  89. })
  90. # 将工具结果喂回去,循环继续
  91. messages.append({"role": "user", "content": results})
  92. # ── 入口 ──────────────────────────────────────────
  93. if __name__ == "__main__":
  94. print("s01: Agent 循环")
  95. print("输入问题,回车发送。输入 q 退出。\n")
  96. history = []
  97. while True:
  98. try:
  99. query = input("\033[36ms01 >> \033[0m")
  100. except (EOFError, KeyboardInterrupt):
  101. break
  102. if query.strip().lower() in ("q", "exit", ""):
  103. break
  104. history.append({"role": "user", "content": query})
  105. agent_loop(history)
  106. # 打印模型最终文本回复
  107. response_content = history[-1]["content"]
  108. if isinstance(response_content, list):
  109. for block in response_content:
  110. if getattr(block, "type", None) == "text":
  111. print(block.text)
  112. print()