Skip to content
Multigrid

Type Hints That Make an AI Codebase Survivable

11 min read · updated August 4, 2026

An LLM codebase is unusually full of nested dictionaries that came off the network, and unusually short of compiler assistance because the interesting values are strings. Typing the four shapes that recur — message, tool schema, tool result, response — recovers most of what is missing.

What types buy in this particular codebase

The generic argument for type hints applies, but three of the payoffs are specific to code that talks to models, and they are the ones worth the effort.

  • The response shape is three levels deep. body["choices"][0]["message"]["content"] has four chances to be wrong and no autocompletion. A TypedDict gives every level a name and catches the typo at check time.
  • Message roles are a closed set that is enforced nowhere. {"role": "assistant "} with a trailing space is a 400 from the provider at runtime, or worse, a silently ignored message. Literal makes it a type error.
  • Optionality is real and pervasive. content is None on a tool call; usage is absent on some streamed responses; finish_reason is None on intermediate frames. Encoding those as | None forces the check at the one place it belongs.

Typing the message list

TypedDict is the right tool here rather than a dataclass, because these values are dictionaries — they go straight into json.dumps — and a dataclass would need converting on every call.

# wire.py
from __future__ import annotations

from typing import Any, Literal, NotRequired, TypedDict

Role = Literal["system", "user", "assistant", "tool"]


class FunctionCall(TypedDict):
    name: str
    arguments: str                 # a JSON *string*, not an object


class ToolCall(TypedDict):
    id: str
    type: Literal["function"]
    function: FunctionCall


class Message(TypedDict):
    role: Role
    content: str | None            # None when the assistant returns tool calls
    name: NotRequired[str]
    tool_calls: NotRequired[list[ToolCall]]
    tool_call_id: NotRequired[str]  # required on role="tool" replies


class Usage(TypedDict):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class Choice(TypedDict):
    index: int
    message: Message
    finish_reason: Literal["stop", "length", "tool_calls", "content_filter"] | None


class ChatResponse(TypedDict):
    id: str
    model: str
    choices: list[Choice]
    usage: NotRequired[Usage]


class ChatRequest(TypedDict):
    model: str
    messages: list[Message]
    temperature: NotRequired[float]
    max_tokens: NotRequired[int]
    stream: NotRequired[bool]
    tools: NotRequired[list[dict[str, Any]]]
    response_format: NotRequired[dict[str, Any]]

NotRequired is in typing from Python 3.11; on 3.10 import it from typing_extensions, which is a dependency worth adding for this alone. The alternative in older code — two TypedDicts with total=False on one and inheritance between them — works and is much harder to read.

The single most useful line in that file is arguments: str. Tool-call arguments arrive as a JSON string that you must parse yourself, and treating them as a dict is a TypeError at runtime that the annotation prevents at check time. It is also a reminder that the string may not parse — the model produced it, so it is subject to everything in parsing model output safely.

Constructors keep the call sites tidy and the roles correct:

def user(content: str) -> Message:
    return {"role": "user", "content": content}


def system(content: str) -> Message:
    return {"role": "system", "content": content}


def tool_reply(call_id: str, content: str) -> Message:
    return {"role": "tool", "content": content, "tool_call_id": call_id}

Tool schemas and tool results

A tool has three typed pieces that are easy to let drift: the schema you send, the Python function you dispatch to, and the result you send back. Keeping them in one object is what stops the schema describing a parameter the function no longer takes.

# tools.py
from dataclasses import dataclass
from typing import Any, Callable, Protocol


class ToolHandler(Protocol):
    def __call__(self, **kwargs: Any) -> str: ...


@dataclass(frozen=True)
class Tool:
    name: str
    description: str
    parameters: dict[str, Any]          # a JSON Schema object
    handler: ToolHandler

    def to_wire(self) -> dict[str, Any]:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


REGISTRY: dict[str, Tool] = {}


def register(tool: Tool) -> Tool:
    REGISTRY[tool.name] = tool
    return tool


def dispatch(call: ToolCall) -> Message:
    """Run one tool call and return the message to append to the conversation."""
    tool = REGISTRY.get(call["function"]["name"])
    if tool is None:
        return tool_reply(call["id"],
                          f"error: no such tool {call['function']['name']!r}")
    try:
        arguments = json.loads(call["function"]["arguments"])
    except json.JSONDecodeError as exc:
        return tool_reply(call["id"], f"error: arguments were not valid JSON: {exc}")
    if not isinstance(arguments, dict):
        return tool_reply(call["id"], "error: arguments must be a JSON object")
    try:
        return tool_reply(call["id"], tool.handler(**arguments))
    except TypeError as exc:
        return tool_reply(call["id"], f"error: wrong arguments: {exc}")
    except Exception as exc:                     # a tool failure is data, not a crash
        return tool_reply(call["id"], f"error: {type(exc).__name__}: {exc}")

Every failure path returns a Message, never raises. That is a type-level decision with a behavioural consequence: the model gets to see “that tool does not exist” and try something else, where an exception would end the loop. dispatch having a single return type is what makes the agent loop above it simple.

Note the deliberate bare except Exception on the last handler. Tool code is arbitrary and a tool raising must not kill the conversation; the exception type and message go back to the model as text. Agent error recovery covers how much detail to hand back, and secure tool calls covers why **arguments into a handler needs the schema to be strict.

A result type instead of an exception or None

Functions that call a model have three outcomes — a value, a failure the caller can act on, and a bug. Returning None collapses the first two; raising for everything makes the ordinary case exceptional. A small tagged union keeps them apart and makes the checker enforce that both are handled.

from dataclasses import dataclass
from typing import Generic, Literal, TypeVar

T = TypeVar("T")


@dataclass(frozen=True)
class Ok(Generic[T]):
    value: T
    usage: Usage | None = None


@dataclass(frozen=True)
class Failed:
    kind: Literal["http", "timeout", "parse", "validation", "budget"]
    detail: str
    retryable: bool


Outcome = Ok[T] | Failed


def classify(text: str) -> Outcome[str]:
    ...


match classify(text):
    case Ok(value=label, usage=usage):
        store(label, usage)
    case Failed(kind="budget", detail=detail):
        halt(detail)
    case Failed(retryable=True) as failure:
        requeue(failure)
    case Failed() as failure:
        dead_letter(failure)

kind as a Literal is what makes this pay. Add a sixth failure kind and every match that does not handle it is reported by the checker, which is precisely the change that otherwise ships as a silent fall-through. The pattern-matching syntax needs Python 3.10; on 3.9 the same union works with isinstance checks and the checker narrows just as well.

Protocols at the boundary

Protocol gives you an interface without an inheritance relationship, which is exactly what you want for the model client: the real one, the fake one in tests and the cached one in a notebook all satisfy it without importing each other.

from typing import Protocol, runtime_checkable


@runtime_checkable
class ChatClient(Protocol):
    def chat(self, request: ChatRequest) -> ChatResponse: ...


class HttpChatClient:
    def __init__(self, client: httpx.Client) -> None:
        self._client = client

    def chat(self, request: ChatRequest) -> ChatResponse:
        response = self._client.post("/chat/completions", json=request)
        response.raise_for_status()
        return response.json()          # cast, not a guarantee — see below


class FakeChatClient:
    def __init__(self, replies: list[str]) -> None:
        self._replies = list(replies)

    def chat(self, request: ChatRequest) -> ChatResponse:
        text = self._replies.pop(0)
        return {"id": "fake", "model": request["model"],
                "choices": [{"index": 0, "finish_reason": "stop",
                             "message": {"role": "assistant", "content": text}}]}

Neither class names the protocol and both satisfy it. Any function annotated client: ChatClient accepts both, which is the seam that testing code that calls a model is built on.

Turning the checker on without a rewrite

Strict mode on an existing codebase produces thousands of errors and gets switched off again. Turn it on for new modules only, then widen.

# pyproject.toml
[tool.mypy]
python_version = "3.11"
warn_unused_ignores = true
warn_redundant_casts = true
warn_return_any = true
# start permissive across the codebase
disallow_untyped_defs = false
ignore_missing_imports = true

# and strict where the wire format lives, because that is where the value is
[[tool.mypy.overrides]]
module = ["mylib.wire", "mylib.tools", "mylib.parsing"]
disallow_untyped_defs = true
strict_optional = true

Add mypy mylib to CI on the first day, even with a permissive configuration, so the count only moves in one direction. Pyright, which is what most editors run for you, is faster and stricter about narrowing; the two disagree at the edges, so pick the one CI enforces and treat the other as advice.

What a type checker cannot tell you

Being clear about this is what keeps the types honest, because the failure mode of a well-typed LLM codebase is believing the annotations about data that came off a socket.

  • response.json() is Any. Annotating the return as ChatResponse is an assertion, not a check — nothing verifies it at runtime. If a provider changes a field, the checker stays green and the code fails at the subscript. The live contract tests in testing code that calls a model exist to cover exactly this gap, and Pydantic covers it for the payloads where the cost of being wrong is high.
  • A str is not a valid prompt. No type distinguishes a 200-token prompt from a 200,000-token one, or trusted text from user-supplied text. NewType can mark the distinction if you propagate it manually, but the checker will not find the place you forgot.
  • Literal["billing", ...] does not constrain the model. It constrains your code. The model can return anything; the type is a claim about what you accept, and something has to validate it at the boundary.
  • Types do not survive **kwargs. The tool dispatch above is deliberately dynamic and the checker can say nothing about whether the arguments match the handler. That is what the TypeError handler is for, and it is why the JSON Schema should be strict.

Related