Pydantic-AI Usage

Pydantic-AI Usage#

Welcome to pydantic-ai, a Python proxy framework designed specifically for generative AI applications, aimed at helping developers build production-grade applications more easily. Developed by the Pydantic team, it inherits Pydantic's type safety and data validation capabilities and is optimized for large language models (LLMs).

This document will introduce the core modules and usage of pydantic-ai, including installation, basic concepts, agent creation, tool usage, and examples of advanced features.

Installation#

Basic Installation#

Install pydantic-ai using pip:

pip install pydantic-ai

Slim Installation#

If you only want to use specific models and avoid installing unnecessary dependencies, you can choose the slim version:

pip install pydantic-ai-slim

Optional Dependencies#

Install support for specific models as needed, for example:

OpenAI model support:

pip install pydantic-ai[openai]
Full model support (including OpenAI, Anthropic, Gemini, etc.):
pip install pydantic-ai[all]

Environment Configuration#

You need to configure the API key before use, for example:

import os os.environ["OPENAI_API_KEY"] = "your API key"

Core Modules and Concepts#

1. Agent#

The Agent is the core class of pydantic-ai, used to interact with LLMs. It encapsulates system prompts, tools, and structured outputs.

Main parameters:
- model: Specifies the LLM model to use (e.g., 'openai' or 'gemini-1.5-flash').
- system_prompt: The system prompt that defines the agent's behavior.
- result_type: The type of structured output, defined using a Pydantic model.

2. BaseModel (Pydantic Model)#

Inherits from Pydantic's BaseModel, used to define the data structures for input and output, ensuring type safety and validation.

3. RunContext#

Provides runtime context information that can be used to access dependencies or input data in tool functions.

4. Tools#

Defined using the @agent.tool decorator, allowing the agent to call custom functions to extend functionality.

Quick Start#

Creating a Simple Agent#

The following example demonstrates how to create an agent that asks a question and receives a structured response:

from pydantic import BaseModel
from pydantic_ai import Agent


# Define output structure
class Answer(BaseModel):
    text: str
    confidence: float


# Create agent
agent = Agent(
    model="openai:gpt-4", system_prompt="Provide concise and accurate answers.", result_type=Answer
)
result = agent.run_sync("What is the capital of China?")
print(result.data)
# Example output: Answer(text='Beijing', confidence=0.99)

Module Details and Usage#

1. Initializing the Agent#

The agent supports various models, here are common configurations:

from pydantic_ai import Agent

# Using the Gemini model
agent = Agent(
    model="gemini-1.5-flash", system_prompt="Answer questions in Chinese, keeping it concise.", retries=3
)
# Using the Ollama local model
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

ollama_model = OpenAIModel(
    model_name="llama3.2", provider=OpenAIProvider(base_url="http://localhost:11434/v1")
)
agent = Agent(ollama_model)

2. Defining Structured Output#

Use Pydantic models to define the output format, ensuring consistency:

from pydantic_ai import Agent
from pydantic import BaseModel, Field


class CityInfo(BaseModel):
    city: str = Field(description="City name")
    country: str = Field(description="Country name")


agent = Agent(model="openai:gpt-4", result_type=CityInfo)
result = agent.run_sync("Where was the 2012 Olympics held?")
print(result.data)  # Example output: CityInfo(city='London', country='UK')

3. Adding Tools#

Tools allow the agent to perform external operations, such as database queries:

from pydantic_ai import Agent, RunContext

agent = Agent(model="openai:gpt-4", system_prompt="Query the weather based on user input.")


@agent.tool
def get_weather(ctx: RunContext[str]) -> str:
    city = ctx.input_data
    return f"The weather in {city} is sunny, with a temperature of 25°C."


result = agent.run_sync("How is the weather in Beijing?")
print(result.data)  # Example output: 'The weather in Beijing is sunny, with a temperature of 25°C.'

4. Dependency Injection#

Specify the dependency type using deps_type to inject runtime data:

from dataclasses import dataclass


@dataclass
class WeatherDB:
    city: str


agent = Agent(
    model="openai:gpt-4", deps_type=WeatherDB, system_prompt="Provide weather information based on the database."
)
result = agent.run_sync("Today's weather", deps=WeatherDB(city="Shanghai"))
print(result.data)

5. Streaming Responses#

Supports streaming output, suitable for real-time applications:

async def stream_example():
    agent = Agent(model="openai:gpt-4", system_prompt="Answer questions step by step.")
    async for chunk in agent.run_stream("Describe the history of Python"):
        print(chunk.data, end="", flush=True)


import asyncio

asyncio.run(stream_example())

Advanced Features#

1. Retry Mechanism#

Configure the retries parameter to automatically retry when validation fails:

agent = Agent(model="gemini-1.5-flash", retries=3, result_type=CityInfo)
result = agent.run_sync("Where was the 2012 Olympics?")

2. Pydantic Logfire Integration#

For debugging and performance monitoring:

import logfire
from pydantic_ai import Agent

logfire.configure()
agent = Agent(model="openai:gpt-4")
logfire.instrument_pydantic()
# Log validation logs result = agent.run_sync("Test question")

3. Multi-Model Fallback#

Automatically switch to a backup model when one model fails:

from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.models.fallback import FallbackModel

openai_model = OpenAIModel("gpt-4o")
backup_model = OpenAIModel("gpt-3.5-turbo")
fallback = FallbackModel(openai_model, backup_model)
agent = Agent(fallback)
result = agent.run_sync("Question")

Example: Bank Support Agent#

Here is a complete example of a bank support agent:

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


@dataclass
class SupportDependencies:
    customer_id: int


class SupportResult(BaseModel):
    advice: str = Field(description="Advice for the customer")
    block_card: bool = Field(description="Whether to block the card")


agent = Agent(
    model="openai:gpt-4",
    deps_type=SupportDependencies,
    result_type=SupportResult,
    system_prompt="As a bank first-level support, analyze customer issues and provide advice.",
)


@agent.tool
def check_account(ctx: RunContext[SupportDependencies]) -> str:
    return f"The account of customer {ctx.deps.customer_id} is normal."


result = agent.run_sync(
    "What should I do if my card is lost?", deps=SupportDependencies(customer_id=12345)
)
print(result.data)
# Example output: SupportResult(advice='Please report your card immediately.', block_card=True)

Using pydantic-ai to Integrate DeepSeek and Ollama's Qwen Model#

pydantic-ai is a powerful Python framework that supports integration with various large language models (LLMs), including models provided by DeepSeek and Ollama. This manual will introduce how to configure and use DeepSeek's deepseek-chat model and Ollama's qwen model, along with an example to demonstrate its usage.

Prerequisites#

Before using, please ensure the following preparations are completed:

Install pydantic-ai
Install pydantic-ai in your Python environment:
```
pip install pydantic-ai
```
DeepSeek Configuration
- Obtain DeepSeek's API key: Visit DeepSeek's official website to register and generate an API key.
- Set environment variables (optional):
```
export DEEPSEEK_API_KEY='your-deepseek-api-key'
```
Ollama Configuration
- Install Ollama: Follow the instructions in the Ollama official documentation to download and install Ollama.
- Download the Qwen model (e.g., qwen2.5):
```
ollama pull qwen2.5
```
- Ensure the Ollama service is running locally:
```
ollama serve
```

Configuring the DeepSeek Model#

Here is an example of how to configure DeepSeek's deepseek-chat model using pydantic-ai:

from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.deepseek import DeepSeekProvider
from pydantic import BaseModel

# Define the return data structure
class BankQueryResult(BaseModel):
    action: str
    details: str

# Configure the DeepSeek model
deepseek_model = OpenAIModel(
    model_name='deepseek-chat',
    provider=DeepSeekProvider(
        api_key='your-deepseek-api-key',  # Replace with your API key
    )
)

# Create agent
agent = Agent(model=deepseek_model, result_type=BankQueryResult)

# Run query
result = agent.run_sync('I want to transfer 100 yuan to my friend, how to do it?')
print(result.data)
# Example output: {'action': 'Transfer', 'details': 'Log in to the banking app, select the transfer function, enter your friend's account and the amount of 100 yuan, and confirm.'}

Configuring Ollama's Qwen Model#

Here is an example of how to configure Ollama's qwen2.5 model using pydantic-ai:

from pydantic_ai import Agent
from pydantic_ai.models.ollama import OllamaModel
from pydantic import BaseModel

# Define the return data structure
class BankQueryResult(BaseModel):
    action: str
    details: str

# Configure Ollama's Qwen model
ollama_model = OllamaModel(
    model_name='qwen2.5',  # Ensure this model has been downloaded via ollama pull
    base_url='http://localhost:11434/v1',  # Default Ollama local service address
)

# Create agent
agent = Agent(model=ollama_model, result_type=BankQueryResult)

# Run query
result = agent.run_sync('How can I check my account balance?')
print(result.data)
# Example output: {'action': 'Check balance', 'details': 'Log in to the banking app, go to the account page, and select the balance inquiry option.'}

Example: Handling Bank Queries with DeepSeek and Ollama's Qwen Model#

Here is an example that combines DeepSeek and Ollama's Qwen model to handle complex bank-related queries:

from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.deepseek import DeepSeekProvider
from pydantic_ai.models.ollama import OllamaModel
from pydantic import BaseModel

# Define the return data structure
class ComplexBankQuery(BaseModel):
    question_type: str
    answer: str
    confidence: float

# Configure the DeepSeek model
deepseek_model = OpenAIModel(
    model_name='deepseek-chat',
    provider=DeepSeekProvider(
        api_key='your-deepseek-api-key',  # Replace with your API key
    )
)

# Configure Ollama's Qwen model
qwen_model = OllamaModel(
    model_name='qwen2.5',
    base_url='http://localhost:11434/v1',
)

# Create two agents
deepseek_agent = Agent(model=deepseek_model, result_type=ComplexBankQuery)
qwen_agent = Agent(model=qwen_model, result_type=ComplexBankQuery)

# Test DeepSeek agent
deepseek_result = deepseek_agent.run_sync('How to apply for a small loan?')
print("DeepSeek Answer:", deepseek_result.data)
# Example output: {'question_type': 'Loan application', 'answer': 'Contact the bank, submit the loan application form and identification, and wait for approval.', 'confidence': 0.95}

# Test Qwen agent
qwen_result = qwen_agent.run_sync('How to apply for a small loan?')
print("Qwen Answer:", qwen_result.data)
# Example output: {'question_type': 'Loan application', 'answer': 'Log in to the banking app, select the small loan option, fill in the information and submit.', 'confidence': 0.92}

Example Notes#

DeepSeek API Key
Ensure your DeepSeek API key is valid; otherwise, it will lead to authentication failures. You can pass it via environment variables or directly in the code.
Ollama Service Running
When using Ollama, make sure the local service is running and the specified model (like qwen2.5) has been downloaded.
Performance Differences
- DeepSeek's deepseek-chat is a cloud model suitable for scenarios requiring high-performance inference but relies on the network.
- Ollama's qwen2.5 is a local model suitable for offline use, but performance is limited by local hardware.
Debugging and Logging
pydantic-ai supports integration with Pydantic Logfire, enabling logs to monitor model behavior:
```
from pydantic_ai import Logfire
Logfire.configure()
```

Notes#

API Key: Ensure the environment variable is correctly configured; otherwise, authentication errors will be thrown.
Model Compatibility: Check if the selected model supports specified features (like streaming output).
Performance Optimization: For high-throughput applications, it is recommended to use local models (like Ollama).

Summary#

pydantic-ai provides a powerful toolkit that enables developers to build complex AI agents in a type-safe manner. By combining agents, tools, and structured outputs, you can quickly develop production-grade applications ranging from simple Q&A to complex business logic. For more information, please refer to the official documentation.