Pydantic AI : Agent Framework

The Pydantic AI Agent Framework is a powerful tool for building agentic AI systems with robust data validation, flexible execution models, and the ability to seamlessly integrate tools. It leverages the structured validation of Pydantic to ensure input and output consistency, making it great for applications ranging from simple data processing to advanced conversational AI.

I have implented a hands-on following the Pydantic documentation and it is pretty easy to follow. JupyterNotebook attached in the end.

Key Features of Pydantic AI Agent Framework

• Structured Validation: Derived from Pydantic, the framework ensures type-safe data handling and validation.
• Execution Flexibility: Supports synchronous, asynchronous, and streamed execution models.
• Tool Integration: Provides decorators and configuration options for integrating custom tools into agents.
• Text Streaming: Enables real-time interaction with progressive output streaming.

Installation and Setup

To use the Pydantic AI framework, install it via pip:

pip install pydantic-ai        

Agent Execution Methods

The framework offers three primary execution methods to suit various use cases:

1. agent.run()

• Type: Coroutine
• Returns: Run Result containing the completed response.
• Use Case: Suitable for asynchronous workflows where you await the result.

result = await agent.run("Tell me about the first color patent in the world")
print(result.data)        

2. agent.run_sync()

• Type: Synchronous function
• Returns: Run Result containing the completed response.
• Internally: Executes agent.run() using asyncio.run().
• Use Case: Ideal for scenarios where synchronous execution is preferred.

result_sync = agent.run_sync("What is the most famous party place in Europe?")
print(result_sync.data)        

3. agent.run_stream()

• Type: Coroutine
• Returns: StreamedRunResult
• Details: Provides methods like .stream() for progressive output streaming.
• Use Case: Suitable for real-time applications or conversational AI where responses are generated progressively.

async def main():
    async with agent.run_stream("What is PayPal Mafia?") as response:
        async for chunk in response.stream():
            print(chunk)        

Text Streaming: Real-Time Interaction

Agents in the Pydantic AI framework support progressive text generation, making them ideal for interactive applications. This feature is especially useful for:

• Conversational AI
• Progressive content delivery
• Real-time dashboards

async def main():
    async with agent.run_stream("Define Pydantic AI") as response:
        async for chunk in response.stream():
            print(chunk)        

Tool Integration

The framework allows tools to be registered and invoked during agent execution. Tools provide modular, reusable functionality, enabling developers to extend the agent’s capabilities.

Tool Registration Methods

1. @agent.tool Decorator: Use for tools needing access to the agent context.
2. @agent.tool_plain Decorator: Use for tools that do not require access to the agent context.
3. tools Keyword Argument: Directly pass tool functions or Tool instances to the Agent.

#Example from pydantic-ai documentation

import random
from pydantic_ai import Agent, Tool

def roll_die() -> str:
    """Roll a six-sided die."""
    return str(random.randint(1, 6))

def get_player_name(ctx) -> str:
    """Get the player's name."""
    return ctx.deps

agent = Agent(
    'openai:gpt-4o',
    deps_type=str,
    tools=[
        Tool(roll_die, takes_ctx=False),
        Tool(get_player_name, takes_ctx=True),
    ],
)
result = agent.run_sync("Guess a number", deps="John")
print(result.data)        

Using Pydantic Data Models

Pydantic’s BaseModel enables structured validation of inputs and outputs for AI systems. This ensures that your AI workflows are robust and less prone to errors.

Example: Defining a Pydantic Model

from pydantic import BaseModel, Field

class AIModelConfig(BaseModel):
    model_name: str
    input_size: int
    output_size: int
    hidden_layers: list[int]
    activation_function: str = "relu"
    learning_rate: float = Field(0.001, gt=0, lt=1)        

Example: Running Agents with Structured Outputs

from pydantic_ai import Agent
from pydantic import BaseModel

class ModelResponse(BaseModel):
    city: str
    country: str
    popular_place: str
    point_of_contact: str

agent = Agent('openai:gpt-4o', result_type=ModelResponse)

result = agent.run_sync("Suggest a party place in Europe for June")
print(result.data)        

Output:

city='Barcelona'
country='Spain'
popular_place='Opium Barcelona'
point_of_contact='[email protected]'        

Common Issues

If you encounter dependency conflicts:

1. Check for incompatible versions in existing libraries.
2. Use virtual environments to isolate dependencies.
3. Run pip install --upgrade for related packages.

Overall, the Pydantic AI Agent Framework is a powerful, flexible, and developer-friendly tool for creating intelligent agents. By combining the validation strength of Pydantic with versatile execution models, it provides a good foundation for building AI-powered systems

Checkout Jupyter Notbook:

Notebook: https://github.com/05satyam/AI-ML/blob/main/agentic-frameworks-and-applications/PydanticAI-Agentic-Framework.ipynb

Checkout AI-ML Cookbook Repository:

Repository: https://github.com/05satyam/AI-ML

Reference

1. Pydantic AI: https://ai.pydantic.dev/

Thank you for reading! ?? ?? Connect with me: Satyam's LinkedIn , Satyam's Github

Also, visit my blogs where I share my work implementations and learning to write: Satyam's Blogs