refac

backend/open_webui/retrieval/loaders/youtube.py

@@ -7,14 +7,14 @@ from langchain_core.documents import Document
 
 log = logging.getLogger(__name__)
 
-ALLOWED_SCHEMES = {"http", "https"}
+ALLOWED_SCHEMES = {'http', 'https'}
 ALLOWED_NETLOCS = {
-    "youtu.be",
-    "m.youtube.com",
-    "youtube.com",
-    "www.youtube.com",
-    "www.youtube-nocookie.com",
-    "vid.plus",
+    'youtu.be',
+    'm.youtube.com',
+    'youtube.com',
+    'www.youtube.com',
+    'www.youtube-nocookie.com',
+    'vid.plus',
 }
 
 
@@ -30,17 +30,17 @@ def _parse_video_id(url: str) -> Optional[str]:
 
     path = parsed_url.path
 
-    if path.endswith("/watch"):
+    if path.endswith('/watch'):
         query = parsed_url.query
         parsed_query = parse_qs(query)
-        if "v" in parsed_query:
-            ids = parsed_query["v"]
+        if 'v' in parsed_query:
+            ids = parsed_query['v']
             video_id = ids if isinstance(ids, str) else ids[0]
         else:
             return None
     else:
-        path = parsed_url.path.lstrip("/")
-        video_id = path.split("/")[-1]
+        path = parsed_url.path.lstrip('/')
+        video_id = path.split('/')[-1]
 
     if len(video_id) != 11:  # Video IDs are 11 characters long
         return None
@@ -54,13 +54,13 @@ class YoutubeLoader:
     def __init__(
         self,
         video_id: str,
-        language: Union[str, Sequence[str]] = "en",
+        language: Union[str, Sequence[str]] = 'en',
         proxy_url: Optional[str] = None,
     ):
         """Initialize with YouTube video ID."""
         _video_id = _parse_video_id(video_id)
         self.video_id = _video_id if _video_id is not None else video_id
-        self._metadata = {"source": video_id}
+        self._metadata = {'source': video_id}
         self.proxy_url = proxy_url
 
         # Ensure language is a list
@@ -70,8 +70,8 @@ class YoutubeLoader:
             self.language = list(language)
 
         # Add English as fallback if not already in the list
-        if "en" not in self.language:
-            self.language.append("en")
+        if 'en' not in self.language:
+            self.language.append('en')
 
     def load(self) -> List[Document]:
         """Load YouTube transcripts into `Document` objects."""
@@ -85,14 +85,12 @@ class YoutubeLoader:
         except ImportError:
             raise ImportError(
                 'Could not import "youtube_transcript_api" Python package. '
-                "Please install it with `pip install youtube-transcript-api`."
+                'Please install it with `pip install youtube-transcript-api`.'
             )
 
         if self.proxy_url:
-            youtube_proxies = GenericProxyConfig(
-                http_url=self.proxy_url, https_url=self.proxy_url
-            )
-            log.debug(f"Using proxy URL: {self.proxy_url[:14]}...")
+            youtube_proxies = GenericProxyConfig(http_url=self.proxy_url, https_url=self.proxy_url)
+            log.debug(f'Using proxy URL: {self.proxy_url[:14]}...')
         else:
             youtube_proxies = None
 
@@ -100,7 +98,7 @@ class YoutubeLoader:
         try:
             transcript_list = transcript_api.list(self.video_id)
         except Exception as e:
-            log.exception("Loading YouTube transcript failed")
+            log.exception('Loading YouTube transcript failed')
             return []
 
         # Try each language in order of priority
@@ -110,14 +108,10 @@ class YoutubeLoader:
                 if transcript.is_generated:
                     log.debug(f"Found generated transcript for language '{lang}'")
                     try:
-                        transcript = transcript_list.find_manually_created_transcript(
-                            [lang]
-                        )
+                        transcript = transcript_list.find_manually_created_transcript([lang])
                         log.debug(f"Found manual transcript for language '{lang}'")
                     except NoTranscriptFound:
-                        log.debug(
-                            f"No manual transcript found for language '{lang}', using generated"
-                        )
+                        log.debug(f"No manual transcript found for language '{lang}', using generated")
                         pass
 
                 log.debug(f"Found transcript for language '{lang}'")
@@ -131,12 +125,10 @@ class YoutubeLoader:
                     log.debug(f"Empty transcript for language '{lang}'")
                     continue
 
-                transcript_text = " ".join(
+                transcript_text = ' '.join(
                     map(
                         lambda transcript_piece: (
-                            transcript_piece.text.strip(" ")
-                            if hasattr(transcript_piece, "text")
-                            else ""
+                            transcript_piece.text.strip(' ') if hasattr(transcript_piece, 'text') else ''
                         ),
                         transcript_pieces,
                     )
@@ -150,9 +142,9 @@ class YoutubeLoader:
                 raise e
 
         # If we get here, all languages failed
-        languages_tried = ", ".join(self.language)
+        languages_tried = ', '.join(self.language)
         log.warning(
-            f"No transcript found for any of the specified languages: {languages_tried}. Verify if the video has transcripts, add more languages if needed."
+            f'No transcript found for any of the specified languages: {languages_tried}. Verify if the video has transcripts, add more languages if needed.'
         )
         raise NoTranscriptFound(self.video_id, self.language, list(transcript_list))