File size: 1,745 Bytes
73e1a92 12bf7c7 73e1a92 12bf7c7 73e1a92 12bf7c7 73e1a92 12bf7c7 73e1a92 12bf7c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | from typing import Any, Optional
from smolagents.tools import Tool
import duckduckgo_search
from googleapiclient.discovery import build
import os
YOUTUBE_API_KEY=os.environ.get('YOUTUBE_API_KEY')
def youtube_search(query):
max_results=10
print(f"query -> {query}")
return_result = []
youtube = build('youtube', 'v3', developerKey=YOUTUBE_API_KEY)
request = youtube.search().list(
q=query,
part='snippet',
maxResults=max_results,
type='video'
)
response = request.execute()
for item in response.get('items', []):
print(f"Title: {item['snippet']['title']}")
print(f"Video ID: {item['id']['videoId']}")
print()
return_result.append({'aquery': query, 'titlevideo': item['snippet']['title'], 'videoId': item['id']['videoId']})
return return_result
class YoutubeSearchTool(Tool):
name = "youtube_search"
#description = "Performs a search with youtube API to select a potentially related song of band based on your query and return top X results."
description = "Performs a search on youtube to select a potentially related song of band based on your query and return top X results. Never use named parameter use on positional parameters"
inputs = {'aquery': {'type': 'string', 'description': 'The search query to perform.'}}
output_type = "object" # TODO or Dict or List ????!???
def __init__(self, max_results=10, **kwargs):
super().__init__()
self.max_results = max_results
def forward(self, aquery: str) -> str:
results = youtube_search(aquery)
if len(results) == 0:
raise Exception("No results found! Try a less restrictive/shorter query.")
return results |