Building your first AI Agent using PydanticAI

AI Agents

An AI agent is a software entity that can perceive its environment through sensors (APIs, user inputs or datastreams), process that information and act upon it autonomously to acheive specific goals or tasks.

AI agents have following characteristics:

• Automony: Can operate without human intervention for certain tasks.
• Reactivity: Can respond to changes in the environment.
• Proactivity: Can initiate action to achieve their goals.
• Social ability: Can interact with other agents or humans, often through natural language or programmatic interfaces.

AI Agents range from simple bots that follow predefined rules to sophisticated systems using machine learning to learn from and adapt to new situations like chantbots, recommendation systems, autonomous vehicles etc.

Pydantic

Pydantic is the most used data validation library for Python. If offers schema validation and serialization for data controlled by type annotations.

Logfire is an application monitoring tool that integrates with many popular Python libraries including FastAPI, OpenAI, and Pydantic. You can use Logfire to monitor Pydantic validations and understand why some inputs fail validation.

PydanticAI

PydanticAI is a Python framework for creation and management of AI agents that use Large Language Models.

It extends the capabilities of the Pydantic Library which is renowned for data validation and settings management using Python type annotations.

PydanticAI focusses on:

1. Agent Creation- Provides abstractions and tools to easily define AI agents that can understand, process and act upon natural language inputs in a structured manner.
2. LLM Integration: Allows to integrate LLMs while handling the complexities of model interactions such as API management, prompt engineering and response parsing.
3. Data Handling : Ensures that data flowing through the system is type-safe and well-structured.

Advantages of PydanticAI:

1. Model Agnostic Design: It works with OpenAI, xAI, Gemini, Open Source Models without requiring significant code changes.
2. Type Safety:Ensures that the data flowing through the systems is type-sage and well structured (response of LLMs fits predefined schemas or formats).
3. Integrates well with other Python libraries like Pydantic.
4. Easy to use.
5. PydanticAI absracts the complexity of dealing with LLM APIs. Simplifies agent development by abstracting away much of the boilerplate code needed for LLM interactions. This includes managing API calls, handling tokens, construvting prompts so that you can focus on logic and behaviour of the agent. Provides built-in support for adding tools or capabilities to agents like external API calls or database queries. Thus provides a clean interface for developers to work with.
6. Supports dependency injection for managing resources or services the agent might need to interact with.
7. Consistency: Enforces structured outputs through Pydantic models.
8. Offers development Tools and utilities for testing AI agents.
9. Developers can attach tools to agents, allowing them to perform actions like data fetch and computation besides generating responses from LLMs.

What is an agent in the context of PydanticAI?

• An agent in PydanticAI is a construct that efficiently interfaces with LLMs to process natural language inputs and produce structured outputs.
• Agents use Pydantic models to define the expected struvture of inputs and outputs, ensuring data integrity and type-safe outputs.
• Agents are task oriented — answering questions, generating content, making decisions based on the input data.
• Agents can encapsulate functions that when triggered by LLM responses, execute specific logic or interact with the external systems like APIs or databases.
• Agents can generate or modify prompts dynamically.

How does PydanticAI use Pydantic models for structured outputs in AI applications.

PydanticAI leverages Pydantic models in several key ways to manage structured outputs in AI applications, particularly when interfacing with Large Language Models (LLMs).

Schema Definition: Pydantic models are used to define the exact schema or structure of the output that an AI agent should produce. This means specifying the data types, nested structures, and validations for each field. For example:

from pydantic import BaseModel
class UserResponse(BaseModel):
 name: str
 age: int
 interests: List[str]        

With this model, when an LLM generates output, PydanticAI ensures that the response adheres to this structure.

2. Automatic Validation: Once defined, Pydantic models automatically validate the data against the schema. This includes type checking, ensuring required fields are present, and applying any custom validators. If the LLM’s response doesn’t match the expected schema, PydanticAI can catch these errors, either by correcting the data or raising an exception for further handling.

3. Type Hinting: Pydantic models provide type hints which can be used during development for better code completion, static type checking with tools like MyPy, and runtime type enforcement, making the integration with Python’s ecosystem more robust.

4. Dynamic and Complex Output Handling: Pydantic models can include optional or dynamic fields, allowing for flexibility in how LLMs can respond while still maintaining structure. For instance, an agent might only return a subset of fields based on the query context.For more complex outputs, you can nest models within models, representing hierarchical data structures that LLMs might need to understand or generate.

5. Integration with LLM Outputs: PydanticAI uses the defined models to parse and structure the sometimes chaotic outputs from LLMs. This parsing can involve converting free-form text into structured data, which is crucial for applications where data integrity is paramount.When constructing prompts for LLMs, PydanticAI can generate or modify prompts to include hints or examples of the expected output structure, guiding the LLM towards producing data that fits the Pydantic model.

6. Error Handling and Debugging: If the LLM’s output doesn’t conform to the expected schema, Pydantic provides detailed error messages about which fields failed validation, aiding in debugging. PydanticAI can implement logic for retrying LLM queries with adjusted prompts if the initial output fails validation, or provide fallback mechanisms for handling malformed data.

from pydantic import BaseModel
from pydanticai import Agent
class BookRecommendation(BaseModel):
    title: str
    author: str
    genre: str
agent = Agent(
    description="Suggests a book based on user input",
    output_model=BookRecommendation
)
# When calling this agent with an LLM, the response would be forced into this structure:
response = agent.invoke("I need a sci-fi book to read")
# response would be validated to ensure it fits the BookRecommendation schema
        

In this example, PydanticAI ensures that the LLM’s response is parsed and validated against the BookRecommendation model.

Installation

PydanticAI requires Python3.8 or above. Check your python version and upgrade if required.

python --version        

Install pydantic-ai using pip

pip install pydantic-ai        

For a leaner installation that includes only the necessary dependencies for your model usage, you can use pydantic-ai-slim. For instance, if you’re only using OpenAI’s models:

pip install pydantic-ai-slim[openai]        

Dependencies

• Python: Version 3.8 or higher.
• Pydantic: PydanticAI is built on top of Pydantic, which is installed automatically when you install PydanticAI. However, ensure you have at least Pydantic version 2.x installed as PydanticAI leverages newer features.
• typing_extensions: This is usually installed with Pydantic, providing runtime support for Python type hints.

Model-Specific Libraries:

• For OpenAI models, openai might be required if you’re not using the slim version.
• For Google’s Vertex AI, google-cloud-aiplatform is necessary.
• Logfire: If you’re interested in debugging and performance monitoring, logfire needs to be installed as an optional dependency.

Optional Dependencies

Depending on the LLM providers or additional features you want to use, you might need to install optional dependencies:

Install OpenAI Support:

pip install pydantic-ai[openai]        

Install Vertex AI Support:

pip install pydantic-ai[vertexai]        

Install Logfire for Debugging:

pip install pydantic-ai[logfire]        

Install Examples for Learning:

pip install pydantic-ai[examples]        

If you’re using multiple models or features, you can chain these extras:

pip install pydantic-ai[openai,vertexai,logfire,examples]        

API Keys

For models like those from OpenAI or Google’s Vertex AI, you’ll need to set up API keys.

export OPENAI_API_KEY='your-api-key'        

export GOOGLE_APPLICATION_CREDENTIALS='path/to/your/credentials.json'        

%env GOOGLE_APPLICATION_CREDENTIALS='path/to/your/credentials.json'        

Ensure these keys are secure and not shared or committed to version control systems.

After installation, you can verify that PydanticAI is installed by running:

python -c "import pydantic_ai; print(pydantic_ai.__version__)"        

Let us see some code examples

First, we will use ollama model llama3.2 for offline simple chat using PydanticAI agent:

• You should install ollama from here.
• Then download the required models from ollama library. I have downloaded llama3.2 on my computer.

#install pydantic-ai if not already installed
#pip install pydantic-ai

import warnings
warnings.filterwarnings("ignore")

from pydantic_ai import Agent
#we will chat with ollama model llama3.2
agent = Agent('ollama:llama3.2')
while True:
    user_input=input("Enter a prompt:")
    if user_input.lower()=='exit':
        print("you exitted")
        break
    else:
        result = agent.run_sync(user_input)
        print(result.data)        

• Save above code in a file called chat1.py and run it at Anaconda prompt

python chat1.py        

Another example

Let us see another example, it is pydantic_model.py file built-in to examples in PydanticAI. We will use this example with gemini-1.5-flash model instead of the default OpenAI models. To run this file with gemini-1.5-flash LLM you will set an environment variable PYDANTIC_AI_MODEL = gemini-1.5-flash the run the file at Anaconda Prompt.

#Save this code as pydantic_model.py
import os
from typing import cast

import logfire
from pydantic import BaseModel

from pydantic_ai import Agent
from pydantic_ai.models import KnownModelName

# 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configured
logfire.configure(send_to_logfire='if-token-present')

# A class inherited from BaseModel helps you define format of LLMs output
class MyModel(BaseModel):
    city: str
    country: str


#you have saved gemini-1.5-flash in PYDANTIC_AI_MODEL environment variable
model = cast(KnownModelName, os.getenv('PYDANTIC_AI_MODEL', 'openai:gpt-4o'))
print(f'Using model: {model}')
agent = Agent(model, result_type=MyModel)

if __name__ == '__main__':
    #The agent.run_sync function sends a prompt to the LLM
    result = agent.run_sync('The windy city in the US of A.')
    print(result.data)
    print(result.cost())        

Running above file:

set PYDANTIC_AI_MODEL=gemini-1.5-flash
python -m pydantic_ai_examples.pydantic_model

or 
set PYDANTIC_AI_MODEL=gemini-1.5-flash
python pydantic_model.py        

Let us build a simple Bank Support AI agent that uses gemini-1.5-flash LLM.

• This is a builtin example in PydanticAI.
• It demonstrates use of dynamic system prompt.
• This agent uses a database to authenticate user credentials.
• Uses BaseModel to define format of output as result_type
• Uses gemini-1.5-flash LLM
• Uses tools (functions) to fetch balance from account of an authenticated user.
• Save below code in a file called bank_support.py
• Run this file at Anaconda Prompt as:

python bank_support.py
or
python -m pydantic_ai_examples.bank_support        

#this code is present in pydantic examples as bank_support.py
from dataclasses import dataclass
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

import warnings
warnings.filterwarnings("ignore")

class DatabaseConn:
    """This is a fake database for example purposes.
    In reality, you'd be connecting to an external database
    (e.g. PostgreSQL) to get information about customers.
    It encapsulates customer name and account balance for customer 
    whose customer id is provided.
    """

    @classmethod
    async def customer_name(cls, *, id: int) -> str | None:
        if id == 123:
            return 'John'

    @classmethod
    async def customer_balance(cls, *, id: int, include_pending: bool) -> float:
        if id == 123:
            return 123.45
        else:
            raise ValueError('Customer not found')


@dataclass
class SupportDependencies:
    customer_id: int
    db: DatabaseConn

#format of outputs, text message, card status and risk associate with acount
class SupportResult(BaseModel):
    support_advice: str = Field(description='Advice returned to the customer')
    block_card: bool = Field(description='Whether to block their')
    risk: int = Field(description='Risk level of query', ge=0, le=10)

#configuration of Agent
support_agent = Agent(
    'gemini-1.5-flash',
    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. '
        "Reply using the customer's name."
    ),
)

#Define a dynamic system prompt
#RunContext is a dependency injection mechanism that injects dependencies into 
#system prompt functions, tools and result validators. This allows these components
#to access external data or services.When defining functions within an agent,
#such as system prompts or tools, RunContext is often the first parameter. 
#This allows these functions to interact with or retrieve data from the 
#dependencies.
@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}"

#Agent uses this function to fetch user's account balance. It is a tool that 
#allows agent to connect with database and fetch user's details.
@support_agent.tool
async def customer_balance(
    ctx: RunContext[SupportDependencies], include_pending: bool) -> str:
    """Returns the customer's current account balance."""
    balance = await ctx.deps.db.customer_balance(
        id=ctx.deps.customer_id,
        include_pending=include_pending,
    )
    return f'${balance:.2f}'

#Specific instance of SupportDependencies class. 
deps = SupportDependencies(customer_id=123, db=DatabaseConn())

#execute a synchronous task using the agent. This method blocks till operation 
#is complete.
result = support_agent.run_sync('What is my balance?', deps=deps)
#printing only the support_advice part of result
print(result.data.support_advice)

#Output
'Hello John, your current account balance, including pending transactions, is $123.45.'

#print full result
print(result.data)

Output: 
support_advice='Hello John, your current account balance, including pending transactions, is $123.45.' block_card=False risk=1

#run the agent again
result = support_agent.run_sync('I just lost my card!', deps=deps)
print(result.data.support_advice)

#Output:
"I'm sorry to hear that, John. We are temporarily blocking your card to prevent unauthorized transactions."


#print full result
print(result.data)

Output: 
support_advice="I'm sorry to hear that, John. We are temporarily blocking your card to prevent unauthorized transactions." block_card=True risk=8        

There you go, your first AI agent is up and running.

References:

https://ai.pydantic.dev/

https://ai.pydantic.dev/examples/