fix: update pytube to yt-dlp (#428)

Add tests and use yt-dlp instead of outdated and failing pytube package

Author: cmaloney111

pyproject.toml

@@ -28,7 +28,6 @@ dependencies = [
     "ipykernel>=6.29.4,<7",
     "tenacity>=8.3.0,<9",
     "pillow-heif>=0.16.0,<0.17",
-    "pytube==15.0.0",
     "anthropic>=0.31.0,<0.32",
     "pydantic>=2.0.0,<3",
     "av>=11.0.0,<12",
@@ -39,6 +38,7 @@ dependencies = [
     "dotenv>=0.9.9,<0.10",
     "pymupdf>=1.23.0,<2",
     "google-genai>=1.0.0,<2",
+    "yt-dlp>=2025.3.31",
 ]
 
 [project.urls]

tests/integ/test_tools.py

@@ -37,6 +37,7 @@
     vit_nsfw_classification,
 )
 
+
 def test_owlv2_object_detection():
     img = ski.data.coins()
     result = owlv2_object_detection(

tests/unit/tools/test_video.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from vision_agent.utils.video import extract_frames_from_video
+from vision_agent.tools import extract_frames_and_timestamps
 
 
 def test_extract_frames_from_video():
@@ -44,6 +45,28 @@ def test_extract_frames_with_input_video_has_no_fps():
     assert len(res) == 0
 
 
+def test_extract_frames_and_timestamps_from_local_video():
+    video_path = _create_video(duration=2)
+    res = extract_frames_and_timestamps(video_path, fps=24)
+    assert isinstance(res, list)
+    assert len(res) == 48
+    assert all("frame" in item and "timestamp" in item for item in res)
+
+
+def test_extract_frames_and_timestamps_from_http():
+    res = extract_frames_and_timestamps(
+        "https://www.w3schools.com/tags/mov_bbb.mp4", fps=0.2
+    )
+    assert isinstance(res, list)
+    assert len(res) == 2
+    assert all("frame" in item and "timestamp" in item for item in res)
+
+
+def test_extract_frames_and_timestamps_invalid_local_file():
+    res = extract_frames_and_timestamps("non_existing_file.mp4", fps=1.0)
+    assert res == []
+
+
 def _create_video(
     *, duration: int = 3, fps: int = 24, fps_video_prop: Optional[int] = 24
 ) -> str:

uv.lock

@@ -19,7 +19,7 @@
 from IPython.display import display
 from PIL import Image, ImageDraw, ImageFont
 from pillow_heif import register_heif_opener  # type: ignore
-from pytube import YouTube  # type: ignore
+import yt_dlp  # type: ignore
 import pymupdf  # type: ignore
 from google import genai  # type: ignore
 from google.genai import types  # type: ignore
@@ -3174,7 +3174,6 @@ def extract_frames_and_timestamps(
         [{"frame": np.ndarray, "timestamp": 0.0}, ...]
     """
     if isinstance(fps, str):
-        # fps could be a string when it's passed in from a web endpoint deployment
         fps = float(fps)
 
     def reformat(
@@ -3194,23 +3193,20 @@ def reformat(
         )
     ):
         with tempfile.TemporaryDirectory() as temp_dir:
-            yt = YouTube(str(video_uri))
-            # Download the highest resolution video
-            video = (
-                yt.streams.filter(progressive=True, file_extension="mp4")
-                .order_by("resolution")
-                .desc()
-                .first()
-            )
-            if not video:
-                raise Exception("No suitable video stream found")
-            video_file_path = video.download(output_path=temp_dir)
+            ydl_opts = {
+                "outtmpl": os.path.join(temp_dir, "%(title)s.%(ext)s"),
+                "format": "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best",
+                "quiet": True,
+            }
+            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+                info = ydl.extract_info(str(video_uri), download=True)
+                video_file_path = ydl.prepare_filename(info)
 
             return reformat(extract_frames_from_video(video_file_path, fps))
+
     elif str(video_uri).startswith(("http", "https")):
         _, image_suffix = os.path.splitext(video_uri)
         with tempfile.NamedTemporaryFile(delete=False, suffix=image_suffix) as tmp_file:
-            # Download the video and save it to the temporary file
             with urllib.request.urlopen(str(video_uri)) as response:
                 tmp_file.write(response.read())
             return reformat(extract_frames_from_video(tmp_file.name, fps))