Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Mini-project 3: build your own @tool decorator framework

You now know what a good tool looks like and how function calling works on the wire. This chapter closes the gap between those two by making you write the thing that sits in between.

By the end you will have a small framework — around two hundred lines — that lets you write this:

@tool(read_only=True, timeout_s=2.0)
def check_stock(sku: str, warehouse: Literal["AMS", "SIN", "any"] = "any") -> dict:
    """Report how many units of one SKU are available to ship today.

    Args:
        sku: The exact SKU identifier, e.g. "SKU-1001". Get it from find_sku.
        warehouse: Which warehouse to check. "any" checks all of them.
    """

…and get a validated JSON Schema, a registry entry, timeout enforcement, structured errors, and result truncation, for free.

This is not a toy. It is a stripped-down version of exactly what LangChain, Google ADK, and the Agents SDK do internally. Writing it once means you will never again wonder what a framework is doing to your function on the way to the model.

Everything runs offline against a scripted mock model. No API key required.

Setup

mkdir -p toolkit-project && cd toolkit-project
pip install jsonschema

We build two files: toolkit.py (the framework) and demo.py (tools, mock model, loop).

Step 1: the problem — schemas drift from code

The naive approach is to write the schema by hand next to the function.

CHECK_STOCK_SCHEMA = {"type": "object", "properties": {"sku": {"type": "string"}}, ...}

def check_stock(sku, warehouse="any"):
    ...

Two representations of the same signature. Add a parameter to the function and the schema silently lies to the model. Rename one and the model calls a field your function does not accept. This drift is a real production bug and it is entirely avoidable, because Python already knows the signature.

So: derive the schema from the function. The type hints give us types, the defaults give us required, and the docstring gives us the descriptions.

Step 2: parsing the docstring

Start toolkit.py. The docstring is where the prompt lives, so we need the summary and the per-parameter descriptions out of it. We support Google style, which is what ADK uses and what most people write anyway.

from __future__ import annotations
import inspect, re

_ARGS_HEADER = re.compile(r"^\s*(Args|Arguments|Parameters)\s*:\s*$", re.M)
_SECTION = re.compile(r"^\s*(Returns|Raises|Yields|Examples?|Notes?)\s*:\s*$", re.M)
_PARAM_LINE = re.compile(r"^\s*(\*{0,2}\w+)\s*(?:\([^)]*\))?\s*:\s*(.*)$")


def parse_docstring(doc: str) -> tuple[str, dict[str, str]]:
    """Split a Google-style docstring into a summary and a param -> description map."""
    doc = inspect.cleandoc(doc or "")
    m = _ARGS_HEADER.search(doc)
    if not m:
        return doc.strip(), {}
    summary = doc[: m.start()].strip()
    rest = doc[m.end():]
    end = _SECTION.search(rest)
    body = rest[: end.start()] if end else rest

    params, current = {}, None
    for line in body.splitlines():
        if not line.strip():
            continue
        pm = _PARAM_LINE.match(line)
        if pm and not line.startswith("        "):
            current = pm.group(1).lstrip("*")
            params[current] = pm.group(2).strip()
        elif current:                      # continuation of the previous param
            params[current] += " " + line.strip()
    return summary, params

Note what the summary includes: everything before Args:, not just the first line. That is deliberate. The “use this when…” paragraph from Chapter 1 is the most valuable sentence in the docstring and it must reach the model.

The indentation check on continuation lines is a small thing that matters — it lets a parameter description wrap across lines without the wrapped part being mistaken for a new parameter.

Step 3: types to JSON Schema

Next, turn annotations into schema fragments. Cover the cases you will actually hit and return {} (meaning “any”) for the rest, rather than crashing.

import typing

_PRIMITIVES = {
    str: {"type": "string"},
    int: {"type": "integer"},
    float: {"type": "number"},
    bool: {"type": "boolean"},
    type(None): {"type": "null"},
}


def schema_for_type(tp) -> dict:
    if tp is inspect.Parameter.empty or tp is typing.Any:
        return {}
    if tp in _PRIMITIVES:
        return dict(_PRIMITIVES[tp])
    origin, args = typing.get_origin(tp), typing.get_args(tp)
    if origin is typing.Literal:
        return {"enum": list(args)}                      # Literal -> enum
    if origin in (list, set, tuple):
        return {"type": "array", "items": schema_for_type(args[0]) if args else {}}
    if origin is dict:
        return {"type": "object"}
    if origin is typing.Union or str(origin) == "<class 'types.UnionType'>":
        non_null = [a for a in args if a is not type(None)]
        sub = [schema_for_type(a) for a in non_null]
        out = sub[0] if len(sub) == 1 else {"anyOf": sub}
        if type(None) in args:
            out = {"anyOf": [out, {"type": "null"}]}
        return out
    return {}

The Literal case is the payoff. Literal["AMS", "SIN", "any"] becomes {"enum": ["AMS", "SIN", "any"]}, which means the model is told the legal values and the validator enforces them — Chapter 1’s two jobs of a schema, from one annotation.

Now assemble the full input schema.

def build_schema(fn) -> dict:
    summary, docs = parse_docstring(fn.__doc__)
    hints = typing.get_type_hints(fn)
    props, required = {}, []
    for name, p in inspect.signature(fn).parameters.items():
        if name in ("self", "cls") or p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD):
            continue
        s = schema_for_type(hints.get(name, p.annotation))
        if desc := docs.get(name):
            s["description"] = desc
        if p.default is not inspect.Parameter.empty:
            s["default"] = p.default
        else:
            required.append(name)
        props[name] = s
    return {
        "summary": summary,
        "input_schema": {
            "type": "object",
            "properties": props,
            "required": required,
            "additionalProperties": False,
        },
    }

typing.get_type_hints rather than reading p.annotation directly, because it resolves string annotations — which is what you get under from __future__ import annotations. additionalProperties: False for the reason from Chapter 2. A parameter with a default is not required; a parameter without one is.

Here is what that produces for a real function:

{
  "summary": "Find product SKUs whose name matches a shopper's words.",
  "input_schema": {
    "type": "object",
    "properties": {
      "query": {"type": "string", "description": "Words from the product name."},
      "limit": {"type": "integer", "description": "Max matches. Defaults to 3.", "default": 3}
    },
    "required": ["query"],
    "additionalProperties": false
  }
}

One function, one docstring, zero duplication.

Step 4: the decorator and the registry

Now the ergonomics. A decorator that builds the schema at import time and files the result in a registry.

import jsonschema


class ToolError(Exception):
    """Raise this inside a tool to return a clean, model-readable failure."""
    def __init__(self, message: str, code: str = "tool_error", retriable: bool = False):
        super().__init__(message)
        self.code, self.retriable = code, retriable


class Tool:
    def __init__(self, fn, name, description, input_schema, timeout_s, read_only):
        self.fn, self.name, self.description = fn, name, description
        self.input_schema, self.timeout_s, self.read_only = input_schema, timeout_s, read_only
        self.validator = jsonschema.Draft202012Validator(input_schema)

    def spec(self) -> dict:
        return {"name": self.name, "description": self.description,
                "input_schema": self.input_schema}


class Registry:
    def __init__(self):
        self._tools: dict[str, Tool] = {}

    def add(self, t: Tool):
        if t.name in self._tools:
            raise ValueError(f"duplicate tool name: {t.name}")
        self._tools[t.name] = t

    def get(self, name): return self._tools.get(name)
    def specs(self):     return [t.spec() for t in self._tools.values()]
    def names(self):     return sorted(self._tools)


REGISTRY = Registry()


def tool(_fn=None, *, name=None, timeout_s=10.0, read_only=False, registry=None):
    """Register a function as a model-callable tool."""
    def wrap(fn):
        built = build_schema(fn)
        if not built["summary"]:
            raise ValueError(
                f"{fn.__name__} needs a docstring; it is the model's only view of this tool")
        t = Tool(fn=fn, name=name or fn.__name__, description=built["summary"],
                 input_schema=built["input_schema"], timeout_s=timeout_s, read_only=read_only)
        (registry or REGISTRY).add(t)
        fn.tool = t
        return fn
    return wrap(_fn) if _fn else wrap

Three deliberate design choices.

A missing docstring is a hard error at import time. Not a warning. A tool with no description is a tool the model cannot use correctly, and failing at import means you find out in seconds rather than in a confusing production trace.

Duplicate names are a hard error too. Two tools with the same name means one silently shadows the other, which is a genuinely miserable bug to track down.

The decorator returns the original function. check_stock("SKU-1001") still works normally in your own code and in your unit tests. Only the metadata is added, hung off fn.tool.

The registry= parameter lets you build isolated registries — useful for tests and for agents that should see different tool sets.

Step 5: dispatch, with everything that can go wrong

This is the heart of the framework. It turns “the model asked for something” into “here is text, and whether it failed” — the uniform contract from Chapter 2.

import concurrent.futures, json, time

_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=8)


def _err(code, message, **extra):
    return {"ok": False, "error": {"code": code, "message": message, **extra}}


def dispatch(call: dict, registry: Registry | None = None) -> dict:
    reg = registry or REGISTRY
    name, args = call.get("name"), call.get("input") or {}

    t = reg.get(name)
    if t is None:
        return _err("unknown_tool",
                    f"No tool named {name!r}. Available tools: {', '.join(reg.names())}. "
                    f"Call one of those instead.")

    errors = sorted(t.validator.iter_errors(args), key=lambda e: list(e.path))
    if errors:
        detail = "; ".join(
            f"{'.'.join(str(p) for p in e.path) or '<root>'}: {e.message}" for e in errors[:5])
        return _err("invalid_arguments",
                    f"Arguments rejected for {name}: {detail}. "
                    f"Re-read the schema and call the tool again with corrected arguments.",
                    schema=t.input_schema)

    started = time.monotonic()
    fut = _POOL.submit(t.fn, **args)
    try:
        value = fut.result(timeout=t.timeout_s)
    except concurrent.futures.TimeoutError:
        return _err("timeout",
                    f"{name} did not finish within {t.timeout_s:g}s. "
                    f"Do not retry immediately; tell the user the backend is slow.")
    except ToolError as e:
        return _err(e.code, str(e), retriable=e.retriable)
    except Exception as e:
        return _err("internal_error",
                    f"{name} failed: {type(e).__name__}: {e}. "
                    f"This is a bug in the tool, not in your arguments; do not retry.")
    return {"ok": True, "value": value,
            "elapsed_ms": round((time.monotonic() - started) * 1000, 1)}

Walk the failure ladder, because each rung exists for a reason from Chapter 2.

Unknown tool returns the list of real names. The model corrects itself on the next turn.

Validation uses iter_errors, not validate, so the model gets all the problems at once instead of fixing them one per round trip. We cap at five so a badly-wrong call cannot produce a wall of text. We include the schema in the error payload, which costs tokens but converts most second attempts into successes.

Timeout via ThreadPoolExecutor. Be honest about the limitation: this bounds how long the agent waits, not how long the function runs. The thread keeps going. For real work, pass a timeout down to your HTTP client or database driver as well — this is the outer guard, not the only one.

ToolError is your clean channel for expected business failures — not found, permission denied, rate limited — with the recovery instruction attached.

Everything else is a bug in your code, and the message says so and tells the model not to retry. That single sentence prevents a large fraction of runaway loops.

Step 6: bounding the output

One more piece, and it belongs in the framework rather than in each tool, because it must apply to tools you did not write.

MAX_RESULT_CHARS = 2000


def render_result(result: dict) -> str:
    text = json.dumps(result, default=str, ensure_ascii=False)
    if len(text) <= MAX_RESULT_CHARS:
        return text
    keep = MAX_RESULT_CHARS - 200
    return (text[:keep] + f"\n...[truncated {len(text) - keep} chars]. "
            "Narrow your arguments (add filters or a smaller limit) and call again.")

default=str keeps a stray datetime from crashing the serializer at the worst possible moment. The truncation notice is actionable, per Chapter 1 — it tells the model what to do rather than just that something was lost.

That is toolkit.py. About 180 lines.

Step 7: real tools

Now demo.py. Three tools, written to the standard of Chapter 1.

from __future__ import annotations
import json, time, typing
from toolkit import REGISTRY, ToolError, dispatch, render_result, tool

INVENTORY = {
    "SKU-1001": {"name": "Astro Zoom Trainers", "warehouse": "AMS", "on_hand": 42},
    "SKU-1002": {"name": "Cirrus Rain Shell",   "warehouse": "AMS", "on_hand": 0},
    "SKU-2001": {"name": "Tidal Swim Goggles",  "warehouse": "SIN", "on_hand": 7},
}


@tool(read_only=True)
def find_sku(query: str, limit: int = 3) -> list[dict]:
    """Find product SKUs whose name matches a shopper's words.

    Use this first when you only have a product name and need a SKU.

    Args:
        query: Words from the product name, e.g. "rain shell".
        limit: Maximum number of matches to return. Defaults to 3.

    Returns:
        A list of {sku, name} objects, most relevant first. Empty if nothing matched.
    """
    q = query.lower()
    hits = [{"sku": k, "name": v["name"]} for k, v in INVENTORY.items() if q in v["name"].lower()]
    return hits[:limit]


@tool(read_only=True, timeout_s=2.0)
def check_stock(sku: str, warehouse: typing.Literal["AMS", "SIN", "any"] = "any") -> dict:
    """Report how many units of one SKU are available to ship today.

    Args:
        sku: The exact SKU identifier, e.g. "SKU-1001". Get it from find_sku.
        warehouse: Which warehouse to check. "any" checks all of them.

    Returns:
        {sku, on_hand, warehouse, shippable}. shippable is false when on_hand is 0.
    """
    row = INVENTORY.get(sku)
    if row is None:
        raise ToolError(
            f"No product with SKU {sku!r}. Ask the shopper for the product name and "
            "call find_sku to get a valid SKU, then try again.",
            code="not_found", retriable=True)
    if warehouse != "any" and row["warehouse"] != warehouse:
        return {"sku": sku, "on_hand": 0, "warehouse": warehouse, "shippable": False}
    return {"sku": sku, "on_hand": row["on_hand"], "warehouse": row["warehouse"],
            "shippable": row["on_hand"] > 0}


@tool(timeout_s=0.5)
def slow_report(rows: int = 100000) -> dict:
    """Generate an inventory report. Deliberately slow, to demonstrate timeouts.

    Args:
        rows: How many rows to generate.
    """
    time.sleep(3)
    return {"rows": rows}

Step 8: a mock model and the loop

To run the loop without an API key, replace the model with something that replays a script. This is not just a convenience — a deterministic fake model is how you write fast unit tests for agent behavior.

class MockModel:
    """Replays a scripted sequence of turns so the loop runs with no network."""
    def __init__(self, script):
        self.script = list(script)

    def __call__(self, messages, tools):
        return self.script.pop(0) if self.script else {"type": "text", "text": "(no more script)"}


def run(model, user_msg, max_steps=6, verbose=True):
    messages = [{"role": "user", "content": user_msg}]
    specs = REGISTRY.specs()
    for step in range(max_steps):
        turn = model(messages, specs)
        if turn["type"] == "text":
            messages.append({"role": "assistant", "content": turn["text"]})
            return turn["text"], messages
        calls = turn["calls"]
        messages.append({"role": "assistant", "content": calls})
        results = []
        for c in calls:
            r = dispatch(c)
            if verbose:
                print(f"  step {step}  {c['name']}({json.dumps(c['input'])}) "
                      f"-> {render_result(r)[:160]}")
            results.append({"tool_use_id": c["id"], "content": render_result(r)})
        messages.append({"role": "user", "content": results})
    return "(step limit reached)", messages

The shape is identical to the real loop from Chapter 2 — same termination condition, same turn cap, same one-message-of-results structure. To go live, swap MockModel for a function that calls client.messages.create(...) with tools=specs and translates the response into the same {"type": ..., "calls": [...]} shape. Nothing else changes. That is the point of keeping the model behind a callable.

Step 9: run it

if __name__ == "__main__":
    print("== registered tools ==")
    for s in REGISTRY.specs():
        print(f"- {s['name']}: {s['description'].splitlines()[0]}")

    print("\n== happy path (parallel calls in one turn) ==")
    script = [
        {"type": "tool_use", "calls": [
            {"id": "a1", "name": "find_sku", "input": {"query": "rain shell"}},
            {"id": "a2", "name": "find_sku", "input": {"query": "goggles"}}]},
        {"type": "tool_use", "calls": [
            {"id": "b1", "name": "check_stock", "input": {"sku": "SKU-1002"}},
            {"id": "b2", "name": "check_stock", "input": {"sku": "SKU-2001", "warehouse": "SIN"}}]},
        {"type": "text",
         "text": "The Cirrus Rain Shell is out of stock; 7 Tidal Swim Goggles ship from SIN."},
    ]
    answer, _ = run(MockModel(script), "Do you have the rain shell and the goggles?")
    print("  final:", answer)

    print("\n== failure modes ==")
    for c in [
        {"id": "c1", "name": "chekc_stock", "input": {"sku": "SKU-1001"}},
        {"id": "c2", "name": "check_stock", "input": {"sku": 1001}},
        {"id": "c3", "name": "check_stock", "input": {"sku": "SKU-1001", "warehouse": "LHR"}},
        {"id": "c4", "name": "check_stock", "input": {}},
        {"id": "c5", "name": "check_stock", "input": {"sku": "SKU-9999"}},
        {"id": "c6", "name": "slow_report", "input": {}},
    ]:
        r = dispatch(c)
        print(f"- {c['name']}{json.dumps(c['input'])}\n"
              f"    {r['error']['code']}: {r['error']['message'][:130]}")

python demo.py

Real output:

== registered tools ==
- find_sku: Find product SKUs whose name matches a shopper's words.
- check_stock: Report how many units of one SKU are available to ship today.
- slow_report: Generate an inventory report. Deliberately slow, to demonstrate timeouts.

== schema for check_stock ==
{
  "type": "object",
  "properties": {
    "sku": {
      "type": "string",
      "description": "The exact SKU identifier, e.g. \"SKU-1001\". Get it from find_sku."
    },
    "warehouse": {
      "enum": ["AMS", "SIN", "any"],
      "description": "Which warehouse to check. \"any\" checks all of them.",
      "default": "any"
    }
  },
  "required": ["sku"],
  "additionalProperties": false
}

== happy path (parallel calls in one turn) ==
  step 0  find_sku({"query": "rain shell"}) -> {"ok": true, "value": [{"sku": "SKU-1002", "name": "Cirrus Rain Shell"}], "elapsed_ms": 1.2}
  step 0  find_sku({"query": "goggles"}) -> {"ok": true, "value": [{"sku": "SKU-2001", "name": "Tidal Swim Goggles"}], "elapsed_ms": 0.1}
  step 1  check_stock({"sku": "SKU-1002"}) -> {"ok": true, "value": {"sku": "SKU-1002", "on_hand": 0, "warehouse": "AMS", "shippable": false}, "elapsed_ms": 0.1}
  step 1  check_stock({"sku": "SKU-2001", "warehouse": "SIN"}) -> {"ok": true, "value": {"sku": "SKU-2001", "on_hand": 7, "warehouse": "SIN", "shippable": true}, "elapsed_ms": 0.1}
  final: The Cirrus Rain Shell is out of stock; 7 Tidal Swim Goggles ship from SIN.

== failure modes ==
- chekc_stock{"sku": "SKU-1001"}
    unknown_tool: No tool named 'chekc_stock'. Available tools: check_stock, find_sku, slow_report. Call one of those instead.
- check_stock{"sku": 1001}
    invalid_arguments: Arguments rejected for check_stock: sku: 1001 is not of type 'string'. Re-read the schema and call the tool again with corrected a
- check_stock{"sku": "SKU-1001", "warehouse": "LHR"}
    invalid_arguments: Arguments rejected for check_stock: warehouse: 'LHR' is not one of ['AMS', 'SIN', 'any']. Re-read the schema and call the tool aga
- check_stock{}
    invalid_arguments: Arguments rejected for check_stock: <root>: 'sku' is a required property. Re-read the schema and call the tool again with correcte
- check_stock{"sku": "SKU-9999"}
    not_found: No product with SKU 'SKU-9999'. Ask the shopper for the product name and call find_sku to get a valid SKU, then try again.
- slow_report{}
    timeout: slow_report did not finish within 0.5s. Do not retry immediately; tell the user the backend is slow.

Read the failure block once more, slowly. Every single message names the problem and prescribes the next action. That is not decoration — it is the difference between a model that recovers on turn two and a model that loops until your turn cap fires.

Note also that Literal["AMS", "SIN", "any"] produced both the enum in the schema and the enforcement that rejected "LHR". One annotation, both jobs.

Extensions worth doing

Try these before moving on; each is under thirty lines and each teaches something.

Repeat detection. Hash (name, sorted(args)) per conversation. On the third identical call, return an error telling the model to stop and explain the problem to the user.

Idempotency. When read_only=False, derive a key from the call hash, cache the result, and return the cached value with "deduplicated": true on a repeat. This is the Chapter 2 defense, implemented once for every tool.

A confirmation gate. Add requires_approval=True to the decorator. Have dispatch return {"ok": False, "error": {"code": "approval_required", ...}} unless an approval token is present in the call. This is how you keep an agent from wiring money at 3 a.m.

Output schemas. Add output_schema= and validate the return value. You will catch backend drift before the model turns it into a confident wrong answer.

Tracing. Log every dispatch as one structured line: tool name, argument hash, outcome code, elapsed milliseconds. When you get to observability in a later part, this is the data you will wish you had been collecting.

Async. Swap the thread pool for asyncio.wait_for and support async def tools. Most real tools are I/O-bound, so this is the version you will actually ship.

What you should be able to do now

  • Generate a correct JSON Schema automatically from a Python function’s signature, type hints, and Google-style docstring, with no hand-maintained duplicate.
  • Explain why Literal[...] is the highest-value annotation in a tool signature, and demonstrate it doing both documentation and enforcement.
  • Write a dispatcher that never raises into the agent loop and returns a uniform structured result for unknown tools, invalid arguments, business failures, timeouts, and internal bugs.
  • Cap and truncate tool output at the framework level so a single misbehaving tool cannot poison the context window.
  • Test an agent loop end to end with a scripted mock model, with no API key and no network.

Further reading