Lesson 07 Sampling & Elicitation
Theme:

Homework Solutions — Sampling & Elicitation

Lesson 07 · MCP Servers · reference implementations + common mistakes

Three reference solutions follow. Each is the smallest correct implementation — production code would add structured logging (Lesson 10) and metric counters, but the shape is what matters for the exercise.

Solution to Task 1 — summarize_all

Goes in pkb/server.py alongside the existing summarize tool. It reuses the lab's summarize_via_sampling helper so the cap check stays in one place: 

@mcp.tool()
async def summarize_all(ctx: Context) -> dict[str, str]:
    """Summarize every note. Refuses if > 5 notes to bound cost."""
    if len(NOTES) > 5:
        raise McpError(ErrorData(
            code=INVALID_PARAMS,
            message=f"summarize_all refuses for {len(NOTES)} notes (cap 5)",
        ))
    return {nid: await summarize_via_sampling(nid, ctx) for nid in NOTES}

Common mistake

Awaiting the sampling calls in parallel via asyncio.gather(...) without a semaphore. The client may rate-limit, and runaway parallel sampling is exactly what the cap was meant to prevent. Sequential await in a dict comprehension is the boring-but-correct choice here.

Solution to Task 2 — elicitation in write_note

Note the optional title kwarg with a default of None: FastMCP honours it in the tool schema (so the LLM can pass a title if it knows one), but the tool gracefully falls back to elicitation when it doesn't. 

@mcp.tool()
async def write_note(body: str, ctx: Context, title: str | None = None) -> str:
    """Create a note. Asks the user for a title if none provided."""
    if title is None:
        result = await ctx.session.elicit(
            message="Pick a title for the new note:",
            requested_schema={
                "type": "object",
                "properties": {"title": {"type": "string"}},
                "required": ["title"],
            },
        )
        if result.action != "accept":
            return "write cancelled by user"
        title = result.content["title"]
    NOTES[title] = body
    return f"saved: {title}"

Common mistake

Assuming result.action == "accept" and indexing result.content unconditionally. Declined or cancelled responses have empty content and crash on the dict access. Always branch on result.action before reading result.content.

Solution to Task 3 — cost-aware short-circuit

This edit is to summarize_via_sampling in pkb/tools.py. Order matters: the unknown-id guard stays first so an invalid request still short-circuits to a structured error, and only then do we check length: 

async def summarize_via_sampling(note_id: str, ctx: Context) -> str:
    if note_id not in NOTES:
        raise McpError(ErrorData(
            code=INVALID_PARAMS,
            message=f"unknown note id: {note_id}",
        ))
    body = NOTES[note_id]
    if len(body) < 80:
        return f"{body}\n\n(Note too short to summarize.)"
    # ... rest unchanged: ctx.session.create_message(...) etc.

Common mistake

Picking the threshold based on words instead of characters. The LLM cost is per-token, which correlates with character count much more reliably than word count for short snippets — a 10-word note can be 30 characters or 90 depending on vocabulary.

What's next

Lesson 08 builds out the full multi-tool flow: write + a composed summarize_and_save that chains sampling and elicitation in a single tool call. Lesson 09 adds an OAuth-gated export tool, so the same patterns scale to authenticated workflows.