The Loop That Makes an Agent

# Colab cell 1 — run once. Set your key in the 🔑 panel (name it
# ANTHROPIC_API_KEY) or just paste it when prompted.
!pip install -q anthropic

import os
try:
    from google.colab import userdata
    os.environ["ANTHROPIC_API_KEY"] = userdata.get("ANTHROPIC_API_KEY")
except Exception:
    from getpass import getpass
    os.environ.setdefault("ANTHROPIC_API_KEY", getpass("Anthropic API key: "))

import anthropic

client = anthropic.Anthropic()
MODEL = "claude-sonnet-5"

# A tiny in-memory "notes database" so the tools do real work — no setup,
# no files. Swap this for a real store later; the loop below never changes.
NOTES = {
    "n042": {"title": "Exponential backoff for retries",
             "body": "On 429/5xx, retry with backoff: wait 2**attempt seconds + "
                     "jitter, cap at 60s, give up after 5 tries. Jitter stops "
                     "retries from synchronizing. Never retry 4xx except 429."},
    "n107": {"title": "Caching layer design",
             "body": "Read-through cache, 5-min TTL, key on the normalized query. "
                     "Invalidate on write. Beware the stampede when a hot key "
                     "expires — use a lock or serve-stale-while-revalidate."},
    "n153": {"title": "Rate limiting",
             "body": "Token bucket beats fixed window. Return 429 + Retry-After so "
                     "clients back off deterministically. Limit per API key, not IP."},
}

def search_notes(query: str) -> str:
    words = [w for w in query.lower().split() if len(w) > 2]
    hits = [f"{nid}: {n['title']}" for nid, n in NOTES.items()
            if any(w in (n["title"] + " " + n["body"]).lower() for w in words)]
    return "\n".join(hits[:5]) if hits else f"No notes matched {query!r}."

def read_note(note_id: str) -> str:
    n = NOTES.get(note_id)
    return f"{n['title']}\n\n{n['body']}" if n else f"No note with id {note_id!r}."

IMPL = {"search_notes": search_notes, "read_note": read_note}

TOOLS = [
    {
        "name": "search_notes",
        "description": (
            "Search the local notes database for a keyword. Use whenever the "
            "user asks about anything that might live in their notes. "
            "Returns up to 5 matching snippets with note ids."
        ),
        "input_schema": {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"],
        },
    },
    {
        "name": "read_note",
        "description": "Read one note in full, by id from search_notes results.",
        "input_schema": {
            "type": "object",
            "properties": {"note_id": {"type": "string"}},
            "required": ["note_id"],
        },
    },
]

def run_agent(question: str, max_iterations: int = 10) -> str:
    messages = [{"role": "user", "content": question}]
    for _ in range(max_iterations):
        resp = client.messages.create(
            model=MODEL, max_tokens=2048,
            tools=TOOLS, messages=messages,
        )
        if resp.stop_reason != "tool_use":
            return resp.content[0].text          # model chose to stop

        messages.append({"role": "assistant", "content": resp.content})
        results = []
        for block in resp.content:
            if block.type == "tool_use":
                output = IMPL[block.name](**block.input)
                results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": output,
                })
        messages.append({"role": "user", "content": results})
    raise RuntimeError("max iterations exceeded")   # we'll fix this in lesson 4

print(run_agent("What did I write about backoff?"))

# Colab cell 2 — run the setup cell above first (it defines client,
# MODEL, TOOLS, IMPL). This just adds print() calls to the same loop.
def run_with_trace(question: str, max_iterations: int = 10) -> str:
    messages = [{"role": "user", "content": question}]
    for i in range(max_iterations):
        resp = client.messages.create(
            model=MODEL, max_tokens=2048, tools=TOOLS, messages=messages,
        )
        if resp.stop_reason != "tool_use":
            print(f"[{i}] final answer after {i} tool iterations")
            return resp.content[0].text

        messages.append({"role": "assistant", "content": resp.content})
        results = []
        for block in resp.content:
            if block.type == "tool_use":
                print(f"[{i}] model chose: {block.name}({block.input})")
                output = IMPL[block.name](**block.input)
                print(f"[{i}]   -> {len(output)} chars back")
                results.append({"type": "tool_result",
                                "tool_use_id": block.id, "content": output})
        messages.append({"role": "user", "content": results})
    raise RuntimeError("max iterations exceeded")

print(run_with_trace("What did I write about backoff?"))
print("---")
print(run_with_trace("Summarize my notes on rate limits and caching"))

# A typical run prints something like:
# Run 1: "What did I write about backoff?"
#   [0] model chose: search_notes({'query': 'backoff'})
#   [1] model chose: read_note({'note_id': 'n042'})
#   [2] final answer after 2 tool iterations
#
# Run 2: "Summarize my notes on rate limits AND caching"
#   [0] model chose: search_notes({'query': 'rate limits'})
#   [0] model chose: search_notes({'query': 'caching'})     # parallel!
#   [1] model chose: read_note({'note_id': 'n153'})
#   [1] model chose: read_note({'note_id': 'n107'})
#   [2] final answer after 2 tool iterations