Solutions — Extending the Hand-Rolled Agent

Lesson 02 · Worked solutions · model code + the common mistake for each task

Below are the reference implementations for the three homework tasks, each shown as the minimum diff against the Lab 02 agent.py. Read them after you have written your own version — the value is in comparing your choices to these, not in copy-pasting. Each solution closes with a Common mistake callout that flags the trap most people fall into the first time.

Solution to Task 1 — Add the search_notes tool

Three additions: the in-process notes dict, the tool function itself, and one new entry in both the dispatch table (TOOL_FUNCS) and the schema list (TOOLS) the API actually sees.

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

def search_notes(query: str) -> str:
    q = query.lower()
    hits = [f"{k}: {v}" for k, v in NOTES.items() if q in k.lower() or q in v.lower()]
    return "\n".join(hits) if hits else "no notes match"

TOOL_FUNCS["search_notes"] = search_notes
TOOLS.append({
    "name": "search_notes",
    "description": "Search the user's notes by substring; returns matching notes.",
    "input_schema": {
        "type": "object",
        "properties": {"query": {"type": "string"}},
        "required": ["query"],
    },
})

Solution to Task 2 — Append a JSON line per turn

Two changes inside run_agent: truncate the log file once at the top, then append one JSON object per line as the very last thing in each loop iteration.

import json

# at the top of run_agent, before the loop:
log_path = "agent.log"
with open(log_path, "w"):  # truncate
    pass

# inside the loop, just after computing `cost`:
log_entry = {
    "turn": turn,
    "stop_reason": resp.stop_reason,
    "tool_calls": [b.name for b in resp.content if b.type == "tool_use"],
    "input_tokens": resp.usage.input_tokens,
    "output_tokens": resp.usage.output_tokens,
    "cost_so_far_usd": round(cost, 6),
}
with open(log_path, "a") as f:
    f.write(json.dumps(log_entry) + "\n")

Solution to Task 3 — Haiku by default, Sonnet as one-shot fallback

Replace the single MODEL constant with a primary/fallback pair and a pricing table keyed by model. Track current_model as loop state. After appending tool results, check whether any of them is an error — if so, upgrade to Sonnet and stay there.

MODEL_PRIMARY = "claude-3-5-haiku-latest"
MODEL_FALLBACK = "claude-3-5-sonnet-latest"
PRICING = {
    MODEL_PRIMARY: (0.80, 4.0),
    MODEL_FALLBACK: (3.0, 15.0),
}

def run_agent(user_question: str) -> tuple[str, float]:
    client = anthropic.Anthropic()
    messages = [{"role": "user", "content": user_question}]
    total_in = total_out = 0.0
    current_model = MODEL_PRIMARY

    for turn in range(MAX_TURNS):
        resp = client.messages.create(
            model=current_model,
            max_tokens=1024,
            system=SYSTEM,
            tools=TOOLS,
            messages=messages,
        )
        in_rate, out_rate = PRICING[current_model]
        total_in += resp.usage.input_tokens
        total_out += resp.usage.output_tokens
        # cost tracking unchanged ...

        # ... existing branch on stop_reason ...

        # after appending tool_results, check for tool errors:
        if any(tr.get("is_error") for tr in tool_results):
            current_model = MODEL_FALLBACK
        messages.append({"role": "user", "content": tool_results})

    raise RuntimeError(f"turn cap ({MAX_TURNS}) reached without end_turn")

What you have built

Together, these three tasks cover the three axes every production agent grows along:

• Capability — one new function plus one schema entry expands what the agent can do, with zero changes to the loop.
• Observability — per-turn JSONL is the smallest thing that’s actually useful for cost dashboards, replays, and regression tests.
• Cost — routing the cheap model by default and upgrading only on proven failure is the canonical production pattern, in miniature.