From a3f55378ca861aad2026d675de5baabdef6023f1 Mon Sep 17 00:00:00 2001 From: Renato Byrro Date: Tue, 18 Mar 2025 15:23:36 -0300 Subject: [PATCH] Add Google News, Jobs, and Maps Tools (#311) Documentation PR: https://github.com/ArcadeAI/docs/pull/190 --- toolkits/search/arcade_search/constants.py | 19 ++ toolkits/search/arcade_search/enums.py | 46 +++ toolkits/search/arcade_search/exceptions.py | 26 ++ toolkits/search/arcade_search/google_data.py | 281 ++++++++++++++++++ .../search/arcade_search/tools/__init__.py | 8 + .../search/arcade_search/tools/google_jobs.py | 65 ++++ .../search/arcade_search/tools/google_maps.py | 100 +++++++ .../search/arcade_search/tools/google_news.py | 47 +++ toolkits/search/arcade_search/utils.py | 158 +++++++++- toolkits/search/evals/eval_google_jobs.py | 157 ++++++++++ .../evals/eval_google_maps_directions.py | 226 ++++++++++++++ toolkits/search/tests/test_google_jobs.py | 91 ++++++ .../tests/test_google_maps_directions.py | 132 ++++++++ toolkits/search/tests/test_google_search.py | 4 +- 14 files changed, 1354 insertions(+), 6 deletions(-) create mode 100644 toolkits/search/arcade_search/constants.py create mode 100644 toolkits/search/arcade_search/exceptions.py create mode 100644 toolkits/search/arcade_search/google_data.py create mode 100644 toolkits/search/arcade_search/tools/google_jobs.py create mode 100644 toolkits/search/arcade_search/tools/google_maps.py create mode 100644 toolkits/search/arcade_search/tools/google_news.py create mode 100644 toolkits/search/evals/eval_google_jobs.py create mode 100644 toolkits/search/evals/eval_google_maps_directions.py create mode 100644 toolkits/search/tests/test_google_jobs.py create mode 100644 toolkits/search/tests/test_google_maps_directions.py diff --git a/toolkits/search/arcade_search/constants.py b/toolkits/search/arcade_search/constants.py new file mode 100644 index 00000000..15dc8b92 --- /dev/null +++ b/toolkits/search/arcade_search/constants.py @@ -0,0 +1,19 @@ +import os + +from arcade_search.enums import GoogleMapsDistanceUnit, GoogleMapsTravelMode + +DEFAULT_GOOGLE_LANGUAGE = os.getenv("ARCADE_GOOGLE_LANGUAGE", "en") + +DEFAULT_GOOGLE_NEWS_LANGUAGE = os.getenv("ARCADE_GOOGLE_NEWS_LANGUAGE", DEFAULT_GOOGLE_LANGUAGE) +DEFAULT_GOOGLE_NEWS_COUNTRY = os.getenv("ARCADE_GOOGLE_NEWS_COUNTRY", None) + +DEFAULT_GOOGLE_JOBS_LANGUAGE = os.getenv("ARCADE_GOOGLE_JOBS_LANGUAGE", DEFAULT_GOOGLE_LANGUAGE) + +DEFAULT_GOOGLE_MAPS_LANGUAGE = os.getenv("ARCADE_GOOGLE_MAPS_LANGUAGE", DEFAULT_GOOGLE_LANGUAGE) +DEFAULT_GOOGLE_MAPS_COUNTRY = os.getenv("ARCADE_GOOGLE_MAPS_COUNTRY", None) +DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT = GoogleMapsDistanceUnit( + os.getenv("ARCADE_GOOGLE_MAPS_DISTANCE_UNIT", GoogleMapsDistanceUnit.KM.value) +) +DEFAULT_GOOGLE_MAPS_TRAVEL_MODE = GoogleMapsTravelMode( + os.getenv("ARCADE_GOOGLE_MAPS_TRAVEL_MODE", GoogleMapsTravelMode.BEST.value) +) diff --git a/toolkits/search/arcade_search/enums.py b/toolkits/search/arcade_search/enums.py index 8e43f80e..8c9fee94 100644 --- a/toolkits/search/arcade_search/enums.py +++ b/toolkits/search/arcade_search/enums.py @@ -1,6 +1,9 @@ from enum import Enum +# ------------------------------------------------------------------------------------------------ +# Google Finance enumerations +# ------------------------------------------------------------------------------------------------ class GoogleFinanceWindow(Enum): ONE_DAY = "1D" FIVE_DAYS = "5D" @@ -12,6 +15,9 @@ class GoogleFinanceWindow(Enum): MAX = "MAX" +# ------------------------------------------------------------------------------------------------ +# Google Flights enumerations +# ------------------------------------------------------------------------------------------------ class GoogleFlightsTravelClass(Enum): ECONOMY = "ECONOMY" PREMIUM_ECONOMY = "PREMIUM_ECONOMY" @@ -64,6 +70,9 @@ class GoogleFlightsSortBy(Enum): return _map[self.value] +# ------------------------------------------------------------------------------------------------ +# Google Hotels enumerations +# ------------------------------------------------------------------------------------------------ class GoogleHotelsSortBy(Enum): RELEVANCE = "RELEVANCE" LOWEST_PRICE = "LOWEST_PRICE" @@ -78,3 +87,40 @@ class GoogleHotelsSortBy(Enum): "MOST_REVIEWED": 13, } return _map[self.value] + + +# ------------------------------------------------------------------------------------------------ +# Google Maps enumerations +# ------------------------------------------------------------------------------------------------ +class GoogleMapsTravelMode(Enum): + BEST = "best" + DRIVING = "driving" + MOTORCYCLE = "motorcycle" + PUBLIC_TRANSPORTATION = "public_transportation" + WALKING = "walking" + BICYCLE = "bicycle" + FLIGHT = "flight" + + def to_api_value(self) -> int: + _map = { + str(self.BEST): 6, + str(self.DRIVING): 0, + str(self.MOTORCYCLE): 9, + str(self.PUBLIC_TRANSPORTATION): 3, + str(self.WALKING): 2, + str(self.BICYCLE): 1, + str(self.FLIGHT): 4, + } + return _map[str(self)] + + +class GoogleMapsDistanceUnit(Enum): + KM = "km" + MILES = "mi" + + def to_api_value(self) -> int: + _map = { + str(self.KM): 0, + str(self.MILES): 1, + } + return _map[str(self)] diff --git a/toolkits/search/arcade_search/exceptions.py b/toolkits/search/arcade_search/exceptions.py new file mode 100644 index 00000000..9e3c1b7c --- /dev/null +++ b/toolkits/search/arcade_search/exceptions.py @@ -0,0 +1,26 @@ +import json +from typing import Optional + +from arcade.sdk.errors import RetryableToolError + +from arcade_search.google_data import COUNTRY_CODES, LANGUAGE_CODES + + +class GoogleRetryableError(RetryableToolError): + pass + + +class CountryNotFoundError(GoogleRetryableError): + def __init__(self, country: Optional[str]) -> None: + valid_countries = json.dumps(COUNTRY_CODES, default=str) + message = f"Country not found: '{country}'." + additional_message = f"Valid countries are: {valid_countries}" + super().__init__(message, additional_prompt_content=additional_message) + + +class LanguageNotFoundError(GoogleRetryableError): + def __init__(self, language: Optional[str]) -> None: + valid_languages = json.dumps(LANGUAGE_CODES, default=str) + message = f"Language not found: '{language}'." + additional_message = f"Valid languages are: {valid_languages}" + super().__init__(message, additional_prompt_content=additional_message) diff --git a/toolkits/search/arcade_search/google_data.py b/toolkits/search/arcade_search/google_data.py new file mode 100644 index 00000000..789e3183 --- /dev/null +++ b/toolkits/search/arcade_search/google_data.py @@ -0,0 +1,281 @@ +COUNTRY_CODES = { + "af": "Afghanistan", + "al": "Albania", + "dz": "Algeria", + "as": "American Samoa", + "ad": "Andorra", + "ao": "Angola", + "ai": "Anguilla", + "aq": "Antarctica", + "ag": "Antigua and Barbuda", + "ar": "Argentina", + "am": "Armenia", + "aw": "Aruba", + "au": "Australia", + "at": "Austria", + "az": "Azerbaijan", + "bs": "Bahamas", + "bh": "Bahrain", + "bd": "Bangladesh", + "bb": "Barbados", + "by": "Belarus", + "be": "Belgium", + "bz": "Belize", + "bj": "Benin", + "bm": "Bermuda", + "bt": "Bhutan", + "bo": "Bolivia", + "ba": "Bosnia and Herzegovina", + "bw": "Botswana", + "bv": "Bouvet Island", + "br": "Brazil", + "io": "British Indian Ocean Territory", + "bn": "Brunei Darussalam", + "bg": "Bulgaria", + "bf": "Burkina Faso", + "bi": "Burundi", + "kh": "Cambodia", + "cm": "Cameroon", + "ca": "Canada", + "cv": "Cape Verde", + "ky": "Cayman Islands", + "cf": "Central African Republic", + "td": "Chad", + "cl": "Chile", + "cn": "China", + "cx": "Christmas Island", + "cc": "Cocos (Keeling) Islands", + "co": "Colombia", + "km": "Comoros", + "cg": "Congo", + "cd": "Congo, the Democratic Republic of the", + "ck": "Cook Islands", + "cr": "Costa Rica", + "ci": "Cote D'ivoire", + "hr": "Croatia", + "cu": "Cuba", + "cy": "Cyprus", + "cz": "Czech Republic", + "dk": "Denmark", + "dj": "Djibouti", + "dm": "Dominica", + "do": "Dominican Republic", + "ec": "Ecuador", + "eg": "Egypt", + "sv": "El Salvador", + "gq": "Equatorial Guinea", + "er": "Eritrea", + "ee": "Estonia", + "et": "Ethiopia", + "fk": "Falkland Islands (Malvinas)", + "fo": "Faroe Islands", + "fj": "Fiji", + "fi": "Finland", + "fr": "France", + "gf": "French Guiana", + "pf": "French Polynesia", + "tf": "French Southern Territories", + "ga": "Gabon", + "gm": "Gambia", + "ge": "Georgia", + "de": "Germany", + "gh": "Ghana", + "gi": "Gibraltar", + "gr": "Greece", + "gl": "Greenland", + "gd": "Grenada", + "gp": "Guadeloupe", + "gu": "Guam", + "gt": "Guatemala", + "gg": "Guernsey", + "gn": "Guinea", + "gw": "Guinea-Bissau", + "gy": "Guyana", + "ht": "Haiti", + "hm": "Heard Island and Mcdonald Islands", + "va": "Holy See (Vatican City State)", + "hn": "Honduras", + "hk": "Hong Kong", + "hu": "Hungary", + "is": "Iceland", + "in": "India", + "id": "Indonesia", + "ir": "Iran, Islamic Republic of", + "iq": "Iraq", + "ie": "Ireland", + "im": "Isle of Man", + "il": "Israel", + "it": "Italy", + "je": "Jersey", + "jm": "Jamaica", + "jp": "Japan", + "jo": "Jordan", + "kz": "Kazakhstan", + "ke": "Kenya", + "ki": "Kiribati", + "kp": "Korea, Democratic People's Republic of", + "kr": "Korea, Republic of", + "kw": "Kuwait", + "kg": "Kyrgyzstan", + "la": "Lao People's Democratic Republic", + "lv": "Latvia", + "lb": "Lebanon", + "ls": "Lesotho", + "lr": "Liberia", + "ly": "Libyan Arab Jamahiriya", + "li": "Liechtenstein", + "lt": "Lithuania", + "lu": "Luxembourg", + "mo": "Macao", + "mk": "Macedonia, the Former Yugosalv Republic of", + "mg": "Madagascar", + "mw": "Malawi", + "my": "Malaysia", + "mv": "Maldives", + "ml": "Mali", + "mt": "Malta", + "mh": "Marshall Islands", + "mq": "Martinique", + "mr": "Mauritania", + "mu": "Mauritius", + "yt": "Mayotte", + "mx": "Mexico", + "fm": "Micronesia, Federated States of", + "md": "Moldova, Republic of", + "mc": "Monaco", + "mn": "Mongolia", + "me": "Montenegro", + "ms": "Montserrat", + "ma": "Morocco", + "mz": "Mozambique", + "mm": "Myanmar", + "na": "Namibia", + "nr": "Nauru", + "np": "Nepal", + "nl": "Netherlands", + "an": "Netherlands Antilles", + "nc": "New Caledonia", + "nz": "New Zealand", + "ni": "Nicaragua", + "ne": "Niger", + "ng": "Nigeria", + "nu": "Niue", + "nf": "Norfolk Island", + "mp": "Northern Mariana Islands", + "no": "Norway", + "om": "Oman", + "pk": "Pakistan", + "pw": "Palau", + "ps": "Palestinian Territory, Occupied", + "pa": "Panama", + "pg": "Papua New Guinea", + "py": "Paraguay", + "pe": "Peru", + "ph": "Philippines", + "pn": "Pitcairn", + "pl": "Poland", + "pt": "Portugal", + "pr": "Puerto Rico", + "qa": "Qatar", + "re": "Reunion", + "ro": "Romania", + "ru": "Russian Federation", + "rw": "Rwanda", + "sh": "Saint Helena", + "kn": "Saint Kitts and Nevis", + "lc": "Saint Lucia", + "pm": "Saint Pierre and Miquelon", + "vc": "Saint Vincent and the Grenadines", + "ws": "Samoa", + "sm": "San Marino", + "st": "Sao Tome and Principe", + "sa": "Saudi Arabia", + "sn": "Senegal", + "rs": "Serbia", + "sc": "Seychelles", + "sl": "Sierra Leone", + "sg": "Singapore", + "sk": "Slovakia", + "si": "Slovenia", + "sb": "Solomon Islands", + "so": "Somalia", + "za": "South Africa", + "gs": "South Georgia and the South Sandwich Islands", + "es": "Spain", + "lk": "Sri Lanka", + "sd": "Sudan", + "sr": "Suriname", + "sj": "Svalbard and Jan Mayen", + "sz": "Swaziland", + "se": "Sweden", + "ch": "Switzerland", + "sy": "Syrian Arab Republic", + "tw": "Taiwan, Province of China", + "tj": "Tajikistan", + "tz": "Tanzania, United Republic of", + "th": "Thailand", + "tl": "Timor-Leste", + "tg": "Togo", + "tk": "Tokelau", + "to": "Tonga", + "tt": "Trinidad and Tobago", + "tn": "Tunisia", + "tr": "Turkiye", + "tm": "Turkmenistan", + "tc": "Turks and Caicos Islands", + "tv": "Tuvalu", + "ug": "Uganda", + "ua": "Ukraine", + "ae": "United Arab Emirates", + "uk": "United Kingdom", + "gb": "United Kingdom", + "us": "United States", + "um": "United States Minor Outlying Islands", + "uy": "Uruguay", + "uz": "Uzbekistan", + "vu": "Vanuatu", + "ve": "Venezuela", + "vn": "Viet Nam", + "vg": "Virgin Islands, British", + "vi": "Virgin Islands, U.S.", + "wf": "Wallis and Futuna", + "eh": "Western Sahara", + "ye": "Yemen", + "zm": "Zambia", + "zw": "Zimbabwe", +} + + +LANGUAGE_CODES = { + "ar": "Arabic", + "bn": "Bengali", + "da": "Danish", + "de": "German", + "el": "Greek", + "en": "English", + "es": "Spanish", + "fi": "Finnish", + "fr": "French", + "hi": "Hindi", + "hu": "Hungarian", + "id": "Indonesian", + "it": "Italian", + "ja": "Japanese", + "ko": "Korean", + "nl": "Dutch", + "ms": "Malay", + "no": "Norwegian", + "pcm": "Nigerian Pidgin", + "pl": "Polish", + "pt": "Portuguese", + "pt-br": "Portuguese (Brazil)", + "pt-pt": "Portuguese (Portugal)", + "ru": "Russian", + "sv": "Swedish", + "tl": "Filipino", + "tr": "Turkish", + "uk": "Ukrainian", + "zh": "Chinese", + "zh-cn": "Chinese (Simplified)", + "zh-tw": "Chinese (Traditional)", +} diff --git a/toolkits/search/arcade_search/tools/__init__.py b/toolkits/search/arcade_search/tools/__init__.py index c99428f1..3d5d54de 100644 --- a/toolkits/search/arcade_search/tools/__init__.py +++ b/toolkits/search/arcade_search/tools/__init__.py @@ -1,6 +1,11 @@ from arcade_search.tools.google_finance import get_stock_historical_data, get_stock_summary from arcade_search.tools.google_flights import search_one_way_flights, search_roundtrip_flights from arcade_search.tools.google_hotels import search_hotels +from arcade_search.tools.google_jobs import search_jobs +from arcade_search.tools.google_maps import ( + get_directions_between_addresses, + get_directions_between_coordinates, +) from arcade_search.tools.google_search import search_google __all__ = [ @@ -10,4 +15,7 @@ __all__ = [ "search_one_way_flights", # Google Flights "search_roundtrip_flights", # Google Flights "search_hotels", # Google Hotels + "get_directions_between_addresses", # Google Maps + "get_directions_between_coordinates", # Google Maps + "search_jobs", # Google Jobs ] diff --git a/toolkits/search/arcade_search/tools/google_jobs.py b/toolkits/search/arcade_search/tools/google_jobs.py new file mode 100644 index 00000000..c11a1d89 --- /dev/null +++ b/toolkits/search/arcade_search/tools/google_jobs.py @@ -0,0 +1,65 @@ +from typing import Annotated, Optional + +from arcade.sdk import ToolContext, tool + +from arcade_search.constants import DEFAULT_GOOGLE_JOBS_LANGUAGE +from arcade_search.exceptions import LanguageNotFoundError +from arcade_search.google_data import LANGUAGE_CODES +from arcade_search.utils import call_serpapi, prepare_params + + +@tool(requires_secrets=["SERP_API_KEY"]) +async def search_jobs( + context: ToolContext, + query: Annotated[ + str, + "Search query. Provide a job title, company name, and/or any keywords in general " + "representing what kind of jobs the user is looking for. E.g. 'software engineer' " + "or 'data analyst at Apple'.", + ], + location: Annotated[ + Optional[str], + "Location to search for jobs. E.g. 'United States' or 'New York, NY'. Defaults to None.", + ] = None, + language: Annotated[ + str, + "2-character language code to use in the Google Jobs search. ", + f"E.g. 'en' for English. Defaults to '{DEFAULT_GOOGLE_JOBS_LANGUAGE}'.", + ] = DEFAULT_GOOGLE_JOBS_LANGUAGE, + limit: Annotated[ + int, + "Maximum number of results to retrieve. Defaults to 10 (max supported by the API).", + ] = 10, + next_page_token: Annotated[ + Optional[str], + "Next page token to paginate results. Defaults to None (start from the first page).", + ] = None, +) -> Annotated[dict, "Google Jobs results"]: + """Search Google Jobs using SerpAPI.""" + if language not in LANGUAGE_CODES: + raise LanguageNotFoundError(language) + + params = prepare_params( + engine="google_jobs", + q=query, + hl=language, + ) + + if location: + params["location"] = location + + if next_page_token: + params["next_page_token"] = next_page_token + + results = call_serpapi(context, params) + jobs_results = results.get("jobs_results", []) + + try: + next_page_token = results["serpapi_pagination"]["next_page_token"] + except KeyError: + next_page_token = None + + return { + "jobs": jobs_results[:limit], + "next_page_token": next_page_token, + } diff --git a/toolkits/search/arcade_search/tools/google_maps.py b/toolkits/search/arcade_search/tools/google_maps.py new file mode 100644 index 00000000..f11f05f9 --- /dev/null +++ b/toolkits/search/arcade_search/tools/google_maps.py @@ -0,0 +1,100 @@ +from typing import Annotated, Optional + +from arcade.sdk import ToolContext, tool + +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.utils import get_google_maps_directions + + +@tool(requires_secrets=["SERP_API_KEY"]) +async def get_directions_between_addresses( + context: ToolContext, + origin_address: Annotated[ + str, "The origin address. Example: '123 Main St, New York, NY 10001'" + ], + destination_address: Annotated[ + str, "The destination address. Example: '456 Main St, New York, NY 10001'" + ], + language: Annotated[ + str, + "2-character language code to use in the Google Maps search. " + f"Defaults to '{DEFAULT_GOOGLE_MAPS_LANGUAGE}'.", + ] = DEFAULT_GOOGLE_MAPS_LANGUAGE, + country: Annotated[ + Optional[str], + "2-character country code to use in the Google Maps search. " + f"Defaults to '{DEFAULT_GOOGLE_MAPS_COUNTRY}'.", + ] = DEFAULT_GOOGLE_MAPS_COUNTRY, + distance_unit: Annotated[ + GoogleMapsDistanceUnit, + f"Distance unit to use in the Google Maps search. Defaults to " + f"'{DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT}'.", + ] = DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, + travel_mode: Annotated[ + GoogleMapsTravelMode, + f"Travel mode to use in the Google Maps search. Defaults to " + f"'{DEFAULT_GOOGLE_MAPS_TRAVEL_MODE}'.", + ] = DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, +) -> Annotated[dict, "The directions from Google Maps"]: + """Get directions from Google Maps.""" + return { + "directions": get_google_maps_directions( + context=context, + origin_address=origin_address, + destination_address=destination_address, + language=language, + country=country, + distance_unit=distance_unit, + travel_mode=travel_mode, + ), + } + + +@tool(requires_secrets=["SERP_API_KEY"]) +async def get_directions_between_coordinates( + context: ToolContext, + origin_latitude: Annotated[str, "The origin latitude. E.g. '40.7128'"], + origin_longitude: Annotated[str, "The origin longitude. E.g. '-74.0060'"], + destination_latitude: Annotated[str, "The destination latitude. E.g. '40.7128'"], + destination_longitude: Annotated[str, "The destination longitude. E.g. '-74.0060'"], + language: Annotated[ + str, + "2-letter language code to use in the Google Maps search. " + f"Defaults to '{DEFAULT_GOOGLE_MAPS_LANGUAGE}'.", + ] = DEFAULT_GOOGLE_MAPS_LANGUAGE, + country: Annotated[ + Optional[str], + f"2-letter country code to use in the Google Maps search. Defaults to " + f"'{DEFAULT_GOOGLE_MAPS_COUNTRY}'.", + ] = DEFAULT_GOOGLE_MAPS_COUNTRY, + distance_unit: Annotated[ + GoogleMapsDistanceUnit, + f"Distance unit to use in the Google Maps search. Defaults to " + f"'{DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT}'.", + ] = DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, + travel_mode: Annotated[ + GoogleMapsTravelMode, + f"Travel mode to use in the Google Maps search. Defaults to " + f"'{DEFAULT_GOOGLE_MAPS_TRAVEL_MODE}'.", + ] = DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, +) -> Annotated[dict, "The directions from Google Maps"]: + """Get directions from Google Maps.""" + return { + "directions": get_google_maps_directions( + context=context, + origin_latitude=origin_latitude, + origin_longitude=origin_longitude, + destination_latitude=destination_latitude, + destination_longitude=destination_longitude, + language=language, + country=country, + distance_unit=distance_unit, + travel_mode=travel_mode, + ), + } diff --git a/toolkits/search/arcade_search/tools/google_news.py b/toolkits/search/arcade_search/tools/google_news.py new file mode 100644 index 00000000..ca58b928 --- /dev/null +++ b/toolkits/search/arcade_search/tools/google_news.py @@ -0,0 +1,47 @@ +from typing import Annotated, Any, Optional + +from arcade.sdk import ToolContext, tool +from arcade.sdk.errors import ToolExecutionError + +from arcade_search.constants import DEFAULT_GOOGLE_NEWS_COUNTRY, DEFAULT_GOOGLE_NEWS_LANGUAGE +from arcade_search.exceptions import CountryNotFoundError, LanguageNotFoundError +from arcade_search.google_data import COUNTRY_CODES, LANGUAGE_CODES +from arcade_search.utils import call_serpapi, extract_news_results, prepare_params + + +@tool(requires_secrets=["SERP_API_KEY"]) +async def search_news_stories( + context: ToolContext, + keywords: Annotated[ + str, + "Keywords to search for news articles. E.g. 'Apple launches new iPhone'.", + ], + country_code: Annotated[ + Optional[str], + "2-character country code to search for news articles. E.g. 'us' (United States). " + f"Defaults to '{DEFAULT_GOOGLE_NEWS_COUNTRY}'.", + ] = None, + language_code: Annotated[ + str, + "2-character language code to search for news articles. E.g. 'en' (English). " + f"Defaults to '{DEFAULT_GOOGLE_NEWS_LANGUAGE}'.", + ] = DEFAULT_GOOGLE_NEWS_LANGUAGE, + limit: Annotated[ + Optional[int], + "Maximum number of news articles to return. Defaults to None " + "(returns all results found by the API).", + ] = None, +) -> Annotated[dict[str, list[dict[str, Any]]], "News results."]: + """Search for news articles related to a given query.""" + if not keywords: + raise ToolExecutionError("Keywords are required to search for news articles.") + + if country_code and country_code not in COUNTRY_CODES: + raise CountryNotFoundError(country_code) + + if language_code not in LANGUAGE_CODES: + raise LanguageNotFoundError(language_code) + + params = prepare_params("google_news", q=keywords, gl=country_code, hl=language_code) + results = call_serpapi(context, params) + return {"news_results": extract_news_results(results, limit=limit)} diff --git a/toolkits/search/arcade_search/utils.py b/toolkits/search/arcade_search/utils.py index 7a559346..68a567ea 100644 --- a/toolkits/search/arcade_search/utils.py +++ b/toolkits/search/arcade_search/utils.py @@ -1,9 +1,22 @@ +import contextlib import re -from typing import Any +from datetime import datetime +from typing import Any, Optional, cast +from zoneinfo import ZoneInfo -import serpapi 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 # ------------------------------------------------------------------------------------------------ @@ -38,10 +51,10 @@ def call_serpapi(context: ToolContext, params: dict) -> dict: The search results as a dictionary. """ api_key = context.get_secret("SERP_API_KEY") - client = serpapi.Client(api_key=api_key) + client = SerpClient(api_key=api_key) try: search = client.search(params) - return search.as_dict() # type: ignore[no-any-return] + 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)) @@ -51,6 +64,122 @@ def call_serpapi(context: ToolContext, params: dict) -> dict: ) +# ------------------------------------------------------------------------------------------------ +# 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 # ------------------------------------------------------------------------------------------------ @@ -72,3 +201,24 @@ def parse_flight_results(results: dict[str, Any]) -> dict[str, Any]: 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 diff --git a/toolkits/search/evals/eval_google_jobs.py b/toolkits/search/evals/eval_google_jobs.py new file mode 100644 index 00000000..16be3864 --- /dev/null +++ b/toolkits/search/evals/eval_google_jobs.py @@ -0,0 +1,157 @@ +from arcade.sdk import ToolCatalog +from arcade.sdk.eval import ( + BinaryCritic, + EvalRubric, + EvalSuite, + ExpectedToolCall, + NoneCritic, + SimilarityCritic, + tool_eval, +) + +import arcade_search +from arcade_search.constants import DEFAULT_GOOGLE_JOBS_LANGUAGE +from arcade_search.tools.google_jobs import search_jobs + +# Evaluation rubric +rubric = EvalRubric( + fail_threshold=0.8, + warn_threshold=0.9, +) + +catalog = ToolCatalog() +# Register the Google Jobs tool +catalog.add_module(arcade_search) + + +@tool_eval() +def google_jobs_eval_suite() -> EvalSuite: + """Create an evaluation suite for the Google Jobs tool.""" + suite = EvalSuite( + name="Google Jobs Tool Evaluation", + system_message="You are an AI assistant that can perform job searches using the provided tools.", + catalog=catalog, + rubric=rubric, + ) + + suite.add_case( + name="Search for 'backend engineer' jobs", + user_message="Search for 'backend engineer' jobs", + expected_tool_calls=[ + ExpectedToolCall( + func=search_jobs, + args={ + "query": "backend engineer", + "location": None, + "language": DEFAULT_GOOGLE_JOBS_LANGUAGE, + "limit": 10, + "next_page_token": None, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="query", weight=0.5), + NoneCritic(critic_field="location", weight=0.1), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="limit", weight=0.1), + NoneCritic(critic_field="next_page_token", weight=0.1), + ], + ) + + suite.add_case( + name="Search for 'senior backend engineer' jobs that are part-time", + user_message="Search for senior backend engineer jobs that are part-time", + expected_tool_calls=[ + ExpectedToolCall( + func=search_jobs, + args={ + "query": "part-time senior backend engineer", + "location": None, + "language": DEFAULT_GOOGLE_JOBS_LANGUAGE, + "limit": 10, + "next_page_token": None, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="query", weight=0.5), + NoneCritic(critic_field="location", weight=0.1), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="limit", weight=0.1), + NoneCritic(critic_field="next_page_token", weight=0.1), + ], + ) + + suite.add_case( + name="Search for 'backend engineer' jobs in San Francisco", + user_message="Search for 'backend engineer' jobs in San Francisco", + expected_tool_calls=[ + ExpectedToolCall( + func=search_jobs, + args={ + "query": "backend engineer", + "location": "San Francisco", + "language": DEFAULT_GOOGLE_JOBS_LANGUAGE, + "limit": 10, + "next_page_token": None, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="query", weight=0.35), + SimilarityCritic(critic_field="location", weight=0.35), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="limit", weight=0.1), + NoneCritic(critic_field="next_page_token", weight=0.1), + ], + ) + + suite.add_case( + name="Get the first 3 jobs for 'backend engineer' in San Francisco", + user_message="Get the first 3 jobs for 'backend engineer' in San Francisco", + expected_tool_calls=[ + ExpectedToolCall( + func=search_jobs, + args={ + "query": "backend engineer", + "location": "San Francisco", + "language": DEFAULT_GOOGLE_JOBS_LANGUAGE, + "limit": 3, + "next_page_token": None, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="query", weight=0.25), + SimilarityCritic(critic_field="location", weight=0.25), + BinaryCritic(critic_field="language", weight=0.125), + BinaryCritic(critic_field="limit", weight=0.25), + NoneCritic(critic_field="next_page_token", weight=0.125), + ], + ) + + suite.add_case( + name="Search for 'engenheiro de software' jobs in Brazil, return results in Portuguese", + user_message="Search for 'engenheiro de software' jobs in Brazil, return results in Portuguese", + expected_tool_calls=[ + ExpectedToolCall( + func=search_jobs, + args={ + "query": "engenheiro de software", + "location": "Brazil", + "language": "pt", + "limit": 10, + "next_page_token": None, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="query", weight=0.25), + SimilarityCritic(critic_field="location", weight=0.125), + BinaryCritic(critic_field="language", weight=0.25), + BinaryCritic(critic_field="limit", weight=0.125), + NoneCritic(critic_field="next_page_token", weight=0.125), + ], + ) + + return suite diff --git a/toolkits/search/evals/eval_google_maps_directions.py b/toolkits/search/evals/eval_google_maps_directions.py new file mode 100644 index 00000000..04a7ebe2 --- /dev/null +++ b/toolkits/search/evals/eval_google_maps_directions.py @@ -0,0 +1,226 @@ +from arcade.sdk import ToolCatalog +from arcade.sdk.eval import ( + BinaryCritic, + EvalRubric, + EvalSuite, + ExpectedToolCall, + SimilarityCritic, + tool_eval, +) + +import arcade_search +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.tools.google_maps import ( + get_directions_between_addresses, + get_directions_between_coordinates, +) + +# Evaluation rubric +rubric = EvalRubric( + fail_threshold=0.8, + warn_threshold=0.9, +) + +catalog = ToolCatalog() +# Register the Google Search tool +catalog.add_module(arcade_search) + + +@tool_eval() +def google_maps_directions_by_addresses_eval_suite() -> EvalSuite: + """Create an evaluation suite for the Google Maps Directions tools.""" + suite = EvalSuite( + name="Google Maps Directions Tool Evaluation", + system_message="You are an AI assistant that can get directions from Google Maps using the provided tools.", + catalog=catalog, + rubric=rubric, + ) + + suite.add_case( + name="Get directions between two addresses", + user_message="Get directions from Google Maps between the following addresses: 1600 Amphitheatre Parkway, Mountain View, CA 94043 and 1 Infinite Loop, Cupertino, CA 95014.", + expected_tool_calls=[ + ExpectedToolCall( + func=get_directions_between_addresses, + args={ + "origin_address": "1600 Amphitheatre Parkway, Mountain View, CA 94043", + "destination_address": "1 Infinite Loop, Cupertino, CA 95014", + "language": DEFAULT_GOOGLE_MAPS_LANGUAGE, + "country": DEFAULT_GOOGLE_MAPS_COUNTRY, + "distance_unit": DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, + "travel_mode": DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="origin_address", weight=0.3), + SimilarityCritic(critic_field="destination_address", weight=0.3), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="country", weight=0.1), + BinaryCritic(critic_field="distance_unit", weight=0.1), + BinaryCritic(critic_field="travel_mode", weight=0.1), + ], + ) + + suite.add_case( + name="Get directions between two addresses with custom distance unit and travel mode", + user_message="Get walking directions from Google Maps between the following addresses in miles: 1600 Amphitheatre Parkway, Mountain View, CA 94043 and 1 Infinite Loop, Cupertino, CA 95014.", + expected_tool_calls=[ + ExpectedToolCall( + func=get_directions_between_addresses, + args={ + "origin_address": "1600 Amphitheatre Parkway, Mountain View, CA 94043", + "destination_address": "1 Infinite Loop, Cupertino, CA 95014", + "language": DEFAULT_GOOGLE_MAPS_LANGUAGE, + "country": DEFAULT_GOOGLE_MAPS_COUNTRY, + "distance_unit": GoogleMapsDistanceUnit.MILES.value, + "travel_mode": GoogleMapsTravelMode.WALKING.value, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="origin_address", weight=0.3), + SimilarityCritic(critic_field="destination_address", weight=0.3), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="country", weight=0.1), + BinaryCritic(critic_field="distance_unit", weight=0.1), + BinaryCritic(critic_field="travel_mode", weight=0.1), + ], + ) + + suite.add_case( + name="Get directions between two addresses in a given country and language", + user_message="Get directions from Google Maps in Portuguese between the following addresses in Brazil: Rua do Amendoim, 1, Belo Horizonte, MG and Av. do Descobrimento, 515, Porto Seguro, BA.", + expected_tool_calls=[ + ExpectedToolCall( + func=get_directions_between_addresses, + args={ + "origin_address": "Rua do Amendoim, 1, Belo Horizonte, MG", + "destination_address": "Av. do Descobrimento, 515, Porto Seguro, BA", + "language": "pt", + "country": "br", + "distance_unit": DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, + "travel_mode": DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="origin_address", weight=0.3), + SimilarityCritic(critic_field="destination_address", weight=0.3), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="country", weight=0.1), + BinaryCritic(critic_field="distance_unit", weight=0.1), + BinaryCritic(critic_field="travel_mode", weight=0.1), + ], + ) + + return suite + + +@tool_eval() +def google_maps_directions_by_coordinates_eval_suite() -> EvalSuite: + """Create an evaluation suite for the Google Maps Directions tools.""" + suite = EvalSuite( + name="Google Maps Directions Tool Evaluation", + system_message="You are an AI assistant that can get directions from Google Maps using the provided tools.", + catalog=catalog, + rubric=rubric, + ) + + suite.add_case( + name="Get directions between two coordinates", + user_message="Get directions from Google Maps between the following coordinates: 37.422740,-122.084961 and 37.331820,-122.031180.", + expected_tool_calls=[ + ExpectedToolCall( + func=get_directions_between_coordinates, + args={ + "origin_latitude": "37.422740", + "origin_longitude": "-122.084961", + "destination_latitude": "37.331820", + "destination_longitude": "-122.031180", + "language": DEFAULT_GOOGLE_MAPS_LANGUAGE, + "country": DEFAULT_GOOGLE_MAPS_COUNTRY, + "distance_unit": DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, + "travel_mode": DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="origin_latitude", weight=0.15), + SimilarityCritic(critic_field="origin_longitude", weight=0.15), + SimilarityCritic(critic_field="destination_latitude", weight=0.15), + SimilarityCritic(critic_field="destination_longitude", weight=0.15), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="country", weight=0.1), + BinaryCritic(critic_field="distance_unit", weight=0.1), + BinaryCritic(critic_field="travel_mode", weight=0.1), + ], + ) + + suite.add_case( + name="Get directions between two coordinates with custom distance unit and travel mode", + user_message="Get walking directions from Google Maps between the following coordinates in miles: 37.422740,-122.084961 and 37.331820,-122.031180.", + expected_tool_calls=[ + ExpectedToolCall( + func=get_directions_between_coordinates, + args={ + "origin_latitude": "37.422740", + "origin_longitude": "-122.084961", + "destination_latitude": "37.331820", + "destination_longitude": "-122.031180", + "language": DEFAULT_GOOGLE_MAPS_LANGUAGE, + "country": DEFAULT_GOOGLE_MAPS_COUNTRY, + "distance_unit": GoogleMapsDistanceUnit.MILES.value, + "travel_mode": GoogleMapsTravelMode.WALKING.value, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="origin_latitude", weight=0.15), + SimilarityCritic(critic_field="origin_longitude", weight=0.15), + SimilarityCritic(critic_field="destination_latitude", weight=0.15), + SimilarityCritic(critic_field="destination_longitude", weight=0.15), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="country", weight=0.1), + BinaryCritic(critic_field="distance_unit", weight=0.1), + BinaryCritic(critic_field="travel_mode", weight=0.1), + ], + ) + + suite.add_case( + name="Get directions between two coordinates in a given country and language", + user_message="Get directions from Google Maps in Portuguese between the following coordinates in Brazil: 37.422740,-122.084961 and 37.331820,-122.031180.", + expected_tool_calls=[ + ExpectedToolCall( + func=get_directions_between_coordinates, + args={ + "origin_latitude": "37.422740", + "origin_longitude": "-122.084961", + "destination_latitude": "37.331820", + "destination_longitude": "-122.031180", + "language": "pt", + "country": "br", + "distance_unit": DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT, + "travel_mode": DEFAULT_GOOGLE_MAPS_TRAVEL_MODE, + }, + ) + ], + critics=[ + SimilarityCritic(critic_field="origin_latitude", weight=0.15), + SimilarityCritic(critic_field="origin_longitude", weight=0.15), + SimilarityCritic(critic_field="destination_latitude", weight=0.15), + SimilarityCritic(critic_field="destination_longitude", weight=0.15), + BinaryCritic(critic_field="language", weight=0.1), + BinaryCritic(critic_field="country", weight=0.1), + BinaryCritic(critic_field="distance_unit", weight=0.1), + BinaryCritic(critic_field="travel_mode", weight=0.1), + ], + ) + + return suite diff --git a/toolkits/search/tests/test_google_jobs.py b/toolkits/search/tests/test_google_jobs.py new file mode 100644 index 00000000..f8ef9333 --- /dev/null +++ b/toolkits/search/tests/test_google_jobs.py @@ -0,0 +1,91 @@ +from unittest.mock import patch + +import pytest +from arcade.core.schema import ToolSecretItem +from arcade.sdk import ToolContext + +from arcade_search.exceptions import LanguageNotFoundError +from arcade_search.tools.google_jobs import search_jobs + + +@pytest.fixture +def mock_context(): + return ToolContext(secrets=[ToolSecretItem(key="serp_api_key", value="fake_api_key")]) + + +@pytest.mark.asyncio +@patch("arcade_search.utils.SerpClient") +async def test_search_jobs_success(mock_serp_client, mock_context): + mock_serp_client_instance = mock_serp_client.return_value + mock_serp_client_instance.search().as_dict.return_value = { + "jobs_results": [ + {"title": "Job 1", "link": "http://example.com/1"}, + {"title": "Job 2", "link": "http://example.com/2"}, + ] + } + + result = await search_jobs(mock_context, "engenheiro de software", "Brazil", "pt", 10, None) + assert result == { + "jobs": [ + {"title": "Job 1", "link": "http://example.com/1"}, + {"title": "Job 2", "link": "http://example.com/2"}, + ], + "next_page_token": None, + } + + +@pytest.mark.asyncio +@patch("arcade_search.utils.SerpClient") +async def test_search_jobs_success_with_custom_language_and_location( + mock_serp_client, mock_context +): + mock_serp_client_instance = mock_serp_client.return_value + mock_serp_client_instance.search().as_dict.return_value = { + "jobs_results": [ + {"title": "Job 1", "link": "http://example.com/1"}, + {"title": "Job 2", "link": "http://example.com/2"}, + ] + } + + result = await search_jobs( + context=mock_context, + query="engenheiro de software", + location="Brazil", + language="pt", + limit=10, + next_page_token=None, + ) + + mock_serp_client_instance.search.assert_called_with({ + "engine": "google_jobs", + "q": "engenheiro de software", + "hl": "pt", + "location": "Brazil", + }) + + assert result == { + "jobs": [ + {"title": "Job 1", "link": "http://example.com/1"}, + {"title": "Job 2", "link": "http://example.com/2"}, + ], + "next_page_token": None, + } + + +@pytest.mark.asyncio +@patch("arcade_search.utils.SerpClient") +async def test_search_jobs_language_not_found_error(mock_serp_client, mock_context): + mock_serp_client_instance = mock_serp_client.return_value + mock_serp_client_instance.search().as_dict.return_value = { + "jobs_results": [ + {"title": "Job 1", "link": "http://example.com/1"}, + {"title": "Job 2", "link": "http://example.com/2"}, + ] + } + + with pytest.raises(LanguageNotFoundError): + await search_jobs( + context=mock_context, + query="backend engineer", + language="invalid_language", + ) diff --git a/toolkits/search/tests/test_google_maps_directions.py b/toolkits/search/tests/test_google_maps_directions.py new file mode 100644 index 00000000..72cf7fb0 --- /dev/null +++ b/toolkits/search/tests/test_google_maps_directions.py @@ -0,0 +1,132 @@ +from unittest.mock import patch + +import pytest +from arcade.core.schema import ToolSecretItem +from arcade.sdk import ToolContext + +from arcade_search.exceptions import CountryNotFoundError, LanguageNotFoundError +from arcade_search.tools.google_maps import ( + get_directions_between_addresses, + get_directions_between_coordinates, +) + + +@pytest.fixture +def mock_context(): + return ToolContext(secrets=[ToolSecretItem(key="serp_api_key", value="fake_api_key")]) + + +@pytest.mark.asyncio +@patch("arcade_search.utils.SerpClient") +async def test_get_directions_between_coordinates_success(mock_serp_client, mock_context): + mock_serp_client_instance = mock_serp_client.return_value + mock_serp_client_instance.search.return_value.as_dict.return_value = { + "directions": [ + { + "arrive_around": 1741789839, + "distance": "100 miles", + "duration": "1 hour", + } + ] + } + + result = await get_directions_between_coordinates( + context=mock_context, + origin_latitude="1", + origin_longitude="2", + destination_latitude="3", + destination_longitude="4", + ) + + assert result == { + "directions": [ + { + "arrive_around": { + "datetime": "2025-03-12T14:30:39+00:00", + "timestamp": 1741789839, + }, + "distance": "100 miles", + "duration": "1 hour", + } + ] + } + + +@pytest.mark.asyncio +@patch("arcade_search.utils.SerpClient") +async def test_get_directions_between_addresses_success(mock_serp_client, mock_context): + mock_serp_client_instance = mock_serp_client.return_value + mock_serp_client_instance.search.return_value.as_dict.return_value = { + "directions": [ + { + "arrive_around": 1741789839, + "distance": "100 miles", + "duration": "1 hour", + } + ] + } + + result = await get_directions_between_addresses( + context=mock_context, + origin_address="1", + destination_address="2", + ) + + assert result == { + "directions": [ + { + "arrive_around": { + "datetime": "2025-03-12T14:30:39+00:00", + "timestamp": 1741789839, + }, + "distance": "100 miles", + "duration": "1 hour", + } + ] + } + + +@pytest.mark.asyncio +@patch("arcade_search.utils.SerpClient") +async def test_get_directions_between_addresses_country_not_found(mock_serp_client, mock_context): + mock_serp_client_instance = mock_serp_client.return_value + mock_serp_client_instance.search.return_value.as_dict.return_value = { + "directions": [ + { + "arrive_around": 1741789839, + "distance": "100 miles", + "duration": "1 hour", + } + ] + } + + with pytest.raises(CountryNotFoundError): + await get_directions_between_addresses( + context=mock_context, + origin_address="1", + destination_address="2", + country="invalid", + ) + + +@pytest.mark.asyncio +@patch("arcade_search.utils.SerpClient") +async def test_get_directions_between_addresses_language_not_found(mock_serp_client, mock_context): + mock_serp_client_instance = mock_serp_client.return_value + mock_serp_client_instance.search.return_value.as_dict.return_value = { + "directions": [ + { + "arrive_around": 1741789839, + "distance": "100 miles", + "duration": "1 hour", + } + ] + } + + with pytest.raises(LanguageNotFoundError): + await get_directions_between_addresses( + context=mock_context, + origin_address="1", + destination_address="2", + language="invalid", + ) diff --git a/toolkits/search/tests/test_google_search.py b/toolkits/search/tests/test_google_search.py index 95ce94a7..fe470fd6 100644 --- a/toolkits/search/tests/test_google_search.py +++ b/toolkits/search/tests/test_google_search.py @@ -16,7 +16,7 @@ def mock_context(): @pytest.mark.asyncio async def test_search_google_success(mock_context): with ( - patch("serpapi.Client") as MockClient, + patch("arcade_search.utils.SerpClient") as MockClient, ): mock_client_instance = MockClient.return_value mock_client_instance.search.return_value.as_dict.return_value = { @@ -39,7 +39,7 @@ async def test_search_google_success(mock_context): @pytest.mark.asyncio async def test_search_google_no_results(mock_context): with ( - patch("serpapi.Client") as MockClient, + patch("arcade_search.utils.SerpClient") as MockClient, ): mock_client_instance = MockClient.return_value mock_client_instance.search.return_value.as_dict.return_value = {"organic_results": []}