Sborole's picture
Update tools/YouTubeTool.py
9598d95 verified
import os
from smolagents import Tool
from googleapiclient.discovery import build
class YouTubeSearchTool(Tool):
"""
A tool to search for YouTube videos using the Google Custom Search Engine (CSE) API,
by specifically restricting results to the youtube.com domain.
"""
name = "youtube_search"
description = "Searches for and returns links and snippets for YouTube videos related to a specific query."
inputs = {
# Corrected the typo/broken line to read correctly
"query": {"type": "string", "description": "The search term to find relevant YouTube videos for."}
}
output_type = "string"
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Retrieve credentials from environment variables
self.api_key = os.getenv("GOOGLE_API_KEY")
self.cse_id = os.getenv("GOOGLE_CSE_ID")
if not self.api_key or not self.cse_id:
# Note: This check relies on environment variables being set up
raise ValueError("GOOGLE_API_KEY or GOOGLE_CSE_ID secret not found.")
# Initialize the Google Custom Search service
self.service = build(
"customsearch", "v1", developerKey=self.api_key
)
def forward(self, query: str) -> str:
"""
Executes a Google search query, restricting results to youtube.com.
Args:
query: The search term provided by the agent.
Returns:
A formatted string of search results (max 3), or an error message.
"""
print(f"Executing YouTube search for: '{query}'")
# Modify the query to ensure results are restricted to YouTube
youtube_query = f"{query} site:youtube.com"
try:
# Execute the search request for up to 3 results
res = self.service.cse().list(
q=youtube_query,
cx=self.cse_id,
num=3
).execute()
items = res.get('items', [])
if not items:
return "XX record info: No YouTube results found."
search_results = []
for i, item in enumerate(items):
search_results.append(
f"RESULT {i+1}: '{item.get('title')}'\n"
f"CONTENT: {item.get('snippet')}\n"
f"SOURCE: {item.get('link')}"
)
return "\n\n---SEPARATOR---\n\n".join(search_results)
except Exception as e:
return f"Error during YouTube Search API call: {e}"