"""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_workflow.pipeline.schemas import DownloadInput, DownloadOutput from kb_workflow.pipeline.processors.base import BaseProcessor if TYPE_CHECKING: from kb_workflow.pipeline.state import WorkflowStateManager from kb_workflow.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_workflow.pipeline.platforms.bilibili import BilibiliAdapter return BilibiliAdapter(state_manager=self.state_manager) elif platform == "zhihu": from kb_workflow.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 (author name, not mid) """ if platform == "bilibili": # Bilibili metadata structure: {metadata: {author: ..., mid: ...}, ...} if "metadata" in metadata: return str(metadata["metadata"].get("author", "")) return str(metadata.get("author", "")) 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 )