import contextlib import re from datetime import datetime from typing import Any, Optional, cast from zoneinfo import ZoneInfo from arcade.sdk import ToolContext from arcade.sdk.errors import ToolExecutionError from serpapi import Client as SerpClient from arcade_search.constants import ( DEFAULT_GOOGLE_MAPS_COUNTRY, DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, DEFAULT_GOOGLE_MAPS_LANGUAGE, DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, ) from arcade_search.enums import GoogleMapsDistanceUnit, GoogleMapsTravelMode from arcade_search.exceptions import CountryNotFoundError, LanguageNotFoundError from arcade_search.google_data import COUNTRY_CODES, LANGUAGE_CODES # ------------------------------------------------------------------------------------------------ # General SerpAPI utils # ------------------------------------------------------------------------------------------------ def prepare_params(engine: str, **kwargs: Any) -> dict[str, Any]: """ Prepares a parameters dictionary for the SerpAPI call. Parameters: engine: The engine name (e.g., "google", "google_finance"). kwargs: Any additional parameters to include. Returns: A dictionary containing the base parameters plus any extras, excluding any parameters whose value is None. """ params = {"engine": engine} params.update({k: v for k, v in kwargs.items() if v is not None}) return params def call_serpapi(context: ToolContext, params: dict) -> dict: """ Execute a search query using the SerpAPI client and return the results as a dictionary. Args: context: The tool context containing required secrets. params: A dictionary of parameters for the SerpAPI search. Returns: The search results as a dictionary. """ api_key = context.get_secret("SERP_API_KEY") client = SerpClient(api_key=api_key) try: search = client.search(params) return cast(dict[str, Any], search.as_dict()) except Exception as e: # SerpAPI error messages sometimes contain the API key, so we need to sanitize it sanitized_e = re.sub(r"(api_key=)[^ &]+", r"\1***", str(e)) raise ToolExecutionError( message="Failed to fetch search results", developer_message=sanitized_e, ) # ------------------------------------------------------------------------------------------------ # Google Maps utils # ------------------------------------------------------------------------------------------------ def get_google_maps_directions( context: ToolContext, origin_address: Optional[str] = None, destination_address: Optional[str] = None, origin_latitude: Optional[str] = None, origin_longitude: Optional[str] = None, destination_latitude: Optional[str] = None, destination_longitude: Optional[str] = None, language: str = DEFAULT_GOOGLE_MAPS_LANGUAGE, country: Optional[str] = DEFAULT_GOOGLE_MAPS_COUNTRY, distance_unit: GoogleMapsDistanceUnit = DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, travel_mode: GoogleMapsTravelMode = DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, ) -> list[dict[str, Any]]: """Get directions from Google Maps. Provide either all(origin_address, destination_address) or all(origin_latitude, origin_longitude, destination_latitude, destination_longitude). Args: context: Tool context containing required Serp API Key secret. origin_address: Origin address. destination_address: Destination address. origin_latitude: Origin latitude. origin_longitude: Origin longitude. destination_latitude: Destination latitude. destination_longitude: Destination longitude. language: Language to use in the Google Maps search. Defaults to 'en' (English). country: 2-letter country code to use in the Google Maps search. Defaults to None (no country is specified). distance_unit: Distance unit to use in the Google Maps search. Defaults to 'km' (kilometers). travel_mode: Travel mode to use in the Google Maps search. Defaults to 'best' (best mode). Returns: The directions from Google Maps. """ language = language.lower() if language not in LANGUAGE_CODES: raise LanguageNotFoundError(language) params = prepare_params( engine="google_maps_directions", hl=language, distance_unit=distance_unit.to_api_value(), travel_mode=travel_mode.to_api_value(), ) if any([ origin_latitude, origin_longitude, destination_latitude, destination_longitude, ]) and any([origin_address, destination_address]): raise ValueError("Either coordinates or addresses must be provided, not both") elif all([origin_latitude, origin_longitude, destination_latitude, destination_longitude]): params["start_coords"] = f"{origin_latitude},{origin_longitude}" params["end_coords"] = f"{destination_latitude},{destination_longitude}" elif all([origin_address, destination_address]): params["start_addr"] = str(origin_address) params["end_addr"] = str(destination_address) else: raise ValueError("Either coordinates or addresses must be provided") if country: country = country.lower() if country not in COUNTRY_CODES: raise CountryNotFoundError(country) params["gl"] = country results = call_serpapi(context, params) directions = cast(list[dict[str, Any]], results.get("directions", [])) for direction in directions: clean_google_maps_direction(direction) if "arrive_around" in direction: direction["arrive_around"] = enrich_google_maps_arrive_around( direction["arrive_around"] ) return directions def clean_google_maps_direction(direction: dict[str, Any]) -> None: for trip in direction.get("trips", []): with contextlib.suppress(KeyError): del trip["start_stop"]["data_id"] del trip["end_stop"]["data_id"] for detail in trip.get("details", []): with contextlib.suppress(KeyError): del detail["geo_photo"] del detail["gps_coordinates"] for stop in trip.get("stops", []): with contextlib.suppress(KeyError): del stop["data_id"] def enrich_google_maps_arrive_around(timestamp: Optional[int]) -> dict[str, Any]: if not timestamp: return {} dt = datetime.fromtimestamp(timestamp, tz=ZoneInfo("UTC")).isoformat() return {"datetime": dt, "timestamp": timestamp} # ------------------------------------------------------------------------------------------------ # Google Flights utils # ------------------------------------------------------------------------------------------------ def parse_flight_results(results: dict[str, Any]) -> dict[str, Any]: """Parse the flight results from the Google Flights API Note: Best flights is not always returned from the API. """ flight_data = {} flights = [] if "best_flights" in results: flights.extend(results["best_flights"]) if "other_flights" in results: flights.extend(results["other_flights"]) if "price_insights" in results: flight_data["price_insights"] = results["price_insights"] flight_data["flights"] = flights return flight_data # ------------------------------------------------------------------------------------------------ # Google News utils # ------------------------------------------------------------------------------------------------ def extract_news_results( results: dict[str, Any], limit: Optional[int] = None ) -> list[dict[str, Any]]: news_results = [] for result in results.get("news_results", []): news_results.append({ "title": result.get("title"), "snippet": result.get("snippet"), "link": result.get("link"), "date": result.get("date"), "source": result.get("source", {}).get("name"), }) if limit: return news_results[:limit] return news_results