Add Youtube tool

agent.py

@@ -16,8 +16,9 @@ from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode
 from langgraph.prebuilt import tools_condition
 from mediawikiapi import MediaWikiAPI
-from transformers import pipeline
 from wikipedia_tool import WikipediaTool
+from yt_tool import speech_recognition_pipe, yt_transcribe
+
 
 @tool
 def read_xlsx_file(file_path: str) -> str:
@@ -83,7 +84,8 @@ class Agent:
             TavilySearch(),
             read_xlsx_file,
             addition,
-            multiple
+            multiple,
+            yt_transcribe
         ]
 
         self.llm_with_tools = llm.bind_tools(self.tools)
@@ -198,12 +200,7 @@ class Agent:
             temp_dir = tempfile.gettempdir()  # Get the temporary directory path
             audio_path = os.path.join(temp_dir, os.path.basename(filename))
 
-            pipe = pipeline(
-                task="automatic-speech-recognition",
-                model="openai/whisper-large-v3"
-            )
-
-            result = pipe(audio_path)
+            result = speech_recognition_pipe(audio_path)
 
             audio_message = HumanMessage(result["text"])
 

requirements.txt

@@ -16,4 +16,5 @@ openpyxl
 protobuf~=5.29.4
 genai~=2.1.0
 transformers~=4.52.4
-torch
+torch
+yt-dlp

yt_tool.py

@@ -0,0 +1,93 @@
+import torch
+import yt_dlp as youtube_dl
+from transformers import pipeline
+from transformers.pipelines.audio_utils import ffmpeg_read
+from langchain_core.tools import tool
+
+import tempfile
+import os
+
+# credit https://huggingface.co/spaces/hf-audio/whisper-large-v3
+
+MODEL_NAME = "openai/whisper-tiny.en"
+BATCH_SIZE = 8
+FILE_LIMIT_MB = 1000
+YT_LENGTH_LIMIT_S = 3600  # limit to 1 hour YouTube files
+
+device = "mps" if torch.mps.is_available() else "cpu"
+
+speech_recognition_pipe = pipeline(
+    task="automatic-speech-recognition",
+    model=MODEL_NAME,
+    chunk_length_s=30,
+    device=device,
+)
+
+def download_yt_audio(yt_url, filename):
+    info_loader = youtube_dl.YoutubeDL()
+
+    try:
+        info = info_loader.extract_info(yt_url, download=False)
+    except youtube_dl.utils.DownloadError as err:
+        raise str(err)
+
+    file_length = info["duration_string"]
+    file_h_m_s = file_length.split(":")
+    file_h_m_s = [int(sub_length) for sub_length in file_h_m_s]
+
+    if len(file_h_m_s) == 1:
+        file_h_m_s.insert(0, 0)
+    if len(file_h_m_s) == 2:
+        file_h_m_s.insert(0, 0)
+    file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2]
+
+    if file_length_s > YT_LENGTH_LIMIT_S:
+        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
+        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
+        raise f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video."
+
+    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
+
+    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+        try:
+            ydl.download([yt_url])
+        except youtube_dl.utils.ExtractorError as err:
+            raise str(err)
+
+def _return_yt_html_embed(yt_url):
+    video_id = yt_url.split("?v=")[-1]
+    HTML_str = (
+        f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
+        " </center>"
+    )
+    return HTML_str
+
+@tool
+def yt_transcribe(yt_url, max_filesize=75.0):
+    """
+        Transcribes the audio from a given YouTube video URL.
+
+        Args:
+            yt_url (str): The URL of the YouTube video.
+            max_filesize (float, optional): The maximum allowed filesize of the video in MB.
+                                             Defaults to 75.0.
+
+        Returns:
+            tuple: A tuple containing:
+                   - str: An HTML embed string for the YouTube video.
+                   - str: The transcribed text of the video's audio.
+        """
+    html_embed_str = _return_yt_html_embed(yt_url)
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        filepath = os.path.join(tmpdirname, "video.mp4")
+        download_yt_audio(yt_url, filepath)
+        with open(filepath, "rb") as f:
+            inputs = f.read()
+
+    inputs = ffmpeg_read(inputs, speech_recognition_pipe.feature_extractor.sampling_rate)
+    inputs = {"array": inputs, "sampling_rate": speech_recognition_pipe.feature_extractor.sampling_rate}
+
+    text = speech_recognition_pipe(inputs, batch_size=8, return_timestamps=True)["text"]
+
+    return html_embed_str, text