stage2_tool_loop.py

"""
Stage 2: Core agentic loop + 单个工具 (bash)。

Stage 1 里一次用户输入 = 一次 API 请求 = 一条回复，三者一一对应。
这一版开始出现分层：一个 turn（用户的一次输入）内部可以包含多个 step
（模型每调用一次工具、拿到结果、再决定下一步，就是一个 step）。

事件命名沿用我们在 DeepSeek Harness 架构文档里读到的那套：
    turn/start -> step/start -> tool/call -> tool/result -> step/end -> ... -> turn/end
真实产品（Claude Code、Codex）内部都有类似的状态机，只是不会把它打印出来给你看。
这里我们故意把它摊开打印，方便你盯着循环的形状。
"""

import os
import subprocess
import sys

from dotenv import load_dotenv
from anthropic import Anthropic

load_dotenv()

MODEL = "claude-sonnet-4-5"
MAX_STEPS_PER_TURN = 20  # 安全阀：防止模型陷入工具调用死循环，真实 harness 都有类似上限

BASH_TOOL = {
    "name": "bash",
    "description": (
        "在本机 shell 中执行一条命令，返回 stdout/stderr。"
        "用于查看文件、跑测试、检查环境等。命令有 15 秒超时。"
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "command": {
                "type": "string",
                "description": "要执行的完整 shell 命令",
            }
        },
        "required": ["command"],
    },
}

MAX_OUTPUT_CHARS = 4000  # 粗糙版的输出截断；Claude Code 官方实测阈值约 25k token，这里先用字符数简化


def run_bash(command: str) -> str:
    print(f"  [tool/call] bash: {command}")
    try:
        result = subprocess.run( # 是用来创建和管理子进程的
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=15,
        )
        output = result.stdout + result.stderr
        if len(output) > MAX_OUTPUT_CHARS:
            output = output[:MAX_OUTPUT_CHARS] + "\n...[输出被截断]"
        if not output.strip():
            output = f"(命令执行完毕，退出码 {result.returncode}，无输出)"
    except subprocess.TimeoutExpired:
        output = "错误：命令执行超过 15 秒，已被终止。"
    print(f"  [tool/result] {output[:200]}{'...' if len(output) > 200 else ''}")
    return output


def run_turn(client: Anthropic, messages: list) -> None:
    """一个 turn：喂入当前用户消息，循环执行 step 直到模型给出最终回复或触顶安全阀。"""
    print("[turn/start]")

    for step_index in range(1, MAX_STEPS_PER_TURN + 1):
        print(f"[step/start] step={step_index}")

        response = client.messages.create(
            model=MODEL,
            max_tokens=1024,
            tools=[BASH_TOOL],
            messages=messages,
        )

        messages.append({"role": "assistant", "content": response.content})

        # 把这个 step 里的文字部分打印出来（模型经常会先说"我要看看 xxx"再调用工具）
        for block in response.content:
            if block.type == "text" and block.text.strip():
                print(f"\nClaude: {block.text}\n")

        if response.stop_reason != "tool_use":
            print(f"[step/end] step={step_index} (无工具调用)")
            print(f"[turn/end] stop_reason={response.stop_reason}\n")
            return

        # 执行这个 step 里所有的工具调用（可能不止一个），把结果打包成一条 user 消息
        tool_results = []
        for block in response.content:
            if block.type != "tool_use":
                continue
            if block.name == "bash":
                output = run_bash(block.input["command"])
            else:
                output = f"未知工具: {block.name}"

            tool_results.append(
                {
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": output,
                }
            )

        messages.append({"role": "user", "content": tool_results})
        print(f"[step/end] step={step_index} (已执行 {len(tool_results)} 个工具调用，继续下一 step)")

    print(f"[turn/end] 触顶 MAX_STEPS_PER_TURN={MAX_STEPS_PER_TURN}，强制结束这个 turn\n")


def main():
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("没有找到 ANTHROPIC_API_KEY。检查 .env 文件。")
        sys.exit(1)

    client = Anthropic(api_key=api_key)
    messages = []

    print("Stage 2: 工具循环 (bash)。输入 exit 退出。")
    print("试试问它 '看看当前目录有什么文件' 或 '这台机器的 python 版本是什么'。\n")

    while True:
        user_input = input("你: ").strip()
        if user_input.lower() in ("exit", "quit"):
            break
        if not user_input:
            continue

        messages.append({"role": "user", "content": user_input})
        run_turn(client, messages)


if __name__ == "__main__":
    main()