multi‑agent flight booking system powered by Pydantic AI and Hugging Face Inference Providers

flight_booking.py

import asyncio
import datetime
from typing import Literal

from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.models.huggingface import HuggingFaceModel
from pydantic_ai.providers.huggingface import HuggingFaceProvider
from rich.console import Console
from rich.prompt import Confirm, Prompt


# Pydantic automatically validates:
#    - Field types (str, int, date)
#    - Field constraints (airport codes, price ranges)
#    - Required vs optional fields
class FlightDetails(BaseModel):
    """
    Structured flight information that AI will return.
    """
    flight_number: str = Field(description="Flight identifier like 'UA123'")
    airline: str = Field(description="Airline name like 'United Airlines'")
    price: int = Field(ge=50, le=5000, description="Price in USD")
    origin: str = Field(description='Three-letter airport code like SFO')
    destination: str = Field(description='Three-letter airport code like JFK')
    departure_time: str = Field(description="Time like '08:00 AM'")
    date: datetime.date = Field(description="Flight date")


# The Literal type restricts seat letters to valid airplane seats.
# Field constraints ensure realistic row numbers.
class SeatPreference(BaseModel):
    """
    Seat selection with automatic validation.
    """
    row: int = Field(ge=1, le=30, description="Seat row number")
    seat: Literal['A', 'B', 'C', 'D', 'E', 'F'] = Field(description="Seat letter")


class NoFlightFound(BaseModel):
    """
    Returned when no suitable flight is found.
    """
    reason: str = Field(description="Why no flight was found")

model = HuggingFaceModel(
    'moonshotai/Kimi-K2-Instruct',
    provider=HuggingFaceProvider(provider_name='together')
)


flight_agent = Agent(
    model,
    output_type=FlightDetails | NoFlightFound,  
    system_prompt="""
    You are a flight search expert. Extract flight information from text and 
    find suitable flights that match the user's criteria.
    
    If this is a follow-up search (you see previous conversation), try to suggest a different flight option 
    than what was previously offered. Consider factors like price, timing, airline, and route.
    
    If this is the first search, find the cheapest option.
    
    If you find a suitable flight, return FlightDetails with: flight number, airline, price, origin, destination, time, and date.
    If no suitable flight matches the criteria, return NoFlightFound with the reason.
    """
)

seat_agent = Agent(
    model,
    output_type=SeatPreference,
    system_prompt="""
    You are a seat selection expert. Understand user preferences and select appropriate seats.
    
    Seat guide:
    - A and F seats: Window seats (great views)
    - C and D seats: Aisle seats (easy access)
    - B and E seats: Middle seats (usually avoided)
    """
)

FLIGHT_DATA = """
🛫 AVAILABLE FLIGHTS - Multiple Routes - October 10, 2025

=== SAN FRANCISCO (SFO) TO ANCHORAGE (ANC) ===
1. Flight SFO-AK123 - Alaska Airlines
   Price: $350 | Departure: 08:00 AM
   Aircraft: Boeing 737-800

2. Flight SFO-AK456 - Alaska Airlines  
   Price: $380 | Departure: 02:30 PM
   Aircraft: Airbus A320

=== CROSS-COUNTRY FLIGHTS ===
3. Flight NYC-LA101 - American Airlines
   Price: $280 | Departure: 09:15 AM
   Route: SFO → ANC (via connection)

4. Flight BOS-SEA303 - JetBlue Airways
   Price: $120 | Departure: 06:00 AM  
   Route: BOS → ANC (via Seattle)
   Note: Budget option with layover

=== PREMIUM OPTIONS ===
5. Flight SFO-AK789 - United Airlines
   Price: $450 | Departure: 11:00 AM
   Aircraft: Boeing 777 | First Class Available
6. Flight DFW-ANC404 - Delta Airlines
   Price: $320 | Departure: 03:45 PM
   Route: SFO → ANC (direct)
"""

REQ_ORIGIN = 'SFO'
REQ_DESTINATION = 'ANC'
REQ_DATE = datetime.date(2025, 10, 10)

console = Console()


def show_welcome():
    """Display welcome message and flight search info."""
    console.print(f"🛫 Flight Booking: {REQ_ORIGIN} → {REQ_DESTINATION} on {REQ_DATE}")

def display_flight_details(flight: FlightDetails):
    """Display flight details simply."""
    console.print(f"✅ Found: {flight.airline} {flight.flight_number} - ${flight.price} {flight.departure_time}")

async def search_with_progress(search_query: str, message_history=None):
    """Search for flights with minimal progress indication."""
    console.print("🔍 Searching flights...")
    result = await flight_agent.run(search_query, message_history=message_history)
    return result

async def complete_booking(flight: FlightDetails, seat: SeatPreference):
    """Complete the booking process."""
    await asyncio.sleep(0.3)
    console.print(f"🎉 Booked! {flight.airline} {flight.flight_number}, Seat {seat.row}{seat.seat} - ${flight.price}")

async def find_seat() -> SeatPreference:
    """Interactive seat selection."""
    console.print("🪑 Seat Selection (A,F=window | C,D=aisle | B,E=middle)")
    
    message_history = None
    attempt = 0

    while True:
        attempt += 1
        answer = Prompt.ask("Seat preference")
        
        if not answer.strip():
            continue
        
        result = await seat_agent.run(answer, message_history=message_history)

        if isinstance(result.output, SeatPreference):
            seat = result.output
            seat_type = "Window" if seat.seat in ['A', 'F'] else "Aisle" if seat.seat in ['C', 'D'] else "Middle"
            console.print(f"✅ Selected: {seat.row}{seat.seat} ({seat_type})")
            return seat
        else:
            message_history = result.all_messages()
            if attempt >= 3:
                return SeatPreference(row=12, seat='C')
            
async def main():
    show_welcome()
    
    message_history = None
    
    while True:
        result = await search_with_progress(
            f'Find me a flight from {REQ_ORIGIN} to {REQ_DESTINATION} on {REQ_DATE}.\n\nAvailable flights:\n{FLIGHT_DATA}',
            message_history=message_history
        )
        
        if not isinstance(result.output, NoFlightFound):
            flight = result.output
            display_flight_details(flight)
            
            answer = Prompt.ask("Book this flight?", choices=["y", "n", "search"], default="y")
            
            if answer == "y":
                seat = await find_seat()
                await complete_booking(flight, seat)
                break
            elif answer == "search":
                console.print("🔄 Searching again...")
                message_history = result.all_messages()
            else:
                console.print("👋 Thanks!")
                break
        else:
            console.print(f"❌ {result.output.reason}")
            if not Confirm.ask("Try again?", default=False):
                break
            message_history = None




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