Agent 循环

"""
Agent 循环:迭代上限、token 预算、工具失败恢复。
依赖:pip install anthropic
"""
import os, time, random
from anthropic import Anthropic
from anthropic import APIStatusError

client = Anthropic(
    base_url="https://buzzai.cc",
    api_key=os.environ["BUZZ_API_KEY"],
)

# === 可调护栏 ===
MAX_ITERATIONS = 30
BUDGET_INPUT_TOKENS  = 1_000_000
BUDGET_OUTPUT_TOKENS =   200_000

# === 工具 ===
TOOLS = [
    {
        "name": "search",
        "description": "搜索知识库,返回 top 匹配。",
        "input_schema": {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"],
        },
    },
    {
        "name": "fetch_url",
        "description": "拉取 URL,返回正文文本。",
        "input_schema": {
            "type": "object",
            "properties": {"url": {"type": "string", "format": "uri"}},
            "required": ["url"],
        },
    },
]


class UnrecoverableToolError(Exception):
    pass


def execute_tool(name, args):
    if name == "search":
        # ... 真实搜索 ...
        return f"Top 3 results for {args['query']!r}: ..."
    if name == "fetch_url":
        import urllib.request
        for attempt in range(3):
            try:
                with urllib.request.urlopen(args["url"], timeout=10) as r:
                    return r.read().decode("utf-8", errors="replace")[:8000]
            except urllib.error.HTTPError as e:
                if e.code in (403, 404):
                    raise UnrecoverableToolError(f"HTTP {e.code} for {args['url']}")
                time.sleep(2 ** attempt + random.random())
        raise RuntimeError(f"fetch_url failed after retries: {args['url']}")
    raise UnrecoverableToolError(f"unknown tool: {name}")


def call_with_retry(**kwargs):
    """API 瞬时错误外层重试(429/500/529)。"""
    for attempt in range(5):
        try:
            return client.messages.create(**kwargs)
        except APIStatusError as e:
            if e.status_code in (429, 500, 503, 529):
                wait = (2 ** attempt) + random.random()
                time.sleep(min(wait, 60))
                continue
            raise
    raise RuntimeError("API failed after retries")


def run_agent(user_request: str):
    messages = [{"role": "user", "content": user_request}]
    total_in = total_out = 0
    iteration = 0
    abort = False

    while iteration < MAX_ITERATIONS:
        iteration += 1

        # 超预算就注入收尾指令
        if total_in > BUDGET_INPUT_TOKENS or total_out > BUDGET_OUTPUT_TOKENS:
            messages.append({
                "role": "user",
                "content": "Token 预算已耗尽。停止调用工具。"
                           "立即给出进展摘要和剩余事项。",
            })

        resp = call_with_retry(
            model="claude-sonnet-4-6",
            max_tokens=4096,
            tools=TOOLS,
            messages=messages,
        )
        total_in  += resp.usage.input_tokens + (resp.usage.cache_read_input_tokens or 0)
        total_out += resp.usage.output_tokens

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

        if resp.stop_reason == "end_turn":
            return _final_text(resp), {"iters": iteration, "in": total_in, "out": total_out}
        if resp.stop_reason == "refusal":
            return "[refused]", {"iters": iteration, "in": total_in, "out": total_out}
        if resp.stop_reason == "max_tokens":
            messages.append({"role": "user", "content": "从刚才停下的地方继续。"})
            continue
        if resp.stop_reason != "tool_use":
            # pause_turn / stop_sequence / 未知 — 直接再循环一次
            continue

        tool_results = []
        for block in resp.content:
            if block.type != "tool_use":
                continue
            try:
                output = execute_tool(block.name, block.input)
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": output,
                })
            except UnrecoverableToolError as e:
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": f"FATAL: {e}",
                    "is_error": True,
                })
                abort = True
            except Exception as e:
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": f"ERROR (recoverable): {e}",
                    "is_error": True,
                })

        messages.append({"role": "user", "content": tool_results})
        if abort:
            return "[aborted: unrecoverable tool error]", {
                "iters": iteration, "in": total_in, "out": total_out,
            }

    # 命中 MAX_ITERATIONS
    messages.append({
        "role": "user",
        "content": "已达迭代上限。停止调用工具,"
                   "总结完成了什么、还剩什么、有什么阻塞。",
    })
    resp = call_with_retry(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        messages=messages,  # 故意去掉 tools,强制文本回复
    )
    return _final_text(resp), {"iters": iteration + 1, "in": total_in, "out": total_out}


def _final_text(resp):
    return "\n".join(b.text for b in resp.content if b.type == "text")


if __name__ == "__main__":
    text, stats = run_agent("找出 2025 年关于 prompt caching 引用最多的三篇论文,总结它们的发现。")
    print(text)
    print(f"\n[stats] iterations={stats['iters']} input_tokens={stats['in']} output_tokens={stats['out']}")