Type-safe Tools with Pydantic AI

Why Type Safety Matters for Agents

In traditional software, type errors surface at compile time or in the first test run. In agent systems, type errors surface in production when the model returns an unexpected structure, the tool fails to parse arguments, or an output gets silently truncated.

Pydantic AI was built specifically to bring Python’s type system to AI agents. The result: agents where every input is validated, every output conforms to a schema, and every tool is a normal Python function.

Defining Your First Pydantic AI Agent

Pydantic AI agents are built from three pieces: an output schema (what you get back), a dependencies dataclass (what you inject in), and an Agent instance that ties them together.

# src/agents/pydantic_agent.py
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext
from dataclasses import dataclass

# 1. Define the output schema: the agent will not return until it produces data matching this shape
class ResearchResult(BaseModel):
    topic: str = Field(description="The topic that was researched")
    summary: str = Field(description="A concise summary (2-3 paragraphs)")
    key_facts: list[str] = Field(description="5-7 key facts")
    confidence: float = Field(ge=0.0, le=1.0, description="Confidence in accuracy (0-1)")  # ge/le enforce range
    sources_searched: list[str] = Field(description="Search queries used")

# 2. Define dependencies: use a dataclass so Pydantic AI can type-check the injection at run time
@dataclass
class AgentDeps:
    search_api_key: str
    max_results: int = 5

# 3. Create the agent: result_type drives automatic schema generation and output validation
research_agent = Agent(
    model="openai:gpt-5.6-sol",          # format is "provider:model-name", Pydantic AI handles the client
    result_type=ResearchResult,      # agent retries until it produces a valid ResearchResult
    system_prompt=(
        "You are a research assistant. Use the search tool to find accurate, "
        "up-to-date information. Always verify facts with multiple searches. "
        "Be specific about what you found and what you're uncertain about."
    ),
    retries=3,  # on validation failure, send the Pydantic error back to the model and retry up to 3 times
)

Registering Tools with @agent.tool

The @agent.tool decorator inspects your function’s type hints and docstring to build the JSON schema automatically, no manual schema writing needed.

import asyncio
import httpx
from pydantic_ai import Agent, RunContext

@research_agent.tool
async def search_web(ctx: RunContext[AgentDeps], query: str) -> list[dict]:
    """Search the web for current information.
    
    Args:
        query: The search query. Be specific for better results.
    
    Returns:
        List of search results with title, url, and snippet.
    """
    # ctx.deps gives access to injected AgentDeps: avoids global state for API keys
    async with httpx.AsyncClient() as client:
        response = await client.get(
            "https://api.search-provider.com/search",
            params={"q": query, "num": ctx.deps.max_results},  # max_results comes from injected deps
            headers={"Authorization": f"Bearer {ctx.deps.search_api_key}"}
        )
        response.raise_for_status()  # raises HTTPStatusError on 4xx/5xx, Pydantic AI surfaces this to the model
        return response.json()["results"]

@research_agent.tool
async def fetch_page(ctx: RunContext[AgentDeps], url: str) -> str:
    """Fetch the full text content of a web page.
    
    Args:
        url: The URL to fetch.
    
    Returns:
        The text content of the page (first 5000 chars).
    """
    async with httpx.AsyncClient() as client:
        response = await client.get(url, follow_redirects=True, timeout=10.0)
        response.raise_for_status()
        # In production, use BeautifulSoup to extract text from HTML
        return response.text[:5000]  # truncate to avoid flooding the model's context window

What Pydantic AI does automatically:

• Inspects query: str → JSON schema {"type": "string"}
• Reads the docstring → tool description for the model
• Injects ctx → carries AgentDeps without it appearing in the LLM’s tool schema
• Validates arguments before your function is called

Running the Agent with Dependency Injection

async def research_topic(topic: str) -> ResearchResult:
    # Construct deps here: keeps secrets out of global scope and makes testing easy (just swap deps)
    deps = AgentDeps(
        search_api_key="your-search-api-key",
        max_results=5,
    )
    
    result = await research_agent.run(
        f"Research the following topic and provide a comprehensive summary: {topic}",
        deps=deps,  # deps are injected into every tool call via RunContext
    )
    
    # result.data is a validated ResearchResult instance: guaranteed to match the schema
    return result.data

# Usage
import asyncio

async def main():
    result = await research_topic("Model Context Protocol (MCP) for AI agents")
    
    print(f"Topic: {result.topic}")
    print(f"Confidence: {result.confidence:.0%}")
    print(f"\nSummary:\n{result.summary}")
    print(f"\nKey Facts:")
    for fact in result.key_facts:
        print(f"  • {fact}")

asyncio.run(main())

Structured Output with Validation

The real power of Pydantic AI is what happens when the model returns malformed data:

class ExtractedInvoice(BaseModel):
    invoice_number: str
    vendor: str
    amount: float = Field(gt=0)                        # gt=0 rejects zero or negative amounts
    currency: str = Field(pattern=r'^[A-Z]{3}$')       # regex ensures ISO 4217 format (e.g. "USD", "GBP")
    due_date: str  # will validate with a validator
    line_items: list[dict]
    
    @field_validator('due_date')
    @classmethod
    def validate_date_format(cls, v: str) -> str:
        from datetime import datetime
        datetime.strptime(v, '%Y-%m-%d')  # raises ValueError if format doesn't match YYYY-MM-DD
        return v

invoice_agent = Agent(
    model="anthropic:claude-sonnet-5",
    result_type=ExtractedInvoice,
    system_prompt="Extract structured invoice data from the provided text.",
)

# If Claude returns amount=-50, the Pydantic validator catches it
# and sends the error back to Claude to correct before your code sees it
#, no manual retry logic needed

Testing Pydantic AI Agents

# tests/test_research_agent.py
import pytest
from pydantic_ai.testing import TestModel, capture_run_messages
from src.agents.pydantic_agent import research_agent, AgentDeps, ResearchResult

@pytest.mark.asyncio
async def test_research_agent_structure():
    deps = AgentDeps(search_api_key="test-key", max_results=3)
    
    # TestModel replaces the real LLM with a deterministic stub: no API calls, no cost, no flakiness
    with research_agent.override(model=TestModel()):
        result = await research_agent.run("What is RAG?", deps=deps)
        
    # Verify the result is a valid ResearchResult: confirms schema and validation logic work
    assert isinstance(result.data, ResearchResult)
    assert result.data.topic != ""
    assert 0 <= result.data.confidence <= 1

@pytest.mark.asyncio
async def test_tool_called_with_correct_args():
    deps = AgentDeps(search_api_key="test-key")
    
    # capture_run_messages records every message exchanged during the run for inspection
    with capture_run_messages() as messages:
        with research_agent.override(model=TestModel()):
            await research_agent.run("Research quantum computing", deps=deps)
    
    # Inspect what tools were called and with what arguments: lets you assert on agent behavior
    tool_calls = [m for m in messages if hasattr(m, 'tool_name')]
    assert any(tc.tool_name == 'search_web' for tc in tool_calls)

When to Choose Pydantic AI vs LangGraph

They complement each other well, a common pattern is using Pydantic AI agents as worker nodes inside a LangGraph workflow, getting structured outputs at each step that LangGraph routes based on.