importdatetimefromdataclassesimportdataclassfromtypingimportLiteralimportlogfirefrompydanticimportBaseModel,Fieldfromrich.promptimportPromptfrompydantic_aiimportAgent,ModelRetry,RunContextfrompydantic_ai.messagesimportModelMessagefrompydantic_ai.usageimportUsage,UsageLimits# 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configuredlogfire.configure(send_to_logfire='if-token-present')classFlightDetails(BaseModel):"""Details of the most suitable flight."""flight_number:strprice:intorigin:str=Field(description='Three-letter airport code')destination:str=Field(description='Three-letter airport code')date:datetime.dateclassNoFlightFound(BaseModel):"""When no valid flight is found."""@dataclassclassDeps:web_page_text:strreq_origin:strreq_destination:strreq_date:datetime.date# This agent is responsible for controlling the flow of the conversation.search_agent=Agent[Deps,FlightDetails|NoFlightFound]('openai:gpt-4o',output_type=FlightDetails|NoFlightFound,# type: ignoreretries=4,system_prompt=('Your job is to find the cheapest flight for the user on the given date. '),instrument=True,)# This agent is responsible for extracting flight details from web page text.extraction_agent=Agent('openai:gpt-4o',output_type=list[FlightDetails],system_prompt='Extract all the flight details from the given text.',)@search_agent.toolasyncdefextract_flights(ctx:RunContext[Deps])->list[FlightDetails]:"""Get details of all flights."""# we pass the usage to the search agent so requests within this agent are countedresult=awaitextraction_agent.run(ctx.deps.web_page_text,usage=ctx.usage)logfire.info('found {flight_count} flights',flight_count=len(result.output))returnresult.output@search_agent.output_validatorasyncdefvalidate_output(ctx:RunContext[Deps],output:FlightDetails|NoFlightFound)->FlightDetails|NoFlightFound:"""Procedural validation that the flight meets the constraints."""ifisinstance(output,NoFlightFound):returnoutputerrors:list[str]=[]ifoutput.origin!=ctx.deps.req_origin:errors.append(f'Flight should have origin {ctx.deps.req_origin}, not {output.origin}')ifoutput.destination!=ctx.deps.req_destination:errors.append(f'Flight should have destination {ctx.deps.req_destination}, not {output.destination}')ifoutput.date!=ctx.deps.req_date:errors.append(f'Flight should be on {ctx.deps.req_date}, not {output.date}')iferrors:raiseModelRetry('\n'.join(errors))else:returnoutputclassSeatPreference(BaseModel):row:int=Field(ge=1,le=30)seat:Literal['A','B','C','D','E','F']classFailed(BaseModel):"""Unable to extract a seat selection."""# This agent is responsible for extracting the user's seat selectionseat_preference_agent=Agent[None,SeatPreference|Failed]('openai:gpt-4o',output_type=SeatPreference|Failed,# type: ignoresystem_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. '),)# in reality this would be downloaded from a booking site,# potentially using another agent to navigate the siteflights_web_page="""1. Flight SFO-AK123- Price: $350- Origin: San Francisco International Airport (SFO)- Destination: Ted Stevens Anchorage International Airport (ANC)- Date: January 10, 20252. Flight SFO-AK456- Price: $370- Origin: San Francisco International Airport (SFO)- Destination: Fairbanks International Airport (FAI)- Date: January 10, 20253. Flight SFO-AK789- Price: $400- Origin: San Francisco International Airport (SFO)- Destination: Juneau International Airport (JNU)- Date: January 20, 20254. Flight NYC-LA101- Price: $250- Origin: San Francisco International Airport (SFO)- Destination: Ted Stevens Anchorage International Airport (ANC)- Date: January 10, 20255. Flight CHI-MIA202- Price: $200- Origin: Chicago O'Hare International Airport (ORD)- Destination: Miami International Airport (MIA)- Date: January 12, 20256. Flight BOS-SEA303- Price: $120- Origin: Boston Logan International Airport (BOS)- Destination: Ted Stevens Anchorage International Airport (ANC)- Date: January 12, 20257. Flight DFW-DEN404- Price: $150- Origin: Dallas/Fort Worth International Airport (DFW)- Destination: Denver International Airport (DEN)- Date: January 10, 20258. Flight ATL-HOU505- Price: $180- Origin: Hartsfield-Jackson Atlanta International Airport (ATL)- Destination: George Bush Intercontinental Airport (IAH)- Date: January 10, 2025"""# restrict how many requests this app can make to the LLMusage_limits=UsageLimits(request_limit=15)asyncdefmain():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=Noneusage:Usage=Usage()# run the agent until a satisfactory flight is foundwhileTrue:result=awaitsearch_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,)ifisinstance(result.output,NoFlightFound):print('No flight found')breakelse:flight=result.outputprint(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,)ifanswer=='buy':seat=awaitfind_seat(usage)awaitbuy_tickets(flight,seat)breakelse:message_history=result.all_messages(output_tool_return_content='Please suggest another flight')asyncdeffind_seat(usage:Usage)->SeatPreference:message_history:list[ModelMessage]|None=NonewhileTrue:answer=Prompt.ask('What seat would you like?')result=awaitseat_preference_agent.run(answer,message_history=message_history,usage=usage,usage_limits=usage_limits,)ifisinstance(result.output,SeatPreference):returnresult.outputelse:print('Could not understand seat preference. Please try again.')message_history=result.all_messages()asyncdefbuy_tickets(flight_details:FlightDetails,seat:SeatPreference):print(f'Purchasing flight {flight_details=!r}{seat=!r}...')if__name__=='__main__':importasyncioasyncio.run(main())