Solutions — Three Extensions of the Research Assistant

Lesson 04 · Building Your Own Agents · worked answers with commentary

Below is one clean answer per homework task, plus the why-this-not-that commentary that makes the difference between a passing solution and an idiomatic one. Compare your attempt against these before you read — even a different-but-working solution is fine, but the commentary often catches subtle traps you may not have noticed.

How to read the solutions

Each section opens with the code, then a common mistake callout, then the broader lesson. The code is meant to drop into the lab’s agent.py with minimal edits — copy-paste, restart the process, and you should see the new behaviour on the next agent.run(...).

Solution to Task 1 — the `save_note` tool

Add a sixth @tool alongside the existing five, and append one line to SYSTEM_PROMPT telling the agent when to call it.

@tool
def save_note(note_id: str, title: str, text: str) -> str:
    """Save a note to the personal knowledge base for future searches.
    Call this AFTER write_report when you've learned something worth keeping.

    Args:
        note_id: short kebab-case id, unique. E.g., 'k6.14-bigtcp-ipv4'.
        title: human-readable title.
        text: 1-3 sentence summary you want pkb_search to find later.
    """
    snippet = text[:120].rstrip() + ("..." if len(text) > 120 else "")
    PKB[note_id] = {"title": title, "snippet": snippet, "text": text}
    return f"saved: {note_id}"

Register it in the agent’s tools=[...] list (so the model can see its schema) and append the following bullet to the system prompt:

4. AFTER write_report, optionally call `save_note(note_id, title, text)` with
   any new fact worth remembering for next time.

Common mistake — saving too early

The natural place to put save_note in the system prompt is before write_report — “save what you learn, then write the report”. Don’t. The model then saves nothing until it’s also ready to answer, which it almost always is at the end anyway — and the few times it isn’t, you save half-baked notes. The post-final pattern (save after the report is in hand) is more reliable and the saved notes are higher quality.

Run the agent twice in the same Python process and watch result.usage.total_cost_usd drop. The second run should hit pkb_search on the notes that the first run saved, skipping web_search entirely or invoking it only to fill gaps.

Solution to Task 2 — the `evaluate_report` function

A tiny pure-Python helper, no SDK needed. Drop it next to the agent (or in a separate eval.py) and call it on every agent.run(...) output.

def evaluate_report(out: dict) -> dict:
    report = out.get("report", {})
    bullets = report.get("bullets", [])
    sources = report.get("sources", [])
    intro = report.get("intro", "")
    return {
        "has_intro": len(intro) >= 20,
        "bullet_count": len(bullets),
        "source_count": len(sources),
        "cost_per_bullet_usd": (
            round(out["cost_usd"] / len(bullets), 4) if bullets else None
        ),
    }


if __name__ == "__main__":
    for q in [
        "What changed in Linux kernel 6.14's networking stack?",
        "Summarise the 6.14 filesystem changes.",
        "What are 6.14's scheduler improvements?",
    ]:
        out = run_research_agent(q)
        print(q)
        print(evaluate_report(out))
        print()

Adapt the keys to whatever shape your agent actually returns — the structure above matches the lab’s run_research_agent wrapper ({"report": {...}, "cost_usd": ...}). If you went straight to agent.run() in your driver, pull report = result.final_output and cost = result.usage.total_cost_usd the same way the homework prompt suggests.

Common mistake — wrong denominator

A tempting variant is to divide by “the cost of just the last call” (the call that produced the report) instead of the total task cost. Don’t. The whole point of the metric is to capture amortised cost across all the model’s exploration — the web_search dead-ends, the failed pkb_fetch ids, the retries. That’s what reaches Module 6’s scorecard, and it’s what genuinely differentiates frameworks.

The deeper lesson: per-output cost is the metric that survives framework swaps. Token counts vary, model choice varies, tool latency varies — but cost-per-useful-bullet is comparable across LangGraph, OpenAI SDK, and CrewAI implementations of the same agent.

Solution to Task 3 — grounding check inside `write_report`

Two pieces: a debug helper (is_grounded) and the modified tool. Plus one module-level dict that stashes the current question where the tool can read it.

def is_grounded(bullets: list[str], sources: list[str], question: str) -> list[bool]:
    haystack = (question + " " + " ".join(sources)).lower()
    return [any(word in haystack for word in b.lower().split() if len(word) > 4)
            for b in bullets]


@tool
def write_report(intro: str, bullets: list[str], sources: list[str]) -> str:
    """Save the final structured report. Bullets MUST be supported by sources."""
    # We use a thread-local stash for the current question; in a real app this
    # would be wired through a Context object. For the lesson we read it from
    # a module-level set by run_research_agent before agent.run().
    question = CURRENT_QUESTION.get("text", "")
    grounded = is_grounded(bullets, sources, question)
    if not all(grounded):
        bad = [b for b, g in zip(bullets, grounded) if not g]
        raise ValueError(f"bullet not backed by a source: {bad}")
    LAST_REPORT.clear()
    LAST_REPORT.update({"intro": intro, "bullets": bullets, "sources": sources})
    return "saved"


CURRENT_QUESTION: dict[str, str] = {}

# In run_research_agent, before agent.run():
CURRENT_QUESTION["text"] = question

The matching rule uses a word-level any-overlap heuristic (any 5+ character word from a bullet must appear somewhere in the sources or the question). That’s laxer than the strict full-substring match the homework prompt sketched — in practice the strict version fires too often on natural rephrasing (the model says “networking performance”; the source says “network throughput”). The word-overlap version is the smallest change that makes the check actually useful instead of constantly tripping. Tune to your taste.

Common mistake — wrong exception type

A natural reach is to raise something exotic like AssertionError, RuntimeError, or even BaseException. Stick to ValueError or a project-specific subclass of Exception. The @tool decorator catches Exception subclasses and routes them to the model as is_error: true; BaseException (and friends like KeyboardInterrupt) bypass that catch and crash the whole run.

The deeper lesson here: tools can enforce invariants. The LLM will see the error, think, and try again — usually by rephrasing bullets to use words that actually appear in its sources. This is one of the most powerful patterns in agentic systems: declarative validation at the tool boundary beats post-hoc cleanup downstream. Every tool you write is an opportunity to push a quality check from a fragile parsing layer out into the tool schema, where the model has to satisfy it before the conversation moves on.

What you’ve built

Across three small extensions you’ve added memory (save_note), measurability (evaluate_report), and quality enforcement (grounded write_report) to the canonical Claude Agent SDK research assistant. That’s the template for production tool-heavy agents in three pieces:

Memory — a writable knowledge store the agent can populate as it learns, so successive runs are cheaper and more informed.
Measurability — a deterministic eval function that produces a comparable per-task scorecard, independent of model or framework choice.
Invariants at the tool boundary — tools that raise on bad inputs so the LLM is forced to fix its own outputs before the user sees them.

Closing tip: You’ve now built the canonical Claude Agent SDK use case. Module 3 introduces LangGraph, where you’ll rebuild the same shape of agent as a drawn graph. The tools will look identical; the assembly will look completely different.

← Back to homework