YouTube Search tools (#316)

This commit is contained in:
Renato Byrro 2025-03-21 19:26:52 -03:00 committed by GitHub
parent a3371d71ac
commit 2338e823cf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 248 additions and 4 deletions

View file

@ -2,18 +2,38 @@ import os
from arcade_search.enums import GoogleMapsDistanceUnit, GoogleMapsTravelMode
# ------------------------------------------------------------------------------------------------
# Google default constants
# ------------------------------------------------------------------------------------------------
DEFAULT_GOOGLE_LANGUAGE = os.getenv("ARCADE_GOOGLE_LANGUAGE", "en")
DEFAULT_GOOGLE_COUNTRY = os.getenv("ARCADE_GOOGLE_COUNTRY")
# ------------------------------------------------------------------------------------------------
# Google News default constants
# ------------------------------------------------------------------------------------------------
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_NEWS_COUNTRY = os.getenv("ARCADE_GOOGLE_NEWS_COUNTRY")
# ------------------------------------------------------------------------------------------------
# Google Jobs default constants
# ------------------------------------------------------------------------------------------------
DEFAULT_GOOGLE_JOBS_LANGUAGE = os.getenv("ARCADE_GOOGLE_JOBS_LANGUAGE", DEFAULT_GOOGLE_LANGUAGE)
# ------------------------------------------------------------------------------------------------
# Google Maps default constants
# ------------------------------------------------------------------------------------------------
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_COUNTRY = os.getenv("ARCADE_GOOGLE_MAPS_COUNTRY")
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)
)
# ------------------------------------------------------------------------------------------------
# YouTube default constants
# ------------------------------------------------------------------------------------------------
YOUTUBE_MAX_DESCRIPTION_LENGTH = 500
DEFAULT_YOUTUBE_SEARCH_LANGUAGE = os.getenv("ARCADE_YOUTUBE_SEARCH_LANGUAGE")
DEFAULT_YOUTUBE_SEARCH_COUNTRY = os.getenv("ARCADE_YOUTUBE_SEARCH_COUNTRY")

View file

@ -0,0 +1,101 @@
from typing import Annotated, Any, Optional, cast
from arcade.sdk import ToolContext, tool
from arcade.sdk.errors import ToolExecutionError
from arcade_search.constants import DEFAULT_YOUTUBE_SEARCH_COUNTRY, DEFAULT_YOUTUBE_SEARCH_LANGUAGE
from arcade_search.utils import (
call_serpapi,
default_country_code,
default_language_code,
extract_video_details,
extract_video_results,
prepare_params,
resolve_country_code,
resolve_language_code,
)
@tool(requires_secrets=["SERP_API_KEY"])
async def search_youtube_videos(
context: ToolContext,
keywords: Annotated[
str,
"The keywords to search for. E.g. 'Python tutorial'.",
],
language_code: Annotated[
Optional[str],
"2-character language code to search for. E.g. 'en' for English. "
f"Defaults to '{default_language_code(DEFAULT_YOUTUBE_SEARCH_LANGUAGE)}'.",
] = None,
country_code: Annotated[
Optional[str],
"2-character country code to search for. E.g. 'us' for United States. "
f"Defaults to '{default_country_code(DEFAULT_YOUTUBE_SEARCH_COUNTRY)}'.",
] = None,
next_page_token: Annotated[
Optional[str],
"The next page token to use for pagination. "
"Defaults to `None` (start from the first page).",
] = None,
) -> Annotated[dict[str, Any], "List of YouTube videos related to the query."]:
"""Search for YouTube videos related to the query."""
language_code = resolve_language_code(language_code, DEFAULT_YOUTUBE_SEARCH_LANGUAGE)
country_code = resolve_country_code(country_code, DEFAULT_YOUTUBE_SEARCH_COUNTRY)
params = prepare_params(
"youtube",
search_query=keywords,
hl=language_code,
gl=country_code,
sp=next_page_token,
)
results = call_serpapi(context, params)
if results.get("error"):
error_msg = cast(str, results.get("error"))
raise ToolExecutionError(error_msg)
return {
"videos": extract_video_results(results),
"next_page_token": results.get("serpapi_pagination", {}).get("next_page_token"),
}
@tool(requires_secrets=["SERP_API_KEY"])
async def get_youtube_video_details(
context: ToolContext,
video_id: Annotated[
str,
"The ID of the YouTube video to get details about. E.g. 'dQw4w9WgXcQ'.",
],
language_code: Annotated[
Optional[str],
"2-character language code to search for. E.g. 'en' for English. "
f"Defaults to '{default_language_code(DEFAULT_YOUTUBE_SEARCH_LANGUAGE)}'.",
] = None,
country_code: Annotated[
Optional[str],
"2-character country code to search for. E.g. 'us' for United States. "
f"Defaults to '{default_country_code(DEFAULT_YOUTUBE_SEARCH_COUNTRY)}'.",
] = None,
) -> Annotated[dict[str, Any], "Details about a YouTube video."]:
"""Get details about a YouTube video."""
language_code = resolve_language_code(language_code, DEFAULT_YOUTUBE_SEARCH_LANGUAGE)
country_code = resolve_country_code(country_code, DEFAULT_YOUTUBE_SEARCH_COUNTRY)
params = prepare_params(
"youtube_video",
v=video_id,
hl=language_code,
gl=country_code,
)
results = call_serpapi(context, params)
if results.get("error"):
error_msg = cast(str, results.get("error"))
raise ToolExecutionError(error_msg)
return {
"video": extract_video_details(results),
}

View file

@ -2,6 +2,7 @@ import contextlib
import re
from datetime import datetime
from typing import Any, Optional, cast
from urllib.parse import parse_qs, urlparse
from zoneinfo import ZoneInfo
from arcade.sdk import ToolContext
@ -9,10 +10,13 @@ from arcade.sdk.errors import ToolExecutionError
from serpapi import Client as SerpClient
from arcade_search.constants import (
DEFAULT_GOOGLE_COUNTRY,
DEFAULT_GOOGLE_LANGUAGE,
DEFAULT_GOOGLE_MAPS_COUNTRY,
DEFAULT_GOOGLE_MAPS_DISTANCE_UNIT,
DEFAULT_GOOGLE_MAPS_LANGUAGE,
DEFAULT_GOOGLE_MAPS_TRAVEL_MODE,
YOUTUBE_MAX_DESCRIPTION_LENGTH,
)
from arcade_search.enums import GoogleMapsDistanceUnit, GoogleMapsTravelMode
from arcade_search.exceptions import CountryNotFoundError, LanguageNotFoundError
@ -64,6 +68,53 @@ def call_serpapi(context: ToolContext, params: dict) -> dict:
)
# ------------------------------------------------------------------------------------------------
# Google general utils
# ------------------------------------------------------------------------------------------------
def default_language_code(default_service_language_code: Optional[str] = None) -> Optional[str]:
if isinstance(default_service_language_code, str):
return default_service_language_code.lower()
elif isinstance(DEFAULT_GOOGLE_LANGUAGE, str):
return DEFAULT_GOOGLE_LANGUAGE.lower()
return None
def default_country_code(default_service_country_code: Optional[str] = None) -> Optional[str]:
if isinstance(default_service_country_code, str):
return default_service_country_code.lower()
elif isinstance(DEFAULT_GOOGLE_COUNTRY, str):
return DEFAULT_GOOGLE_COUNTRY.lower()
return None
def resolve_language_code(
language_code: Optional[str] = None,
default_service_language_code: Optional[str] = None,
) -> Optional[str]:
language_code = language_code or default_language_code(default_service_language_code)
if isinstance(language_code, str):
language_code = language_code.lower()
if language_code not in LANGUAGE_CODES:
raise LanguageNotFoundError(language_code)
return language_code
def resolve_country_code(
country_code: Optional[str] = None,
default_service_country_code: Optional[str] = None,
) -> Optional[str]:
country_code = country_code or default_country_code(default_service_country_code)
if isinstance(country_code, str):
country_code = country_code.lower()
if country_code not in COUNTRY_CODES:
raise CountryNotFoundError(country_code)
return country_code
# ------------------------------------------------------------------------------------------------
# Google Maps utils
# ------------------------------------------------------------------------------------------------
@ -75,7 +126,7 @@ def get_google_maps_directions(
origin_longitude: Optional[str] = None,
destination_latitude: Optional[str] = None,
destination_longitude: Optional[str] = None,
language: str = DEFAULT_GOOGLE_MAPS_LANGUAGE,
language: Optional[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,
@ -104,7 +155,8 @@ def get_google_maps_directions(
Returns:
The directions from Google Maps.
"""
language = language.lower()
if isinstance(language, str):
language = language.lower()
if language not in LANGUAGE_CODES:
raise LanguageNotFoundError(language)
@ -222,3 +274,74 @@ def extract_news_results(
if limit:
return news_results[:limit]
return news_results
# ------------------------------------------------------------------------------------------------
# YouTube utils
# ------------------------------------------------------------------------------------------------
def extract_video_id_from_link(link: Optional[str]) -> Optional[str]:
if not isinstance(link, str):
return None
parsed_url = urlparse(link)
query_params = parse_qs(parsed_url.query)
return query_params.get("v", [""])[0]
def extract_video_description(
video: dict[str, Any],
max_description_length: int = YOUTUBE_MAX_DESCRIPTION_LENGTH,
) -> Optional[str]:
description = video.get("description", "")
if isinstance(description, dict):
description = description.get("content", "")
if isinstance(description, str):
too_long = len(description) > max_description_length
if too_long:
description = description[:max_description_length] + " [truncated]"
if description is not None:
description = str(description).strip()
return cast(Optional[str], description)
def extract_video_results(
results: dict[str, Any],
max_description_length: int = YOUTUBE_MAX_DESCRIPTION_LENGTH,
) -> list[dict[str, Any]]:
videos = []
for video in results.get("video_results", []):
videos.append({
"id": extract_video_id_from_link(video.get("link")),
"title": video.get("title"),
"description": extract_video_description(video, max_description_length),
"link": video.get("link"),
"published_date": video.get("published_date"),
"duration": video.get("duration"),
"channel": {
"name": video.get("channel", {}).get("name"),
"link": video.get("channel", {}).get("link"),
},
})
return videos
def extract_video_details(video: dict[str, Any]) -> dict[str, Any]:
return {
"id": extract_video_id_from_link(video.get("link")),
"title": video.get("title"),
"description": extract_video_description(video, YOUTUBE_MAX_DESCRIPTION_LENGTH),
"published_date": video.get("published_date"),
"channel": {
"name": video.get("channel", {}).get("name"),
"link": video.get("channel", {}).get("link"),
},
"like_count": video.get("extracted_likes"),
"view_count": video.get("extracted_views"),
"live": video.get("live", False),
}