Get Started with Pydantic AI

pip install pydantic-ai openai 

CEREBRAS_API_KEY=your-cerebras-api-key-here

export CEREBRAS_API_KEY=your-cerebras-api-key-here

import os
import asyncio
from openai import AsyncOpenAI
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.cerebras import CerebrasProvider

# Create an OpenAI client with the integration header
openai_client = AsyncOpenAI(
    base_url='https://api.cerebras.ai/v1',
    api_key=os.getenv('CEREBRAS_API_KEY'),
    default_headers={'X-Cerebras-3rd-Party-Integration': 'pydantic-ai'}
)

# Create a Cerebras provider with the custom client
provider = CerebrasProvider(openai_client=openai_client)

# Create a Cerebras model instance
model = OpenAIChatModel('gpt-oss-120b', provider=provider)

# Create an agent with the Cerebras model
agent = Agent(
    model=model,
    system_prompt='You are a helpful assistant that provides clear, concise answers.'
)

async def main():
    # Run the agent asynchronously
    result = await agent.run('What is the capital of France?')
    print(result.output)

asyncio.run(main())

import os
import asyncio
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel

# Define your output structure with validation
class CityInfo(BaseModel):
    city: str = Field(description="The name of the city")
    country: str = Field(description="The country where the city is located")
    population: int = Field(description="Approximate population")
    famous_for: list[str] = Field(description="List of things the city is famous for")

# Create the Cerebras model
model = OpenAIChatModel(
    'gpt-oss-120b',
    provider='cerebras'
)

# Create an agent with structured output
agent = Agent(
    model=model,
    output_type=CityInfo,  # Changed from result_type to output_type
    system_prompt='Extract city information from the user query. Provide accurate, factual data.'
)

# Get structured, validated output
async def main():
    result = await agent.run('Tell me about Paris')
    print(result.output)

import asyncio
asyncio.run(main())

import os
import asyncio
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel

# Create the Cerebras model
model = OpenAIChatModel(
    'gpt-oss-120b',
    provider='cerebras'
)

agent = Agent(
    model=model,
    system_prompt='You are a helpful assistant.'
)

async def main():
    results = await asyncio.gather(
        agent.run('What is 2+2?'),
        agent.run('What is the capital of Spain?'),
        agent.run('Who wrote Romeo and Juliet?')
    )
    for i, result in enumerate(results, 1):
        print(f"Answer {i}: {result.output}")

asyncio.run(main())

import os
import asyncio
from datetime import datetime
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIChatModel

model = OpenAIChatModel(
    'gpt-oss-120b',
    provider='cerebras'
)

agent = Agent(
    model=model,
    system_prompt='You are a helpful assistant with access to tools.'
)

@agent.tool
def get_current_time(ctx: RunContext[None]) -> str:
    """Get the current time in a readable format."""
    return datetime.now().strftime('%I:%M %p')

@agent.tool
def calculate(ctx: RunContext[None], expression: str) -> str:
    """Safely evaluate a mathematical expression.
    
    Args:
        expression: A mathematical expression like '15 * 23'
    """
    try:
        # In production, use a safe eval library like numexpr
        result = eval(expression, {'__builtins__': {}}, {})
        return str(result)
    except Exception as e:
        return f"Error: {str(e)}"

# The agent can now use these tools automatically
import asyncio

async def main():
    result = await agent.run('What time is it and what is 15 * 23?')
    print(result.output)

asyncio.run(main())

import os
import asyncio
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel

model = OpenAIChatModel(
    'gpt-oss-120b',
    provider='cerebras'
)

agent = Agent(
    model=model,
    system_prompt='You are a creative writing assistant.'
)

# Stream the response token by token
async def main():
    async with agent.run_stream('Write a short poem about AI') as response:
        async for chunk in response.stream_text():
            print(chunk, end='', flush=True)
        print()  # New line at the end

asyncio.run(main())