Homework — Extend Your Hand-Rolled Agent

Lesson 02 · Homework · ~60 min · three coding extensions of agent.py

In Lab 02 you built an agent.py that wires together a single Claude call, two mock tools, and the four stop conditions. The three tasks below take that file and stretch it in directions you will meet again in production: adding a new tool, logging every turn for later analysis, and swapping models with a graceful fallback. Each task is self-contained — you can do them in any order, but doing them in sequence on the same file is the most honest exercise.

How to do this homework

Keep your Lab 02 agent.py on hand and edit it directly — do not start a fresh file. Run the script after every change with a real question. When you are stuck, peek at the solutions page, but write the code yourself first.

Watch your spend

Task 3 deliberately triggers an error and a model upgrade. Keep the cost cap from Lab 02 in place (MAX_COST_USD = 0.10 or lower) so a misbehaving loop cannot run away with your wallet. If you raised the cap during the lab, lower it again before starting.

Task 1 — Add a third tool: `search_notes(query)`

Extend agent.py with a new tool that searches a small in-process dictionary of notes. Use this exact dict so your output matches the solution:

NOTES = {
    "kyiv-trip": "Pack thermals, the forecast was rainy.",
    "groceries": "milk, bread, oats",
}

Write search_notes(query: str) -> str that returns matching notes (substring match against keys and values is fine). Register the tool in both TOOLS (the schema list sent to Claude) and TOOL_FUNCS (the in-process dispatch table). Then run the agent with a question that requires the new tool, for example:

python agent.py "What did I write about my Kyiv trip?"

✓ Self-check: the agent calls search_notes, gets the note body back, and synthesises a one-sentence answer that mentions thermals.

Task 2 — Add a structured log

After every turn of the loop — before you continue to the next iteration — append one JSON line to a file named agent.log. Each line must contain:

{
  "turn": 1,
  "stop_reason": "tool_use",
  "tool_calls": ["get_weather"],
  "input_tokens": 432,
  "output_tokens": 87,
  "cost_so_far_usd": 0.0021
}

Use json.dumps(...) + a newline and open the file in append mode ("a"). Do not pretty-print — one JSON object per line so the file stays grep-friendly and jq-friendly.

✓ Self-check: running one question produces a multi-line agent.log with one JSON object per turn, and the cost_so_far_usd field is monotonically increasing across lines.

Task 3 — Switch to a cheaper model with a fallback

Change the top-level constant MODEL to "claude-3-5-haiku-latest" and update the per-million-token pricing constants accordingly. Use these Haiku 2026 Q2 numbers:

PRICE_IN_PER_M = 0.80 (USD per 1M input tokens)
PRICE_OUT_PER_M = 4.00 (USD per 1M output tokens)

Re-run the same question from Task 1. The cost summary at the end should drop noticeably compared to the Sonnet baseline.

Then add a one-shot fallback: if any turn's response contains a tool_result block with is_error=True, switch MODEL to "claude-3-5-sonnet-latest" for the rest of the loop. This is a one-way upgrade — once you switch up, you stay on Sonnet, you do not flip back and forth per turn.

To exercise the fallback, deliberately break one of your tools (e.g. make get_weather raise RuntimeError("boom") on a specific city, and wrap the dispatch in a try/except that returns a tool_result with is_error=True).

✓ Self-check: the agent uses Haiku by default; a deliberately-broken tool causes the loop to upgrade to Sonnet for the next turn and successfully recover — the final answer still comes back, and your agent.log shows the model swap.

What you have learned

Adding search_notes shows you how cheap it is to give an agent a new capability once the loop is in place — one function, one schema entry, one dispatch row. The structured log is the seed of every observability story you will hear about in Module 7: per-turn JSON is enough to feed a dashboard, a regression test, or a cost alarm. And the Haiku-with-Sonnet-fallback is the production pattern in miniature: run the cheap model by default, escalate only when the cheap model proves it cannot cope.

View solutions →