starmen1265
60b0eaec93
- When content_type is video and no official subtitle available: - Download video using yt-dlp - Extract audio for Whisper fallback - Update content_file with actual video path This enables the pipeline to handle videos that don't have official Bilibili subtitles. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
339 lines
12 KiB
Python
339 lines
12 KiB
Python
"""DownloadProcessor - Content download processor
|
|
|
|
Handles metadata extraction and content downloading for all platforms.
|
|
Works with PlatformAdapter implementations (BilibiliAdapter, ZhihuAdapter).
|
|
|
|
Purpose: Unify metadata and content fetching into a single processor
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional, TYPE_CHECKING, Dict, Any
|
|
|
|
from kb_pipeline.schemas import DownloadInput, DownloadOutput
|
|
from .base import BaseProcessor
|
|
|
|
if TYPE_CHECKING:
|
|
from kb_pipeline.state import WorkflowStateManager
|
|
from kb_pipeline.platforms.base import PlatformAdapter
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DownloadProcessor(BaseProcessor):
|
|
"""Content download processor.
|
|
|
|
Fetches metadata and content from platforms using PlatformAdapter.
|
|
Supports video, article, answer, and zhuanlan content types.
|
|
|
|
Attributes:
|
|
STEP_NAME: "download" - identifies this processor step
|
|
"""
|
|
|
|
STEP_NAME = "download"
|
|
|
|
async def process(self, input_data: DownloadInput) -> DownloadOutput:
|
|
"""Process the download request.
|
|
|
|
Args:
|
|
input_data: DownloadInput with content_id, platform, content_dir, content_type
|
|
|
|
Returns:
|
|
DownloadOutput with metadata_file, content_file, title, author, etc.
|
|
"""
|
|
content_dir = Path(input_data.content_dir)
|
|
content_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Check if step already completed (idempotency)
|
|
if self.state_manager:
|
|
state = self.state_manager.load()
|
|
if state and state.is_step_completed(self.STEP_NAME):
|
|
logger.info(f"Step {self.STEP_NAME} already completed, skipping")
|
|
# Return a success output based on existing files
|
|
return self._get_existing_output(input_data, content_dir)
|
|
|
|
# Get appropriate adapter
|
|
try:
|
|
adapter = self._get_adapter(input_data.platform)
|
|
except ValueError as e:
|
|
return self._record_failure(str(e))
|
|
|
|
try:
|
|
# Fetch metadata (pass source_url for zhihu answers)
|
|
if input_data.platform == "zhihu":
|
|
metadata = await adapter.fetch_metadata(
|
|
input_data.content_id,
|
|
input_data.content_type,
|
|
source_url=input_data.source_url
|
|
)
|
|
else:
|
|
metadata = await adapter.fetch_metadata(
|
|
input_data.content_id,
|
|
input_data.content_type
|
|
)
|
|
|
|
# Save metadata
|
|
metadata_filename = adapter.get_output_filename(
|
|
input_data.content_id,
|
|
input_data.content_type,
|
|
"metadata.json"
|
|
)
|
|
metadata_path = content_dir / metadata_filename
|
|
self._save_metadata(metadata, metadata_path)
|
|
|
|
# Extract title and author from metadata
|
|
title = self._extract_title(metadata, input_data.platform)
|
|
author = self._extract_author(metadata, input_data.platform)
|
|
duration = self._extract_duration(metadata, input_data.content_type)
|
|
|
|
# Fetch content (for text-based content)
|
|
content_file = ""
|
|
if input_data.content_type in ["article", "answer", "zhuanlan"]:
|
|
if input_data.platform == "zhihu":
|
|
content = await adapter.fetch_content(
|
|
input_data.content_id,
|
|
input_data.content_type,
|
|
source_url=input_data.source_url
|
|
)
|
|
else:
|
|
content = await adapter.fetch_content(
|
|
input_data.content_id,
|
|
input_data.content_type
|
|
)
|
|
content_filename = adapter.get_output_filename(
|
|
input_data.content_id,
|
|
input_data.content_type,
|
|
"content.md"
|
|
)
|
|
content_path = content_dir / content_filename
|
|
self._save_content(content, content_path)
|
|
content_file = str(content_path)
|
|
|
|
# Download video for ASR (Whisper needs video file)
|
|
if input_data.content_type == "video" and input_data.platform == "bilibili":
|
|
download_result = await adapter.download_video(
|
|
input_data.content_id,
|
|
content_dir
|
|
)
|
|
if download_result["success"]:
|
|
content_file = download_result["filepath"]
|
|
logger.info(f"Video downloaded for ASR: {content_file}")
|
|
else:
|
|
logger.warning(f"Video download failed: {download_result['message']}")
|
|
|
|
# Record success
|
|
result = DownloadOutput(
|
|
success=True,
|
|
metadata_file=str(metadata_path),
|
|
content_file=content_file,
|
|
title=title,
|
|
author=author,
|
|
duration=duration
|
|
)
|
|
|
|
if self.state_manager:
|
|
state = self.state_manager.load()
|
|
if state:
|
|
summary = f"title={title}, type={input_data.content_type}"
|
|
state.mark_step_completed(self.STEP_NAME, summary)
|
|
self.state_manager.save(state)
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"Download failed: {e}")
|
|
return self._record_failure(str(e))
|
|
|
|
def _get_adapter(self, platform: str) -> PlatformAdapter:
|
|
"""Get the appropriate PlatformAdapter for the platform.
|
|
|
|
Args:
|
|
platform: Platform name ("bilibili" or "zhihu")
|
|
|
|
Returns:
|
|
PlatformAdapter instance
|
|
|
|
Raises:
|
|
ValueError: If platform is not supported
|
|
"""
|
|
if platform == "bilibili":
|
|
from kb_pipeline.platforms.bilibili import BilibiliAdapter
|
|
return BilibiliAdapter(state_manager=self.state_manager)
|
|
elif platform == "zhihu":
|
|
from kb_pipeline.platforms.zhihu import ZhihuAdapter
|
|
return ZhihuAdapter()
|
|
else:
|
|
raise ValueError(f"Unsupported platform: {platform}")
|
|
|
|
def _save_metadata(self, metadata: Dict[str, Any], path: Path) -> None:
|
|
"""Save metadata to JSON file.
|
|
|
|
Args:
|
|
metadata: Metadata dictionary
|
|
path: Output file path
|
|
"""
|
|
with open(path, 'w', encoding='utf-8') as f:
|
|
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
|
logger.info(f"Saved metadata to {path}")
|
|
|
|
def _save_content(self, content: str, path: Path) -> None:
|
|
"""Save content to Markdown file.
|
|
|
|
Args:
|
|
content: Markdown content string
|
|
path: Output file path
|
|
"""
|
|
with open(path, 'w', encoding='utf-8') as f:
|
|
f.write(content)
|
|
logger.info(f"Saved content to {path}")
|
|
|
|
def _extract_title(self, metadata: Dict[str, Any], platform: str) -> str:
|
|
"""Extract title from metadata.
|
|
|
|
Args:
|
|
metadata: Metadata dictionary
|
|
platform: Platform name
|
|
|
|
Returns:
|
|
Title string
|
|
"""
|
|
if platform == "bilibili":
|
|
# Bilibili metadata structure: {metadata: {title: ...}, ...}
|
|
if "metadata" in metadata:
|
|
return metadata["metadata"].get("title", "")
|
|
return metadata.get("title", "")
|
|
else:
|
|
return metadata.get("title", "")
|
|
|
|
def _extract_author(self, metadata: Dict[str, Any], platform: str) -> str:
|
|
"""Extract author from metadata.
|
|
|
|
Args:
|
|
metadata: Metadata dictionary
|
|
platform: Platform name
|
|
|
|
Returns:
|
|
Author string
|
|
"""
|
|
if platform == "bilibili":
|
|
# Bilibili metadata structure: {metadata: {mid: ...}, ...}
|
|
if "metadata" in metadata:
|
|
return str(metadata["metadata"].get("mid", ""))
|
|
return str(metadata.get("mid", ""))
|
|
else:
|
|
return metadata.get("author", "")
|
|
|
|
def _extract_duration(self, metadata: Dict[str, Any], content_type: str) -> Optional[int]:
|
|
"""Extract duration from metadata for video content.
|
|
|
|
Args:
|
|
metadata: Metadata dictionary
|
|
content_type: Content type
|
|
|
|
Returns:
|
|
Duration in seconds for videos, None otherwise
|
|
"""
|
|
if content_type != "video":
|
|
return None
|
|
|
|
# Bilibili video duration
|
|
if "metadata" in metadata and "duration" in metadata["metadata"]:
|
|
duration = metadata["metadata"]["duration"]
|
|
# Handle formatted duration string (e.g., "12:34" or "01:23:45")
|
|
if isinstance(duration, str):
|
|
if not duration.strip():
|
|
return None
|
|
parts = duration.split(':')
|
|
try:
|
|
if len(parts) == 2:
|
|
return int(parts[0]) * 60 + int(parts[1])
|
|
elif len(parts) == 3:
|
|
return int(parts[0]) * 3600 + int(parts[1]) * 60 + int(parts[2])
|
|
except (ValueError, IndexError):
|
|
return None
|
|
elif isinstance(duration, (int, float)):
|
|
return int(duration)
|
|
|
|
return None
|
|
|
|
def _record_failure(self, error_message: str) -> DownloadOutput:
|
|
"""Record failure to state manager and return failure output.
|
|
|
|
Args:
|
|
error_message: Error message
|
|
|
|
Returns:
|
|
DownloadOutput with success=False
|
|
"""
|
|
if self.state_manager:
|
|
state = self.state_manager.load()
|
|
if state:
|
|
state.mark_step_failed(self.STEP_NAME, error_message)
|
|
self.state_manager.save(state)
|
|
|
|
return DownloadOutput(
|
|
success=False,
|
|
error_message=error_message,
|
|
metadata_file="",
|
|
content_file="",
|
|
title="",
|
|
author=""
|
|
)
|
|
|
|
def _get_existing_output(
|
|
self,
|
|
input_data: DownloadInput,
|
|
content_dir: Path
|
|
) -> DownloadOutput:
|
|
"""Get output from existing files when step was already completed.
|
|
|
|
Args:
|
|
input_data: DownloadInput
|
|
content_dir: Content directory path
|
|
|
|
Returns:
|
|
DownloadOutput based on existing files
|
|
"""
|
|
adapter = self._get_adapter(input_data.platform)
|
|
|
|
metadata_filename = adapter.get_output_filename(
|
|
input_data.content_id,
|
|
input_data.content_type,
|
|
"metadata.json"
|
|
)
|
|
metadata_path = content_dir / metadata_filename
|
|
|
|
# Try to read metadata for title/author
|
|
title = ""
|
|
author = ""
|
|
duration = None
|
|
if metadata_path.exists():
|
|
try:
|
|
with open(metadata_path, 'r', encoding='utf-8') as f:
|
|
metadata = json.load(f)
|
|
title = self._extract_title(metadata, input_data.platform)
|
|
author = self._extract_author(metadata, input_data.platform)
|
|
duration = self._extract_duration(metadata, input_data.content_type)
|
|
except Exception:
|
|
pass
|
|
|
|
content_filename = adapter.get_output_filename(
|
|
input_data.content_id,
|
|
input_data.content_type,
|
|
"content.md"
|
|
)
|
|
content_path = content_dir / content_filename
|
|
content_file = str(content_path) if content_path.exists() else ""
|
|
|
|
return DownloadOutput(
|
|
success=True,
|
|
metadata_file=str(metadata_path),
|
|
content_file=content_file,
|
|
title=title,
|
|
author=author,
|
|
duration=duration
|
|
)
|