Type-Safe Tool Definitions with Pydantic AI

Define agent tools with typed parameters, validated args, clear docstrings, and dependency injection so the model cannot invent shapes your code cannot run.

Summary

Register Python functions with @agent.tool / @agent.tool_plain (or Tool(...)) so Pydantic AI builds JSON schemas from type hints, validates model-supplied args, and returns results into the agent loop.

Recipe

1. Write a pure-looking function with explicit parameter types and a return type.
2. Document the tool and each parameter (Google/Numpy/Sphinx style) so descriptions enter the schema.
3. Use @agent.tool_plain when the tool needs no deps; use @agent.tool with RunContext[Deps] when it does.
4. Prefer small JSON-serializable returns; structure complex returns as Pydantic models when useful.
5. Keep side effects idempotent where retries are possible; validate destructive args before execute.
6. Limit the tool allowlist per agent role; do not expose kitchen-sink admin tools by default.
7. Test tools as ordinary functions, then with TestModel / eval runs for call patterns.

Working Example

from dataclasses import dataclass
from datetime import date
 
from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
 
 
@dataclass
class BookingDeps:
    user_id: str
    # db: DatabaseClient  # inject real clients in production
 
 
class BookingResult(BaseModel):
    confirmation_id: str
    nights: int
    total_usd: float
 
 
class HotelAgentOut(BaseModel):
    message: str
    booking: BookingResult | None = None
 
 
agent = Agent(
    "openai:gpt-4o-mini",  # verify at build
    deps_type=BookingDeps,
    output_type=HotelAgentOut,
    instructions=(
        "Help users check availability and book stays. "
        "Use tools for facts; finish with HotelAgentOut."
    ),
)
 
 
@agent.tool_plain
def list_room_types() -> list[str]:
    """Return the room types this property offers."""
    return ["standard", "deluxe", "suite"]
 
 
@agent.tool
def quote_stay(
    ctx: RunContext[BookingDeps],
    room_type: str,
    check_in: date,
    nights: int,
) -> dict:
    """Price a stay for the current user.
 
    Args:
        room_type: One of standard, deluxe, suite.
        check_in: Arrival date (YYYY-MM-DD).
        nights: Number of nights between 1 and 30.
    """
    if nights < 1 or nights > 30:
        return {"error": "nights must be between 1 and 30", "user_id": ctx.deps.user_id}
    rates = {"standard": 120.0, "deluxe": 180.0, "suite": 280.0}
    if room_type not in rates:
        # Returning an error string is ok; raising ModelRetry can re-ask
        return {"error": f"unknown room_type {room_type}", "user_id": ctx.deps.user_id}
    total = rates[room_type] * nights
    return {
        "user_id": ctx.deps.user_id,
        "room_type": room_type,
        "check_in": check_in.isoformat(),
        "nights": nights,
        "total_usd": total,
    }
 
 
@agent.tool
def confirm_booking(
    ctx: RunContext[BookingDeps],
    room_type: str,
    check_in: date,
    nights: int,
) -> BookingResult:
    """Finalize a booking after the user agrees to the quote.
 
    Args:
        room_type: Selected room type.
        check_in: Arrival date.
        nights: Number of nights.
    """
    # Call your booking service with ctx.deps.user_id here.
    quote = quote_stay(ctx, room_type, check_in, nights)
    if "error" in quote:
        raise ValueError(quote["error"])
    return BookingResult(
        confirmation_id=f"BKG-{ctx.deps.user_id[:4].upper()}-001",
        nights=nights,
        total_usd=float(quote["total_usd"]),
    )
 
 
result = agent.run_sync(
    "Quote a deluxe room for 2 nights starting 2026-08-01, then book it if under $400.",
    deps=BookingDeps(user_id="user-42"),
)
print(result.output)

Notes:

• Parameter types become the tool JSON schema.
• Docstrings become tool and parameter descriptions for the model.
• Invalid args fail validation and can be reflected back for retry.

Deep Dive

Registration styles

Schema generation rules

• All parameters except RunContext are model-visible.
• Single object parameters (Pydantic model / dataclass / TypedDict) can simplify to that object schema.
• Return values must be JSON-serializable via Pydantic's serialization rules.
• Advanced returns can use ToolReturn for model-visible content plus metadata (for example UI data chunks).

Type safety across the stack

1. Write-time: type checkers verify RunContext[Deps] usage.
2. Call-time: Pydantic validates model args before your function runs.
3. Return-time: structured returns give the model cleaner observations.
4. Final output: output_type remains a separate contract for the end of the run.

Tools and final output are both structured surfaces. Do not conflate them: tools intermediate, output_type terminates (in default configurations).

Designing for least privilege

• One agent should not receive every internal API.
• Prefer read tools separate from write tools.
• Gate irreversible tools with approval flows when using Vercel AI streaming (requires_approval / deferred tools - verify at build).

Errors and retries

• Validation errors on args are framework-mediated retries.
• Application errors can return structured error objects or raise ModelRetry with repair instructions.
• Cap tool retries so a broken integration cannot loop forever.

Testing strategy

• Unit-test pure logic without the model.
• Use Pydantic AI test/function models to assert which tools are offered and how schemas look.
• Eval real trajectories for unnecessary tool calls and wrong argument patterns.

Gotchas

1. Missing type hints. Untyped parameters produce weak schemas and weak validation.
2. Huge tool returns. Dumping entire DB rows blows the context window; return digests + ids.
3. Side effects on speculative calls. Models may call write tools while exploring; require confirmation for high risk.
4. Docstring drift. Outdated docs mis-train the model; treat them as public API.
5. Overloading one tool. Giant multi-purpose tools hurt selection accuracy; split by intent.
6. Leaking secrets via deps. Do not put raw API keys into tool return values or prompts.
7. Ignoring RunContext typing. Wrong deps_type vs annotation is a static bug you want caught early.

Alternatives

FAQs

Related

• Pydantic AI & Instructor Basics - first tool + output demo
• Defining Agent Output Schemas with Pydantic Models - final output contracts
• Handling Validation Failures and Retry Loops Gracefully - retries on bad args/results
• Pydantic AI's Native Vercel AI SDK Data Stream Support - streaming tool events to UI
• Why Structured Output Matters More as Agents Get Autonomous - structure at every boundary
• Pydantic AI & Instructor Best Practices - section checklist

Stack versions: Pins from the category manifest (verify at build): OpenRouter (~315+ models, July 2026 pricing/fees); LangGraph 1.0+; CrewAI 1.14+; Microsoft Agent Framework 1.0; Vercel AI SDK 6; Pydantic AI (latest); LlamaIndex (latest); OpenAI Agents SDK (latest + MCP); MCP (Linux Foundation governance); A2A (HTTP+SSE+JSON-RPC 2.0); Solana @solana/web3.js + @solana/spl-token.