|
from typing import Any, Dict, Generator, List, Optional, Sequence, Union |
|
from urllib.parse import parse_qs, urlparse |
|
from langchain_core.documents import Document |
|
|
|
|
|
ALLOWED_SCHEMES = {"http", "https"} |
|
ALLOWED_NETLOCS = { |
|
"youtu.be", |
|
"m.youtube.com", |
|
"youtube.com", |
|
"www.youtube.com", |
|
"www.youtube-nocookie.com", |
|
"vid.plus", |
|
} |
|
|
|
|
|
def _parse_video_id(url: str) -> Optional[str]: |
|
"""Parse a YouTube URL and return the video ID if valid, otherwise None.""" |
|
parsed_url = urlparse(url) |
|
|
|
if parsed_url.scheme not in ALLOWED_SCHEMES: |
|
return None |
|
|
|
if parsed_url.netloc not in ALLOWED_NETLOCS: |
|
return None |
|
|
|
path = parsed_url.path |
|
|
|
if path.endswith("/watch"): |
|
query = parsed_url.query |
|
parsed_query = parse_qs(query) |
|
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] |
|
|
|
if len(video_id) != 11: |
|
return None |
|
|
|
return video_id |
|
|
|
|
|
class YoutubeLoader: |
|
"""Load `YouTube` video transcripts.""" |
|
|
|
def __init__( |
|
self, |
|
video_id: str, |
|
language: Union[str, Sequence[str]] = "en", |
|
): |
|
"""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.language = language |
|
if isinstance(language, str): |
|
self.language = [language] |
|
else: |
|
self.language = language |
|
|
|
def load(self) -> List[Document]: |
|
"""Load YouTube transcripts into `Document` objects.""" |
|
try: |
|
from youtube_transcript_api import ( |
|
NoTranscriptFound, |
|
TranscriptsDisabled, |
|
YouTubeTranscriptApi, |
|
) |
|
except ImportError: |
|
raise ImportError( |
|
'Could not import "youtube_transcript_api" Python package. ' |
|
"Please install it with `pip install youtube-transcript-api`." |
|
) |
|
|
|
try: |
|
transcript_list = YouTubeTranscriptApi.list_transcripts(self.video_id) |
|
except Exception as e: |
|
print(e) |
|
return [] |
|
|
|
try: |
|
transcript = transcript_list.find_transcript(self.language) |
|
except NoTranscriptFound: |
|
transcript = transcript_list.find_transcript(["en"]) |
|
|
|
transcript_pieces: List[Dict[str, Any]] = transcript.fetch() |
|
|
|
transcript = " ".join( |
|
map( |
|
lambda transcript_piece: transcript_piece["text"].strip(" "), |
|
transcript_pieces, |
|
) |
|
) |
|
return [Document(page_content=transcript, metadata=self._metadata)] |
|
|
