TOPTRACING

PydanticAI: Agent Framework / shim to use Pydantic with LLMs

Dec 5, 2024

0
https://ai.pydantic.dev/https://github.com/pydantic/pydantic-ai

PydanticAI is a Python Agent Framework designed to make it less painful to build production grade applications with Generative AI.

Why use PydanticAI

  • Built by the team behind Pydantic (the validation layer of the OpenAI SDK, the Anthropic SDK, LangChain, LlamaIndex, AutoGPT, Transformers, CrewAI, Instructor and many more)
  • Model-agnostic — currently OpenAI, Gemini, and Groq are supported, Anthropic is coming soon. And there is a simple interface to implement support for other models.
  • Type-safe
  • Control flow and agent composition is done with vanilla Python, allowing you to make use of the same Python development best practices you'd use in any other (non-AI) project
  • Structured response validation with Pydantic
  • Streamed responses, including validation of streamed structured responses with Pydantic
  • Novel, type-safe dependency injection system, useful for testing and eval-driven iterative development
  • Logfire integration for debugging and monitoring the performance and general behavior of your LLM-powered application

Tools & Dependency Injection Example

Here is a concise example using PydanticAI to build a support agent for a bank:

from dataclasses import dataclass

from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

from bank_database import DatabaseConn


@dataclass
class SupportDependencies:  
    customer_id: int
    db: DatabaseConn  


class SupportResult(BaseModel):  
    support_advice: str = Field(description='Advice returned to the customer')
    block_card: bool = Field(description="Whether to block the customer's card")
    risk: int = Field(description='Risk level of query', ge=0, le=10)


support_agent = Agent(  
    'openai:gpt-4o',  
    deps_type=SupportDependencies,
    result_type=SupportResult,  
    system_prompt=(  
        'You are a support agent in our bank, give the '
        'customer support and judge the risk level of their query.'
    ),
)


@support_agent.system_prompt  
async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
    customer_name = await ctx.deps.db.customer_name(id=ctx.deps.customer_id)
    return f"The customer's name is {customer_name!r}"


@support_agent.tool  
async def customer_balance(
    ctx: RunContext[SupportDependencies], include_pending: bool
) -> float:
    """Returns the customer's current account balance."""  
    return await ctx.deps.db.customer_balance(
        id=ctx.deps.customer_id,
        include_pending=include_pending,
    )

...  

async def main():
    deps = SupportDependencies(customer_id=123, db=DatabaseConn())
    result = await support_agent.run('What is my balance?', deps=deps)  
    print(result.data)  
    """
    support_advice='Hello John, your current account balance, including pending transactions, is $123.45.' block_card=False risk=1
    """

    result = await support_agent.run('I just lost my card!', deps=deps)
    print(result.data)
    """
    support_advice="I'm sorry to hear that, John. We are temporarily blocking your card to prevent unauthorized transactions." block_card=True risk=8
    """

The code included here is incomplete for the sake of brevity (the definition of DatabaseConn is missing); you can find the complete example here.

Instrumentation with Pydantic Logfire

To understand the flow of the above runs, we can watch the agent in action using Pydantic Logfire.

To do this, we need to set up logfire, and add the following to our code:

...
from bank_database import DatabaseConn

import logfire
logfire.configure()  
logfire.instrument_asyncpg()  
...

That's enough to get the following view of your agent in action:


See Monitoring and Performance to learn more.

Usage

Examples are distributed with pydantic-ai so you can run them either by cloning the pydantic-ai repo or by simply installing pydantic-ai from PyPI with pip or uv.

Installing required dependencies

Either way you'll need to install extra dependencies to run some examples, you just need to install the examples optional dependency group.

If you've installed pydantic-ai via pip, you can install the extra dependencies with:

pip install 'pydantic-ai[examples]'

If you clone the repo, you should instead use uv sync --extra examples to install extra dependencies.

Setting model environment variables

These examples will need you to set up authentication with one or more of the LLMs, see the model configuration docs for details on how to do this.

TL;DR: in most cases you'll need to set one of the following environment variables:

OpenAI

export OPENAI_API_KEY=your-api-key

Google Gemini

export GEMINI_API_KEY=your-api-key

Running Examples

To run the examples (this will work whether you installed pydantic_ai, or cloned the repo), run:

python -m pydantic_ai_examples.<example_module_name>

For examples, to run the very simple pydantic_model example:

python -m pydantic_ai_examples.pydantic_model

If you like one-liners and you're using uv, you can run a pydantic-ai example with zero setup:

OPENAI_API_KEY='your-api-key' \
  uv run --with 'pydantic-ai[examples]' \
  -m pydantic_ai_examples.pydantic_model

You'll probably want to edit examples in addition to just running them. You can copy the examples to a new directory with:

python -m pydantic_ai_examples --copy-to examples/

Agents

Agents are PydanticAI's primary interface for interacting with LLMs.

In some use cases a single Agent will control an entire application or component, but multiple agents can also interact to embody more complex workflows.

The Agent class has full API documentation, but conceptually you can think of an agent as a container for:

  • A system prompt — a set of instructions for the LLM written by the developer
  • One or more retrieval tool — functions that the LLM may call to get information while generating a response
  • An optional structured result type — the structured datatype the LLM must return at the end of a run
  • A dependency type constraint — system prompt functions, tools and result validators may all use dependencies when they're run
  • Agents may optionally also have a default LLM model associated with them; the model to use can also be specified when running the agent

In typing terms, agents are generic in their dependency and result types, e.g., an agent which required dependencies of type Foobar and returned results of type list[str] would have type cAgent[Foobar, list[str]]. In practice, you shouldn't need to care about this, it should just mean your IDE can tell you when you have the right type, and if you choose to use static type checking it should work well with PydanticAI.

Here's a toy example of an agent that simulates a roulette wheel:

from pydantic_ai import Agent, RunContext

roulette_agent = Agent(  
    'openai:gpt-4o',
    deps_type=int,
    result_type=bool,
    system_prompt=(
        'Use the `roulette_wheel` function to see if the '
        'customer has won based on the number they provide.'
    ),
)


@roulette_agent.tool
async def roulette_wheel(ctx: RunContext[int], square: int) -> str:  
    """check if the square is a winner"""
    return 'winner' if square == ctx.deps else 'loser'


# Run the agent
success_number = 18  
result = roulette_agent.run_sync('Put my money on square eighteen', deps=success_number)
print(result.data)  
#> True

result = roulette_agent.run_sync('I bet five is the winner', deps=success_number)
print(result.data)
#> False

Agents are intended to be instantiated once (frequently as module globals) and reused throughout your application, similar to a small FastAPI app or an APIRouter.

For more detail you can see Agents.