6af49ef068
Addresses general improvements to all toolkits including changing ruff from python 3.9 to python 3.10 which is the reason for the removal of Optional[] among others. Also, turns out that our `make install` for toolkits wasn't correctly checking for whether poetry was installed (&> /dev/null syntax isn't supported by our check-toolkits GitHub action, so we were installing poetry twice. I replaced with the more portable >/dev/null 2>&1) Question: Should we also change ruff to py310 for the `arcade/` package in a later PR? ------------------- CU-86b4gzyp6
110 lines
4.4 KiB
Python
110 lines
4.4 KiB
Python
from typing import Annotated, Any
|
|
|
|
from arcade.sdk import ToolContext, tool
|
|
|
|
from arcade_search.enums import GoogleFlightsMaxStops, GoogleFlightsSortBy, GoogleFlightsTravelClass
|
|
from arcade_search.utils import call_serpapi, parse_flight_results, prepare_params
|
|
|
|
|
|
@tool(requires_secrets=["SERP_API_KEY"])
|
|
async def search_roundtrip_flights(
|
|
context: ToolContext,
|
|
departure_airport_code: Annotated[
|
|
str, "The departure airport code. An uppercase 3-letter code"
|
|
],
|
|
arrival_airport_code: Annotated[str, "The arrival airport code. An uppercase 3-letter code"],
|
|
outbound_date: Annotated[str, "Flight outbound date in YYYY-MM-DD format"],
|
|
return_date: Annotated[str | None, "Flight return date in YYYY-MM-DD format"],
|
|
currency_code: Annotated[
|
|
str | None, "Currency of the returned prices. Defaults to 'USD'"
|
|
] = "USD",
|
|
travel_class: Annotated[
|
|
GoogleFlightsTravelClass,
|
|
"Travel class of the flight. Defaults to 'ECONOMY'",
|
|
] = GoogleFlightsTravelClass.ECONOMY,
|
|
num_adults: Annotated[int | None, "Number of adult passengers. Defaults to 1"] = 1,
|
|
num_children: Annotated[int | None, "Number of child passengers. Defaults to 0"] = 0,
|
|
max_stops: Annotated[
|
|
GoogleFlightsMaxStops,
|
|
"Maximum number of stops (layovers) for the flight. Defaults to any number of stops",
|
|
] = GoogleFlightsMaxStops.ANY,
|
|
sort_by: Annotated[
|
|
GoogleFlightsSortBy,
|
|
"The sorting order of the results. Defaults to TOP_FLIGHTS.",
|
|
] = GoogleFlightsSortBy.TOP_FLIGHTS,
|
|
) -> Annotated[dict[str, Any], "Flight search results from the Google Flights API"]:
|
|
"""Retrieve flight search results using Google Flights"""
|
|
# Prepare the request
|
|
params = prepare_params(
|
|
"google_flights",
|
|
departure_id=departure_airport_code,
|
|
arrival_id=arrival_airport_code,
|
|
outbound_date=outbound_date,
|
|
return_date=return_date,
|
|
currency=currency_code,
|
|
travel_class=travel_class.to_api_value(),
|
|
adults=num_adults,
|
|
children=num_children,
|
|
stops=max_stops.to_api_value(),
|
|
sort_by=sort_by.to_api_value(),
|
|
deep_search=True, # Same search depth of the Google Flights page in the browser
|
|
)
|
|
|
|
# Execute the request
|
|
results = call_serpapi(context, params)
|
|
|
|
# Parse the results
|
|
flights = parse_flight_results(results)
|
|
|
|
return flights
|
|
|
|
|
|
@tool(requires_secrets=["SERP_API_KEY"])
|
|
async def search_one_way_flights(
|
|
context: ToolContext,
|
|
departure_airport_code: Annotated[
|
|
str, "The departure airport code. An uppercase 3-letter code"
|
|
],
|
|
arrival_airport_code: Annotated[str, "The arrival airport code. An uppercase 3-letter code"],
|
|
outbound_date: Annotated[str, "Flight departure date in YYYY-MM-DD format"],
|
|
currency_code: Annotated[
|
|
str | None, "Currency of the returned prices. Defaults to 'USD'"
|
|
] = "USD",
|
|
travel_class: Annotated[
|
|
GoogleFlightsTravelClass,
|
|
"Travel class of the flight. Defaults to 'ECONOMY'",
|
|
] = GoogleFlightsTravelClass.ECONOMY,
|
|
num_adults: Annotated[int | None, "Number of adult passengers. Defaults to 1"] = 1,
|
|
num_children: Annotated[int | None, "Number of child passengers. Defaults to 0"] = 0,
|
|
max_stops: Annotated[
|
|
GoogleFlightsMaxStops,
|
|
"Maximum number of stops (layovers) for the flight. Defaults to any number of stops",
|
|
] = GoogleFlightsMaxStops.ANY,
|
|
sort_by: Annotated[
|
|
GoogleFlightsSortBy,
|
|
"The sorting order of the results. Defaults to TOP_FLIGHTS.",
|
|
] = GoogleFlightsSortBy.TOP_FLIGHTS,
|
|
) -> Annotated[dict[str, Any], "Flight search results from the Google Flights API"]:
|
|
"""Retrieve flight search results for a one-way flight using Google Flights"""
|
|
params = prepare_params(
|
|
"google_flights",
|
|
departure_id=departure_airport_code,
|
|
arrival_id=arrival_airport_code,
|
|
outbound_date=outbound_date,
|
|
currency=currency_code,
|
|
travel_class=travel_class.to_api_value(),
|
|
adults=num_adults,
|
|
children=num_children,
|
|
stops=max_stops.to_api_value(),
|
|
sort_by=sort_by.to_api_value(),
|
|
type=2, # indicates one-way
|
|
deep_search=True, # Same search depth as the Google Flights page in the browser
|
|
)
|
|
|
|
# Execute the request
|
|
results = call_serpapi(context, params)
|
|
|
|
# Parse the results
|
|
flights = parse_flight_results(results)
|
|
|
|
return flights
|