[docs]classFlightSearchSchema(BaseModel):"""Schema for the AmadeusFlightSearch tool."""originLocationCode:str=Field(description=(" The three letter International Air Transport "" Association (IATA) Location Identifier for the "" search's origin airport. "))destinationLocationCode:str=Field(description=(" The three letter International Air Transport "" Association (IATA) Location Identifier for the "" search's destination airport. "))departureDateTimeEarliest:str=Field(description=(" The earliest departure datetime from the origin airport "" for the flight search in the following format: "' "YYYY-MM-DDTHH:MM:SS", where "T" separates the date and time '' components. For example: "2023-06-09T10:30:00" represents '" June 9th, 2023, at 10:30 AM. "))departureDateTimeLatest:str=Field(description=(" The latest departure datetime from the origin airport "" for the flight search in the following format: "' "YYYY-MM-DDTHH:MM:SS", where "T" separates the date and time '' components. For example: "2023-06-09T10:30:00" represents '" June 9th, 2023, at 10:30 AM. "))page_number:int=Field(default=1,description="The specific page number of flight results to retrieve",)
[docs]classAmadeusFlightSearch(AmadeusBaseTool):# type: ignore[override, override]"""Tool for searching for a single flight between two airports."""name:str="single_flight_search"description:str=(" Use this tool to search for a single flight between the origin and "" destination airports at a departure between an earliest and "" latest datetime. ")args_schema:Type[FlightSearchSchema]=FlightSearchSchemadef_run(self,originLocationCode:str,destinationLocationCode:str,departureDateTimeEarliest:str,departureDateTimeLatest:str,page_number:int=1,run_manager:Optional[CallbackManagerForToolRun]=None,)->list:try:fromamadeusimportResponseErrorexceptImportErrorase:raiseImportError("Unable to import amadeus, please install with `pip install amadeus`.")fromeRESULTS_PER_PAGE=10# Authenticate and retrieve a clientclient=self.client# Check that earliest and latest dates are in the same dayearliestDeparture=dt.strptime(departureDateTimeEarliest,"%Y-%m-%dT%H:%M:%S")latestDeparture=dt.strptime(departureDateTimeLatest,"%Y-%m-%dT%H:%M:%S")ifearliestDeparture.date()!=latestDeparture.date():logger.error(" Error: Earliest and latest departure dates need to be the "" same date. If you're trying to search for round-trip "" flights, call this function for the outbound flight first, "" and then call again for the return flight. ")return[None]# Collect all results from the Amadeus Flight Offers Search APIresponse=Nonetry:response=client.shopping.flight_offers_search.get(originLocationCode=originLocationCode,destinationLocationCode=destinationLocationCode,departureDate=latestDeparture.strftime("%Y-%m-%d"),adults=1,)exceptResponseErroraserror:print(error)# noqa: T201# Generate output dictionaryoutput=[]ifresponseisnotNone:forofferinresponse.data:itinerary:Dict={}itinerary["price"]={}itinerary["price"]["total"]=offer["price"]["total"]currency=offer["price"]["currency"]currency=response.result["dictionaries"]["currencies"][currency]itinerary["price"]["currency"]={}itinerary["price"]["currency"]=currencysegments=[]forsegmentinoffer["itineraries"][0]["segments"]:flight={}flight["departure"]=segment["departure"]flight["arrival"]=segment["arrival"]flight["flightNumber"]=segment["number"]carrier=segment["carrierCode"]carrier=response.result["dictionaries"]["carriers"][carrier]flight["carrier"]=carriersegments.append(flight)itinerary["segments"]=[]itinerary["segments"]=segmentsoutput.append(itinerary)# Filter out flights after latest departure timeforindex,offerinenumerate(output):offerDeparture=dt.strptime(offer["segments"][0]["departure"]["at"],"%Y-%m-%dT%H:%M:%S")ifofferDeparture>latestDeparture:output.pop(index)# Return the paginated resultsstartIndex=(page_number-1)*RESULTS_PER_PAGEendIndex=startIndex+RESULTS_PER_PAGEreturnoutput[startIndex:endIndex]