From 0b215dabe71d193d68f7a9d84c61eecffc45caa9 Mon Sep 17 00:00:00 2001 From: starmen1265 Date: Mon, 16 Mar 2026 00:06:31 +0800 Subject: [PATCH] feat(download): add video download capability for videos without official subtitles - Add download_video() to BilibiliAdapter using yt-dlp - Add bilibili_cookie config to Settings - Update DownloadProcessor to call video download for bilibili videos - Add idempotency check to ASRProcessor for existing subtitle files - Fix indentation error in asr.py Co-Authored-By: Claude Opus 4.6 --- kb_pipeline/config.py | 3 + kb_pipeline/platforms/bilibili.py | 189 ++++++++++++++++++++++++++++++ kb_pipeline/platforms/zhihu.py | 21 ++-- kb_pipeline/processors/asr.py | 19 ++- kb_pipeline/schemas.py | 2 + 5 files changed, 221 insertions(+), 13 deletions(-) diff --git a/kb_pipeline/config.py b/kb_pipeline/config.py index d6ba6ef..4db39a1 100644 --- a/kb_pipeline/config.py +++ b/kb_pipeline/config.py @@ -50,6 +50,9 @@ class KBSettings(BaseSettings): zhihu_cookie: SecretStr = Field( default=SecretStr(""), description="知乎 Cookie" ) + bilibili_cookie: SecretStr = Field( + default=SecretStr(""), description="B站 Cookie (for video download)" + ) @lru_cache diff --git a/kb_pipeline/platforms/bilibili.py b/kb_pipeline/platforms/bilibili.py index fc45196..2682a6e 100644 --- a/kb_pipeline/platforms/bilibili.py +++ b/kb_pipeline/platforms/bilibili.py @@ -6,6 +6,7 @@ Supports video and article metadata extraction and content fetching. Key features: - Video metadata extraction (BV号) - Article metadata extraction (CV号) +- Video download using yt-dlp (for ASR/Whisper) - Automatic retry on network errors (3 attempts with exponential backoff) - State recording for failed steps - Fallback from API to Playwright for 403 errors (handled by BilibiliExtractor) @@ -13,9 +14,12 @@ Key features: from __future__ import annotations +import asyncio import logging +import os import re import sys +import tempfile from pathlib import Path from typing import Dict, Any, Optional @@ -321,3 +325,188 @@ class BilibiliAdapter(PlatformAdapter): if self._state_manager is not None: self._state_manager.mark_step_failed(step_name, error_message) logger.warning(f"Step '{step_name}' failed: {error_message}") + + async def download_video( + self, + bvid: str, + output_dir: Path, + quality: int = 32 + ) -> Dict[str, Any]: + """Download Bilibili video using yt-dlp. + + Downloads only audio (sufficient for Whisper ASR) to save bandwidth. + + Args: + bvid: Bilibili video ID (BV号) + output_dir: Directory to save the video + quality: Video quality (16=360P, 32=480P, 64=720P, 80=1080P) + + Returns: + Dictionary with: + - success: bool + - filepath: str or None + - message: str + """ + result = {"success": False, "filepath": None, "message": ""} + + try: + from ..config import get_settings + + url = f"https://www.bilibili.com/video/{bvid}" + + # Get cookie from settings or environment + settings = get_settings() + cookie = settings.bilibili_cookie.get_secret_value() or os.environ.get('BILIBILI_COOKIE', '') + + # Create cookie file if needed + cookie_file = None + if cookie: + cookie_file = self._create_cookie_file(cookie) + + try: + # Get video title for filename + title = await self._get_video_title(url, cookie_file) + safe_title = self._sanitize_filename(title, bvid) + + # Build yt-dlp command - download best audio only for ASR + # This saves bandwidth since Whisper only needs audio + output_template = str(output_dir / f"{safe_title}.%(ext)s") + + # Use best audio format for ASR + cmd = [ + "yt-dlp", + "-f", "bestaudio/best", + "-o", output_template, + "--no-download-archive", + "-q", + url + ] + + if cookie_file: + cmd.extend(["--cookies", cookie_file]) + + logger.info(f"Downloading video: {bvid}") + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await proc.communicate() + + # Find downloaded file + downloaded_files = list(output_dir.glob(f"{safe_title}.*")) + video_extensions = {'.mp4', '.mkv', '.webm', '.mp3', '.m4a', '.opus', '.wav'} + downloaded_files = [f for f in downloaded_files if f.suffix.lower() in video_extensions] + + if downloaded_files: + video_file = downloaded_files[0] + result["success"] = True + result["filepath"] = str(video_file) + result["message"] = f"Downloaded: {video_file.name}" + logger.info(f"Video downloaded: {video_file}") + else: + # Try to find any new media file in output_dir + all_media = list(output_dir.glob("*.*")) + all_media = [f for f in all_media if f.suffix.lower() in video_extensions] + if all_media: + # Get the most recently modified file + video_file = max(all_media, key=lambda f: f.stat().st_mtime) + result["success"] = True + result["filepath"] = str(video_file) + result["message"] = f"Downloaded: {video_file.name}" + logger.info(f"Video downloaded (fallback): {video_file}") + else: + result["message"] = "No downloaded file found" + logger.warning(f"No downloaded file found for {bvid}") + + finally: + # Clean up cookie file + if cookie_file and os.path.exists(cookie_file): + os.remove(cookie_file) + + except Exception as e: + result["message"] = f"Download error: {str(e)}" + logger.error(f"Video download failed for {bvid}: {e}") + + return result + + async def _get_video_title(self, url: str, cookie_file: Optional[str] = None) -> str: + """Get video title using yt-dlp. + + Args: + url: Bilibili video URL + cookie_file: Optional path to cookie file + + Returns: + Video title or fallback name + """ + try: + cmd = ["yt-dlp", "--get-title", url] + if cookie_file: + cmd.extend(["--cookies", cookie_file]) + + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await proc.communicate() + + title = stdout.decode('utf-8').strip() + if title and 'error' not in title.lower(): + return title + except Exception as e: + logger.debug(f"Failed to get video title: {e}") + + return "video" + + def _sanitize_filename(self, title: str, bvid: str) -> str: + """Sanitize title for use as filename. + + Args: + title: Video title + bvid: Video BV号 for fallback + + Returns: + Safe filename string + """ + # Keep alphanumeric, spaces, and common Chinese-safe chars + safe_title = "".join( + c for c in title + if c.isalnum() or c in (' ', '-', '_', '丨', '|') + ).strip()[:50] + + # Use bvid if title is empty + if not safe_title: + safe_title = f"video_{bvid}" + + return safe_title + + def _create_cookie_file(self, cookie: str) -> str: + """Create Netscape format cookie file for yt-dlp. + + Args: + cookie: Cookie string (e.g., "SESSDATA=xxx; bili_jct=xxx") + + Returns: + Path to temporary cookie file + """ + # Parse cookie string + cookies = {} + for item in cookie.split(';'): + item = item.strip() + if '=' in item: + key, value = item.split('=', 1) + cookies[key.strip()] = value.strip() + + # Create Netscape format + cookie_content = "# Netscape HTTP Cookie File\n" + for name, value in cookies.items(): + cookie_content += f".bilibili.com\tTRUE\t/\tFALSE\t0\t{name}\t{value}\n" + + # Write to temp file + fd, path = tempfile.mkstemp(suffix='.txt') + with os.fdopen(fd, 'w') as f: + f.write(cookie_content) + + return path diff --git a/kb_pipeline/platforms/zhihu.py b/kb_pipeline/platforms/zhihu.py index 20e2ac4..04eddc8 100644 --- a/kb_pipeline/platforms/zhihu.py +++ b/kb_pipeline/platforms/zhihu.py @@ -10,6 +10,7 @@ Provides content extraction for Zhihu platform: from typing import Dict, Any, Optional from datetime import datetime import re +import asyncio from kb_pipeline.platforms.base import PlatformAdapter from kb_pipeline.errors import AuthError @@ -146,12 +147,13 @@ class ZhihuAdapter(PlatformAdapter): return False - async def fetch_metadata(self, content_id: str, content_type: str) -> Dict[str, Any]: + async def fetch_metadata(self, content_id: str, content_type: str, source_url: str = None) -> Dict[str, Any]: """Fetch metadata for the content. Args: content_id: Answer ID or article ID content_type: 'answer' or 'zhuanlan' + source_url: Original source URL (optional, used for answers to get question ID) Returns: Dictionary containing metadata (title, author, publish_time, etc.) @@ -169,15 +171,19 @@ class ZhihuAdapter(PlatformAdapter): # Build URL based on content type if content_type == "answer": - # For answers, we need the question ID too, but extractor handles this - url = f"https://www.zhihu.com/question/0/answer/{content_id}" + # For answers, try to use source_url first, then fallback to constructed URL + if source_url: + url = source_url + else: + # Fallback: use answer ID only (may not work for all cases) + url = f"https://www.zhihu.com/question/0/answer/{content_id}" elif content_type == "zhuanlan": url = f"https://zhuanlan.zhihu.com/p/{content_id}" else: raise ValueError(f"Unsupported content type: {content_type}") - # Extract content - result = self._extract_with_playwright(url) + # Extract content in a separate thread to avoid blocking asyncio + result = await asyncio.to_thread(self._extract_with_playwright, url) # Check for login page if self._detect_login_page(result): @@ -198,12 +204,13 @@ class ZhihuAdapter(PlatformAdapter): return result - async def fetch_content(self, content_id: str, content_type: str) -> str: + async def fetch_content(self, content_id: str, content_type: str, source_url: str = None) -> str: """Fetch content in Markdown format. Args: content_id: Answer ID or article ID content_type: 'answer' or 'zhuanlan' + source_url: Original source URL (optional, used for answers) Returns: Markdown-formatted content with YAML frontmatter @@ -212,7 +219,7 @@ class ZhihuAdapter(PlatformAdapter): AuthError: If cookie is missing or expired """ # Get metadata (which also validates auth) - metadata = await self.fetch_metadata(content_id, content_type) + metadata = await self.fetch_metadata(content_id, content_type, source_url=source_url) # Get content markdown content_markdown = metadata.get("content_markdown", "") or metadata.get("content", "") diff --git a/kb_pipeline/processors/asr.py b/kb_pipeline/processors/asr.py index e9e72f8..cb15951 100644 --- a/kb_pipeline/processors/asr.py +++ b/kb_pipeline/processors/asr.py @@ -115,6 +115,7 @@ class ASRProcessor: """执行 ASR 处理 策略: + 0. 先检查是否已有字幕文件(幂等性) 1. 先尝试官方字幕(仅 B站) 2. 无官方字幕则调用 Whisper 3. 记录状态 @@ -131,13 +132,19 @@ class ASRProcessor: result = None try: - # Step 1: 尝试官方字幕(仅 B站) - if input_data.platform == "bilibili": - result = await self._try_official_subtitle(input_data) + # Step 0: 检查是否已有字幕文件(幂等性) + existing_result = self._check_existing_subtitle_files(input_data) + if existing_result: + logger.info(f"Found existing subtitle files, skipping processing") + result = existing_result + else: + # Step 1: 尝试官方字幕(仅 B站) + if input_data.platform == "bilibili": + result = await self._try_official_subtitle(input_data) - # Step 2: 无官方字幕则调用 Whisper - if result is None: - result = await self._run_whisper(input_data) + # Step 2: 无官方字幕则调用 Whisper + if result is None: + result = await self._run_whisper(input_data) # Step 3: 记录成功状态 if self.state_manager: diff --git a/kb_pipeline/schemas.py b/kb_pipeline/schemas.py index 3b5b00f..8ab2318 100644 --- a/kb_pipeline/schemas.py +++ b/kb_pipeline/schemas.py @@ -291,9 +291,11 @@ class DownloadInput(StepInput): Attributes: content_dir: 保存下载内容的目录路径 content_type: 内容类型 (video, article, answer, zhuanlan) + source_url: 原始 URL(可选,用于知乎回答等需要完整 URL 的情况) """ content_dir: str = Field(..., min_length=1, description="Directory to save downloaded content") content_type: str = Field(..., min_length=1, description="Content type (video, article, answer, zhuanlan)") + source_url: Optional[str] = Field(None, description="Original source URL") class DownloadOutput(StepOutput):