Agent Loop

"""
Agent loop with iteration cap, token budget, and tool failure recovery.
Requires: 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"],
)

# === Configurable guards ===
MAX_ITERATIONS = 30
BUDGET_INPUT_TOKENS  = 1_000_000
BUDGET_OUTPUT_TOKENS =   200_000

# === Tools ===
TOOLS = [
    {
        "name": "search",
        "description": "Search a knowledge base. Returns top matches.",
        "input_schema": {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"],
        },
    },
    {
        "name": "fetch_url",
        "description": "Fetch a URL and return its text content.",
        "input_schema": {
            "type": "object",
            "properties": {"url": {"type": "string", "format": "uri"}},
            "required": ["url"],
        },
    },
]


class UnrecoverableToolError(Exception):
    pass


def execute_tool(name, args):
    if name == "search":
        # ... do real 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):
    """Outer retry for transient API errors (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

        # Inject wrap-up if we are out of budget
        if total_in > BUDGET_INPUT_TOKENS or total_out > BUDGET_OUTPUT_TOKENS:
            messages.append({
                "role": "user",
                "content": "Token budget exhausted. Stop calling tools. "
                           "Reply now with a summary of progress and what remains.",
            })

        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 from where you stopped."})
            continue
        if resp.stop_reason != "tool_use":
            # pause_turn, stop_sequence, or unknown — re-loop with empty user nudge
            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,
            }

    # Hit MAX_ITERATIONS
    messages.append({
        "role": "user",
        "content": "You hit the iteration limit. Stop calling tools and "
                   "summarise what you accomplished, what's left, and any blockers.",
    })
    resp = call_with_retry(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        messages=messages,  # tools omitted to force a text response
    )
    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("Find the three most cited 2025 papers on prompt caching and summarise their findings.")
    print(text)
    print(f"\n[stats] iterations={stats['iters']} input_tokens={stats['in']} output_tokens={stats['out']}")