Pydantic AI

from obiguard import AsyncObiguard
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

# Set up Obiguard client with tracing and metadata
obiguard_client = AsyncObiguard(
    obiguard_api_key="vk-obg***",  # Your Obiguard virtual key
)

# Define structured output using Pydantic
class F1GrandPrix(BaseModel):
    gp_name: str = Field(description="Grand Prix name, e.g. `Emilia Romagna Grand Prix`")
    year: int = Field(description="The year of the Grand Prix")
    constructor_winner: str = Field(description="The winning constructor of the Grand Prix")
    podium: list[str] = Field(description="Names of the podium drivers (1st, 2nd, 3rd)")

# Create the agent with structured output type
f1_gp_agent = Agent[None, F1GrandPrix](
    model=OpenAIModel(
        model_name="gpt-4o",
        provider=OpenAIProvider(openai_client=obiguard_client),
    ),
    output_type=F1GrandPrix,
    system_prompt="Assist the user by providing data about the specified Formula 1 Grand Prix"
)

# Run the agent
async def main():
    result = await f1_gp_agent.run("Las Vegas 2023")
    print(result.output)

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

gp_name='Las Vegas Grand Prix'
year=2023
constructor_winner='Red Bull Racing'
podium=['Max Verstappen', 'Charles Leclerc', 'Sergio Pérez']

result = f1_gp_agent.run_sync("Las Vegas 2023")
print(result.output)

from obiguard import AsyncObiguard
from pydantic_ai import Agent, ImageUrl
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

# Set up Obiguard client
obiguard_client = AsyncObiguard(
    obiguard_api_key="sk-obg***",  # Your Obiguard API key
)

# Create a vision-capable agent
vision_agent = Agent(
    model=OpenAIModel(
        model_name="gpt-4o",  # Vision-capable model
        provider=OpenAIProvider(openai_client=obiguard_client),
    ),
    system_prompt="Analyze images and provide detailed descriptions."
)

# Process an image
result = vision_agent.run_sync([
    'What company is this logo from?',
    ImageUrl(url='https://iili.io/3Hs4FMg.png'),
])
print(result.output)

import random
from obiguard import AsyncObiguard
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

# Set up Obiguard client
obiguard_client = AsyncObiguard(
    obiguard_api_key="sk-obg***",  # Your Obiguard API key
)

# Create an agent with dependency injection (player name)
dice_agent = Agent(
    model=OpenAIModel(
        model_name="gpt-4o",
        provider=OpenAIProvider(openai_client=obiguard_client),
    ),
    deps_type=str,  # Dependency type (player name as string)
    system_prompt=(
        "You're a dice game host. Roll the die and see if it matches "
        "the user's guess. If so, tell them they're a winner. "
        "Use the player's name in your response."
    ),
)

# Define a plain tool (no context needed)
@dice_agent.tool_plain
def roll_die() -> str:
    """Roll a six-sided die and return the result."""
    return str(random.randint(1, 6))

# Define a tool that uses the dependency
@dice_agent.tool
def get_player_name(ctx: RunContext[str]) -> str:
    """Get the player's name."""
    return ctx.deps

# Run the agent
dice_result = dice_agent.run_sync('My guess is 4', deps='Anne')
print(dice_result.output)

import datetime
from dataclasses import dataclass
from typing import Literal

from pydantic import BaseModel, Field
from rich.prompt import Prompt

from pydantic_ai import Agent, ModelRetry, RunContext
from pydantic_ai.messages import ModelMessage
from pydantic_ai.usage import Usage, UsageLimits
from obiguard import AsyncObiguard

# Set up Obiguard clients with shared trace ID for connected tracing
obiguard_client = AsyncObiguard(
    obiguard_api_key="sk-obg***",  # Your Obiguard API key
)

# Define structured output types
class FlightDetails(BaseModel):
    """Details of the most suitable flight."""
    flight_number: str
    price: int
    origin: str = Field(description='Three-letter airport code')
    destination: str = Field(description='Three-letter airport code')
    date: datetime.date

class NoFlightFound(BaseModel):
    """When no valid flight is found."""

class SeatPreference(BaseModel):
    row: int = Field(ge=1, le=30)
    seat: Literal['A', 'B', 'C', 'D', 'E', 'F']

class Failed(BaseModel):
    """Unable to extract a seat selection."""

# Dependencies for flight search
@dataclass
class Deps:
    web_page_text: str
    req_origin: str
    req_destination: str
    req_date: datetime.date

# This agent is responsible for controlling the flow of the conversation
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

search_agent = Agent[Deps, FlightDetails | NoFlightFound](
    model=OpenAIModel(
        model_name="gpt-4o",
        provider=OpenAIProvider(openai_client=obiguard_client),
    ),
    output_type=FlightDetails | NoFlightFound,  # type: ignore
    retries=4,
    system_prompt=(
        'Your job is to find the cheapest flight for the user on the given date. '
    ),
    instrument=True,  # Enable instrumentation for better tracing
)

# This agent is responsible for extracting flight details from web page text
extraction_agent = Agent(
    model=OpenAIModel(
        model_name="gpt-4o",
        provider=OpenAIProvider(openai_client=obiguard_client),
    ),
    output_type=list[FlightDetails],
    system_prompt='Extract all the flight details from the given text.',
)

# This agent is responsible for extracting the user's seat selection
seat_preference_agent = Agent[None, SeatPreference | Failed](
    model=OpenAIModel(
        model_name="gpt-4o",
        provider=OpenAIProvider(openai_client=obiguard_client),
    ),
    output_type=SeatPreference | Failed,  # type: ignore
    system_prompt=(
        "Extract the user's seat preference. "
        'Seats A and F are window seats. '
        'Row 1 is the front row and has extra leg room. '
        'Rows 14, and 20 also have extra leg room. '
    ),
)

@search_agent.tool
async def extract_flights(ctx: RunContext[Deps]) -> list[FlightDetails]:
    """Get details of all flights."""
    # Pass the usage to track nested agent calls
    result = await extraction_agent.run(ctx.deps.web_page_text, usage=ctx.usage)
    return result.output

@search_agent.output_validator
async def validate_output(
    ctx: RunContext[Deps], output: FlightDetails | NoFlightFound
) -> FlightDetails | NoFlightFound:
    """Procedural validation that the flight meets the constraints."""
    if isinstance(output, NoFlightFound):
        return output

    errors: list[str] = []
    if output.origin != ctx.deps.req_origin:
        errors.append(
            f'Flight should have origin {ctx.deps.req_origin}, not {output.origin}'
        )
    if output.destination != ctx.deps.req_destination:
        errors.append(
            f'Flight should have destination {ctx.deps.req_destination}, not {output.destination}'
        )
    if output.date != ctx.deps.req_date:
        errors.append(f'Flight should be on {ctx.deps.req_date}, not {output.date}')

    if errors:
        raise ModelRetry('\n'.join(errors))
    else:
        return output

# Sample flight data (in a real application, this would be from a web scraper)
flights_web_page = """
1. Flight SFO-AK123
- Price: $350
- Origin: San Francisco International Airport (SFO)
- Destination: Ted Stevens Anchorage International Airport (ANC)
- Date: January 10, 2025

2. Flight SFO-AK456
- Price: $370
- Origin: San Francisco International Airport (SFO)
- Destination: Fairbanks International Airport (FAI)
- Date: January 10, 2025

... more flights ...
"""

# Main application flow
async def main():
    # Restrict how many requests this app can make to the LLM
    usage_limits = UsageLimits(request_limit=15)

    deps = Deps(
        web_page_text=flights_web_page,
        req_origin='SFO',
        req_destination='ANC',
        req_date=datetime.date(2025, 1, 10),
    )
    message_history: list[ModelMessage] | None = None
    usage: Usage = Usage()

    # Run the agent until a satisfactory flight is found
    while True:
        result = await search_agent.run(
            f'Find me a flight from {deps.req_origin} to {deps.req_destination} on {deps.req_date}',
            deps=deps,
            usage=usage,
            message_history=message_history,
            usage_limits=usage_limits,
        )
        if isinstance(result.output, NoFlightFound):
            print('No flight found')
            break
        else:
            flight = result.output
            print(f'Flight found: {flight}')
            answer = Prompt.ask(
                'Do you want to buy this flight, or keep searching? (buy/*search)',
                choices=['buy', 'search', ''],
                show_choices=False,
            )
            if answer == 'buy':
                seat = await find_seat(usage, usage_limits)
                await buy_tickets(flight, seat)
                break
            else:
                message_history = result.all_messages(
                    output_tool_return_content='Please suggest another flight'
                )

async def find_seat(usage: Usage, usage_limits: UsageLimits) -> SeatPreference:
    message_history: list[ModelMessage] | None = None
    while True:
        answer = Prompt.ask('What seat would you like?')
        result = await seat_preference_agent.run(
            answer,
            message_history=message_history,
            usage=usage,
            usage_limits=usage_limits,
        )
        if isinstance(result.output, SeatPreference):
            return result.output
        else:
            print('Could not understand seat preference. Please try again.')
            message_history = result.all_messages()

async def buy_tickets(flight_details: FlightDetails, seat: SeatPreference):
    print(f'Purchasing flight {flight_details=!r} {seat=!r}...')