diff --git a/kb_pipeline/processors/download.py b/kb_pipeline/processors/download.py new file mode 100644 index 0000000..55494e8 --- /dev/null +++ b/kb_pipeline/processors/download.py @@ -0,0 +1,298 @@ +"""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 + 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"]: + 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) + + # 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"]: + return metadata["metadata"]["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 + ) diff --git a/tests/test_processors/test_download.py b/tests/test_processors/test_download.py index dc1f0e3..d87542b 100644 --- a/tests/test_processors/test_download.py +++ b/tests/test_processors/test_download.py @@ -161,6 +161,7 @@ class TestDownloadProcessorMetadata: "url": "https://www.zhihu.com/question/123/answer/456", "content_markdown": "Test content" }) + adapter.fetch_content = AsyncMock(return_value="---\ntitle: Test\n---\nTest content") adapter.get_output_filename = MagicMock( return_value="zhihu_456_metadata.json" ) @@ -189,11 +190,13 @@ class TestDownloadProcessorMetadata: content_type="video" ) - with patch.object( - DownloadProcessor, '_get_adapter', return_value=mock_bilibili_adapter - ): - processor = DownloadProcessor(state_manager=mock_state_manager) - result = await processor.process(input_data) + with patch("kb_pipeline.processors.base.get_settings") as mock_settings: + mock_settings.return_value = MagicMock() + with patch.object( + DownloadProcessor, '_get_adapter', return_value=mock_bilibili_adapter + ): + processor = DownloadProcessor(state_manager=mock_state_manager) + result = await processor.process(input_data) assert result.success is True assert result.title == "Test Video Title" @@ -219,11 +222,13 @@ class TestDownloadProcessorMetadata: content_type="answer" ) - with patch.object( - DownloadProcessor, '_get_adapter', return_value=mock_zhihu_adapter - ): - processor = DownloadProcessor(state_manager=mock_state_manager) - result = await processor.process(input_data) + with patch("kb_pipeline.processors.base.get_settings") as mock_settings: + mock_settings.return_value = MagicMock() + with patch.object( + DownloadProcessor, '_get_adapter', return_value=mock_zhihu_adapter + ): + processor = DownloadProcessor(state_manager=mock_state_manager) + result = await processor.process(input_data) assert result.success is True assert result.title == "Test Answer Title" @@ -253,11 +258,13 @@ class TestDownloadProcessorMetadata: content_type="video" ) - with patch.object( - DownloadProcessor, '_get_adapter', return_value=mock_bilibili_adapter - ): - processor = DownloadProcessor(state_manager=mock_state_manager) - result = await processor.process(input_data) + with patch("kb_pipeline.processors.base.get_settings") as mock_settings: + mock_settings.return_value = MagicMock() + with patch.object( + DownloadProcessor, '_get_adapter', return_value=mock_bilibili_adapter + ): + processor = DownloadProcessor(state_manager=mock_state_manager) + result = await processor.process(input_data) assert result.success is True @@ -281,6 +288,7 @@ class TestDownloadProcessorMetadata: side_effect=NetworkError( step_name="download", input_summary="BV1xx", + status_code=503, cause=Exception("Network timeout") ) ) @@ -297,14 +305,16 @@ class TestDownloadProcessorMetadata: content_type="video" ) - with patch.object( - DownloadProcessor, '_get_adapter', return_value=mock_bilibili_adapter - ): - processor = DownloadProcessor(state_manager=mock_state_manager) - result = await processor.process(input_data) + with patch("kb_pipeline.processors.base.get_settings") as mock_settings: + mock_settings.return_value = MagicMock() + with patch.object( + DownloadProcessor, '_get_adapter', return_value=mock_bilibili_adapter + ): + processor = DownloadProcessor(state_manager=mock_state_manager) + result = await processor.process(input_data) assert result.success is False - assert "Network timeout" in result.error_message + assert "503" in result.error_message # NetworkError message includes status code # Verify failure was recorded mock_state.mark_step_failed.assert_called_once()