Building a Documentation Generator with Pydantic AI: Complete Guide

Building a Documentation Generator with Pydantic AI: Complete Guide

Documentation is one of those tasks every developer knows they should do better, yet it consistently falls behind feature work. What if you could point an agent at your codebase and have it produce structured, accurate, and consistent documentation automatically? That is exactly what Pydantic AI makes possible. In this guide, you will build a complete documentation generator that inspects Python source files, extracts meaningful information, and produces clean Markdown docs using a typed, schema-driven AI agent.

What Is Pydantic AI?

Pydantic AI is a Python framework for building production-grade AI agents. It is built by the same team behind Pydantic, the data validation library that powers FastAPI. The core idea is simple but powerful: instead of parsing free-form LLM output with regex and hope, you define typed schemas (Pydantic models) that describe exactly what the agent should return. The framework then handles validation, retries, and structured output for you.

Key features that make it ideal for a documentation generator:

• Type-safe outputs — Define a Pydantic model for your docs; the agent returns validated instances.
• Multiple model support — Swap between OpenAI, Anthropic, Gemini, Groq, Ollama, and others with one line.
• Dependency injection — Pass tools, databases, or file readers into agents cleanly.
• Tool calling — Let the agent call back into your code to inspect files, read ASTs, or query metadata.
• Streaming and validation — Stream partial results while still enforcing schema correctness.

Why a Documentation Generator?

Manual documentation suffers from three chronic problems: it drifts from the code, it is inconsistent in tone and structure, and it is expensive to maintain. An AI-driven generator addresses all three. By reading the actual source on every run, the docs stay in sync. By using a fixed schema, structure stays consistent. And by automating the bulk of the writing, you free maintainers to focus on review and refinement rather than blank-page authoring.

The approach we will take is schema-first: we decide what a "ModuleDoc" looks like before we ever prompt the model. This guarantees downstream tooling — static site generators, search indexers, linters — can consume our output reliably.

Project Setup

Create a new project directory and install the dependencies. We will use Pydantic AI with the OpenAI provider, but you can substitute any supported provider.

mkdir docgen-pydantic-ai
cd docgen-pydantic-ai
python -m venv .venv
source .venv/bin/activate  # on Windows: .venv\Scripts\activate

pip install pydantic-ai pydantic rich

Set your API key as an environment variable:

export OPENAI_API_KEY="sk-..."

Create the project layout:

docgen-pydantic-ai/
├── docgen/
│   ├── __init__.py
│   ├── models.py
│   ├── agent.py
│   ├── tools.py
│   └── runner.py
├── sample_pkg/
│   ├── __init__.py
│   └── calculator.py
└── main.py

Defining the Documentation Schema

The schema is the contract between the AI and the rest of your system. We will model documentation at three levels: a module, the functions/classes inside it, and individual parameters. Nesting Pydantic models gives the agent a clear target and gives us automatic validation.

Create docgen/models.py:

from __future__ import annotations
from typing import Literal, Optional
from pydantic import BaseModel, Field


class ParameterDoc(BaseModel):
    name: str = Field(..., description="Parameter name as it appears in the signature.")
    type_annotation: Optional[str] = Field(
        None, description="The type annotation, e.g. 'int', 'list[str]'."
    )
    description: str = Field(..., description="What this parameter represents.")
    required: bool = Field(True, description="Whether the parameter is required.")


class FunctionDoc(BaseModel):
    name: str
    kind: Literal["function", "method", "classmethod", "staticmethod"]
    signature: str = Field(..., description="Full signature, e.g. 'def add(a: int, b: int) -> int'.")
    summary: str = Field(..., description="One-sentence summary of what it does.")
    description: str = Field(..., description="Longer explanation, may include examples.")
    parameters: list[ParameterDoc] = Field(default_factory=list)
    returns: Optional[str] = Field(None, description="Description of the return value.")
    raises: list[str] = Field(
        default_factory=list,
        description="List of exceptions that may be raised, e.g. 'ValueError: when x is negative'.",
    )
    example: Optional[str] = Field(None, description="A short usage example in Python.")


class ModuleDoc(BaseModel):
    module_path: str = Field(..., description="Dython dotted path, e.g. 'sample_pkg.calculator'.")
    title: str = Field(..., description="Human-friendly module title.")
    summary: str = Field(..., description="One-paragraph overview of the module's purpose.")
    functions: list[FunctionDoc] = Field(default_factory=list)
    notes: list[str] = Field(
        default_factory=list,
        description="Implementation notes, caveats, or design decisions worth documenting.",
    )

Notice how every field carries a description. Pydantic AI feeds these into the model's prompt, so well-written descriptions dramatically improve output quality.

Building the Tools

Agents are more reliable when they can call tools to fetch ground truth rather than hallucinating from memory. We will give our agent two tools: one to list Python files in a package, and one to read a file's contents. We will also extract the AST so the agent has structured information about functions and signatures.

Create docgen/tools.py:

from __future__ import annotations
import ast
from pathlib import Path
from dataclasses import dataclass


@dataclass
class DocDeps:
    """Dependencies injected into the agent."""
    root: Path


def list_python_files(deps: DocDeps) -> list[str]:
    """Return relative paths of all .py files under the root, excluding __pycache__."""
    files = []
    for path in sorted(deps.root.rglob("*.py")):
        if "__pycache__" in path.parts:
            continue
        files.append(str(path.relative_to(deps.root)))
    return files


def read_source(deps: DocDeps, relative_path: str) -> str:
    """Read the raw source code of a file."""
    return (deps.root / relative_path).read_text(encoding="utf-8")


def extract_signatures(deps: DocDeps, relative_path: str) -> list[dict]:
    """Parse the file with ast and return function/class signatures."""
    source = read_source(deps, relative_path)
    tree = ast.parse(source)
    signatures = []

    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            args = [a.arg for a in node.args.args]
            returns = ast.unparse(node.returns) if node.returns else None
            signatures.append({
                "name": node.name,
                "kind": "function",
                "args": args,
                "returns": returns,
                "lineno": node.lineno,
            })
        elif isinstance(node, ast.ClassDef):
            methods = []
            for item in node.body:
                if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    methods.append(item.name)
            signatures.append({
                "name": node.name,
                "kind": "class",
                "methods": methods,
                "lineno": node.lineno,
            })

    return signatures

Keeping these as plain functions makes them easy to test independently of the agent. The DocDeps dataclass is the dependency container Pydantic AI will pass in.

Creating the Agent

Now we wire the schema and tools together into an agent. The agent's system prompt sets expectations, and the result_type enforces our schema.

Create docgen/agent.py:

from __future__ import annotations
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIModel

from .models import ModuleDoc
from .tools import DocDeps, list_python_files, read_source, extract_signatures


SYSTEM_PROMPT = """\
You are a senior technical writer and Python expert.

Your job is to produce accurate, concise documentation for a Python module.

Rules:
1. Always call the available tools to read the actual source code before writing docs.
2. Never invent functions, parameters, or behaviors that are not present in the code.
3. Use the extracted signatures as ground truth for names, arguments, and return types.
4. Write the summary in plain English, aimed at a developer reading the docs for the first time.
5. For each function, include a realistic example only if the function's purpose is non-obvious.
6. Note any edge cases, side effects, or non-thread-safe behavior in the 'notes' list.
7. If a function raises exceptions, list them in 'raises' with a brief reason.
"""


def build_agent(model_name: str = "gpt-4o-mini") -> Agent[DocDeps, ModuleDoc]:
    model = OpenAIModel(model_name)
    agent = Agent(
        model,
        deps_type=DocDeps,
        result_type=ModuleDoc,
        system_prompt=SYSTEM_PROMPT,
    )

    @agent.tool
    def list_files(ctx: RunContext[DocDeps]) -> list[str]:
        """List all Python files in the target package."""
        return list_python_files(ctx.deps)

    @agent.tool
    def get_source(ctx: RunContext[DocDeps], relative_path: str) -> str:
        """Read the full source code of a Python file."""
        return read_source(ctx.deps, relative_path)

    @agent.tool
    def get_signatures(ctx: RunContext[DocDeps], relative_path: str) -> list[dict]:
        """Return parsed function and class signatures for a file."""
        return extract_signatures(ctx.deps, relative_path)

    return agent

A few important details here. The Agent generic parameters [DocDeps, ModuleDoc] declare the dependency type and the result type. Tools are registered with @agent.tool, and each receives a RunContext whose .deps attribute holds our DocDeps instance. The docstrings on the tool functions become part of the tool schema the model sees, so write them carefully.

Rendering Markdown

Once we have a validated ModuleDoc, we want to render it to Markdown. Because the schema is typed, rendering is straightforward and deterministic.

Create docgen/runner.py:

from __future__ import annotations
from pathlib import Path

from .models import ModuleDoc


def render_markdown(doc: ModuleDoc) -> str:
    lines: list[str] = []
    lines.append(f"# {doc.title}")
    lines.append("")
    lines.append(f"_Module: `{doc.module_path}`_")
    lines.append("")
    lines.append(doc.summary)
    lines.append("")

    if doc.functions:
        lines.append("## Functions")
        lines.append("")
        for fn in doc.functions:
            lines.append(f"### `{fn.name}`")
            lines.append("")
            lines.append(f"python\n{fn.signature}\n")
            lines.append("")
            lines.append(f"**{fn.summary}**")
            lines.append("")
            lines.append(fn.description)
            lines.append("")

            if fn.parameters:
                lines.append("| Parameter | Type | Required | Description |")
                lines.append("|---|---|---|---|")
                for p in fn.parameters:
                    req = "yes" if p.required else "no"
                    lines.append(
                        f"| `{p.name}` | `{p.type_annotation or '—'}` | {req} | {p.description} |"
                    )
                lines.append("")

            if fn.returns:
                lines.append(f"**Returns:** {fn.returns}")
                lines.append("")

            if fn.raises:
                lines.append("**Raises:**")
                lines.append("")
                for r in fn.raises:
                    lines.append(f"- {r}")
                lines.append("")

            if fn.example:
                lines.append("**Example:**")
                lines.append("")
                lines.append(f"python\n{fn.example}\n")
                lines.append("")

    if doc.notes:
        lines.append("## Notes")
        lines.append("")
        for note in doc.notes:
            lines.append(f"- {note}")
        lines.append("")

    return "\n".join(lines)


def write_markdown(doc: ModuleDoc, output_dir: Path) -> Path:
    output_dir.mkdir(parents=True, exist_ok=True)
    filename = doc.module_path.replace(".", "_") + ".md"
    path = output_dir / filename
    path.write_text(render_markdown(doc), encoding="utf-8")
    return path

Putting It All Together

Now we need a sample package to document and an entry point that runs the agent over each file.

Create sample_pkg/calculator.py:

"""A small calculator module used to demonstrate the documentation generator."""


class Calculator:
    """A stateful calculator that accumulates results."""

    def __init__(self, initial: float = 0.0) -> None:
        self.value = initial

    def add(self, x: float) -> float:
        """Add x to the current value and return the new total."""
        self.value += x
        return self.value

    def reset(self) -> None:
        """Reset the accumulator to zero."""
        self.value = 0.0


def divide(numerator: float, denominator: float) -> float:
    """Divide numerator by denominator.

    Raises ValueError if denominator is zero.
    """
    if denominator == 0:
        raise ValueError("denominator must not be zero")
    return numerator / denominator


def average(values: list[float]) -> float:
    """Return the arithmetic mean of a non-empty list of numbers."""
    if not values:
        raise ValueError("values must not be empty")
    return sum(values) / len(values)

Create main.py:

from __future__ import annotations
import asyncio
from pathlib import Path

from docgen.agent import build_agent
from docgen.runner import write_markdown
from docgen.tools import DocDeps, list_python_files


async def generate_for_file(agent, deps: DocDeps, relative_path: str) -> None:
    module_path = relative_path.replace("/", ".").removesuffix(".py")
    if module_path.endswith(".__init__"):
        module_path = module_path[: -len(".__init__")]

    prompt = (
        f"Document the Python module at '{relative_path}'. "
        f"Its dotted path is '{module_path}'. "
        f"Read the source and signatures, then produce the ModuleDoc."
    )

    print(f"Generating docs for {relative_path} ...")
    result = await agent.run(prompt, deps=deps)
    out = write_markdown(result.data, Path("docs"))
    print(f"  -> wrote {out}")


async def main() -> None:
    root = Path("sample_pkg")
    deps = DocDeps(root=root)
    agent = build_agent()

    files = list_python_files(deps)
    if not files:
        print("No Python files found.")
        return

    for f in files:
        await generate_for_file(agent, deps, f)

    print("Done.")


if __name__ == "__main__":
    asyncio.run(main())

Run it:

python main.py

You should see output like:

Generating docs for calculator.py ...
  -> wrote docs/sample_pkg_calculator.md
Done.

Open docs/sample_pkg_calculator.md and you will find a structured page with a summary, a function table for each callable, return descriptions, raised exceptions, and notes — all derived from the actual source.

Best Practices

• Keep tools deterministic. Tools should return facts (file contents, AST nodes), not opinions. Let the model do the writing; let your code do the reading.
• Constrain the schema aggressively. Use Literal types, min_length, and enums wherever possible. The narrower the schema, the fewer hallucinations you get.
• Provide rich field descriptions. Each Field(description=...) becomes part of the prompt. Treat them as instructions to the model.
• Validate before trusting. Pydantic AI validates automatically, but add a post-processing pass to check that every documented function actually exists in the source. Mismatches usually mean the model invented something.
• Cache aggressively. Re-running the agent on unchanged files wastes tokens. Hash each file and skip generation when the hash matches a stored value.
• Use a cheaper model for iteration. Develop your prompts with gpt-4o-mini or a local Ollama model, then switch to a stronger model for final generation runs.
• Review, do not blindly publish. Treat AI-generated docs as a pull request: a human should approve before they ship. The schema guarantees structure, not correctness.
• Log the full conversation. Pydantic AI exposes message history. Persist it so you can debug bad outputs and improve prompts over time.

Extending the Generator

Once the core works, there are several high-value extensions:

• Cross-references. Add a post-processing step that turns function names in descriptions into Markdown links to their sections.
• Multi-file modules. Run the agent per file but generate an index page that links them together.
• Change-driven generation. Integrate with git diff so only modified files are re-documented.
• Doctest extraction. Parse the example field and run it through doctest in CI to catch drift.
• Provider portability. Swap OpenAIModel for AnthropicModel or OllamaModel with no other code changes — useful for offline or cost-sensitive environments.

Conclusion

Building a documentation generator with Pydantic AI demonstrates the framework's core strength: combining the flexibility of LLMs with the rigor of typed schemas. By defining a ModuleDoc model up front, giving the agent deterministic tools to read real source code, and rendering validated output to Markdown, you get documentation that is structured, consistent, and grounded in what the code actually does. The schema acts as a contract that keeps the model honest, the tools keep it accurate, and the rendering layer keeps the output consumable by humans and machines alike. From here, the natural next steps are caching, change detection, and wiring the generator into your CI pipeline so documentation updates itself on every commit — turning a chronic maintenance burden into an automated, reviewable workflow.