""" YouTube pull + Whisper transcribe adapter. Ported from xtract/main.py (VideoDownloader + SubtitleExtractor classes). Adapted to work with the second-brain Source model and Config. """ from __future__ import annotations from datetime import datetime from pathlib import Path from typing import Optional import whisper import yt_dlp from second_brain.config import Config from second_brain.models import Source, SourceStatus class VideoDownloader: """Download a YouTube video (or playlist) via yt-dlp. Ported from xtract VideoDownloader — same core yt-dlp options, adapted to write into config.media_dir. """ def __init__(self, config: Config) -> None: self.config = config def download(self, url: str) -> list[Path]: """Download video(s) and return paths to the downloaded files.""" config = self.config config.media_dir.mkdir(parents=True, exist_ok=True) ydl_opts = { "format": "bestvideo[height<=720]+bestaudio/best[height<=720]", "outtmpl": str(config.media_dir / "%(id)s.%(ext)s"), "writeinfojson": True, "ignoreerrors": True, "no_warnings": False, "quiet": False, } downloaded: list[Path] = [] with yt_dlp.YoutubeDL(ydl_opts) as ydl: print(f"[yt-dlp] Fetching: {url}") info = ydl.extract_info(url, download=True) if info is None: print(f"[yt-dlp] Error: could not extract info from {url}") return downloaded if "entries" in info: for entry in info["entries"]: if entry: p = self._find_media_file(entry.get("id", "")) if p: downloaded.append(p) else: p = self._find_media_file(info.get("id", "")) if p: downloaded.append(p) return downloaded def _find_media_file(self, video_id: str) -> Optional[Path]: for f in self.config.media_dir.glob(f"{video_id}.*"): if f.suffix in {".mp4", ".mkv", ".webm", ".m4a"}: return f return None def extract_metadata(self, url: str) -> dict: """Fetch video metadata without downloading the media file.""" ydl_opts = {"quiet": True, "skip_download": True} with yt_dlp.YoutubeDL(ydl_opts) as ydl: info = ydl.extract_info(url, download=False) or {} return { "title": info.get("title"), "author": info.get("uploader"), "duration_seconds": info.get("duration"), "published_at": _parse_upload_date(info.get("upload_date")), } class SubtitleExtractor: """Transcribe a video file to an SRT subtitle file using OpenAI Whisper. Ported from xtract SubtitleExtractor — identical core logic, writes SRTs to config.subtitles_dir. """ def __init__(self, config: Config) -> None: self.config = config self._model = None def _load_model(self) -> None: if self._model is None: print(f"[Whisper] Loading {self.config.whisper_model} model…") self._model = whisper.load_model(self.config.whisper_model) print("[Whisper] Model ready") def extract(self, media_file: Path) -> Optional[Path]: """Transcribe *media_file* and return the path to the SRT file.""" self.config.subtitles_dir.mkdir(parents=True, exist_ok=True) srt_file = self.config.subtitles_dir / f"{media_file.stem}.srt" if srt_file.exists(): print(f"[Whisper] Already transcribed: {srt_file.name}") return srt_file self._load_model() print(f"[Whisper] Transcribing: {media_file.name}") try: result = self._model.transcribe( str(media_file), verbose=False, language=None, ) self._write_srt(result, srt_file) print(f"[Whisper] Saved: {srt_file.name}") return srt_file except Exception as exc: print(f"[Whisper] Error transcribing {media_file.name}: {exc}") return None def extract_text(self, media_file: Path) -> Optional[str]: """Return the plain-text transcript (no SRT formatting).""" self._load_model() try: result = self._model.transcribe(str(media_file), verbose=False) return " ".join(seg["text"].strip() for seg in result["segments"]) except Exception as exc: print(f"[Whisper] Error: {exc}") return None @staticmethod def _write_srt(result: dict, output_file: Path) -> None: with open(output_file, "w", encoding="utf-8") as fh: for i, seg in enumerate(result["segments"], start=1): start = _fmt_timestamp(seg["start"]) end = _fmt_timestamp(seg["end"]) fh.write(f"{i}\n{start} --> {end}\n{seg['text'].strip()}\n\n") @staticmethod def read_srt(srt_file: Path) -> str: """Read an SRT file and return its raw text content.""" return srt_file.read_text(encoding="utf-8") # --------------------------------------------------------------------------- # High-level adapter # --------------------------------------------------------------------------- class YouTubeAdapter: """Orchestrates pull + transcribe for a YouTube Source record.""" def __init__(self, config: Config) -> None: self.config = config self.downloader = VideoDownloader(config) self.transcriber = SubtitleExtractor(config) def pull(self, source: Source) -> bool: """Download the video. Updates source in-place, returns success.""" try: meta = self.downloader.extract_metadata(source.url) source.title = source.title or meta.get("title") source.author = meta.get("author") source.duration_seconds = meta.get("duration_seconds") source.published_at = meta.get("published_at") files = self.downloader.download(source.url) if not files: source.error_message = "yt-dlp returned no files" source.status = SourceStatus.FAILED return False source.media_path = str(files[0]) source.status = SourceStatus.PULLED return True except Exception as exc: source.error_message = str(exc) source.status = SourceStatus.FAILED return False def transcribe(self, source: Source) -> bool: """Transcribe the downloaded video. Updates source in-place.""" if not source.media_path: source.error_message = "No media_path; run pull first" return False media_file = Path(source.media_path) srt_file = self.transcriber.extract(media_file) if srt_file is None: source.error_message = "Whisper transcription failed" source.status = SourceStatus.FAILED return False source.transcript_path = str(srt_file) source.transcript_text = SubtitleExtractor.read_srt(srt_file) source.status = SourceStatus.TRANSCRIBED return True # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _fmt_timestamp(seconds: float) -> str: h = int(seconds // 3600) m = int((seconds % 3600) // 60) s = int(seconds % 60) ms = int((seconds % 1) * 1000) return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}" def _parse_upload_date(upload_date: Optional[str]) -> Optional[datetime]: """Parse yt-dlp's YYYYMMDD upload_date string.""" if not upload_date or len(upload_date) != 8: return None try: return datetime.strptime(upload_date, "%Y%m%d") except ValueError: return None __all__ = ["VideoDownloader", "SubtitleExtractor", "YouTubeAdapter"]