"""Pipeline engine - orchestrates multiple processors The Pipeline class coordinates the workflow: 1. Parse URL to extract platform and content ID 2. Run processors in order (parse, download, asr, analyze) 3. Handle errors and state persistence 4. Return results Key design decisions: - Sequential execution (simpler, easier to debug) - State persistence after each step (via WorkflowStateManager) - Simple progress logging """ from __future__ import annotations import logging from pathlib import Path from typing import Dict, List, Optional from kb_workflow.pipeline.config import get_settings from kb_workflow.pipeline.link_parser import LinkParser, ParsedLink from kb_workflow.pipeline.processors.base import BaseProcessor from kb_workflow.pipeline.schemas import PipelineConfig, PipelineInput, PipelineOutput, SRTInput logger = logging.getLogger(__name__) class Pipeline: """Pipeline engine - orchestrates multiple processors Orchestrates the workflow: 1. Parse URL to extract platform and content ID 2. Run processors in order (parse, download, asr, analyze) 3. Handle errors and state persistence 4. Return results Key design decisions: - Sequential execution (simpler, easier to debug) - State persistence after each step (via WorkflowStateManager) - Simple progress logging """ def __init__( self, processors: List[BaseProcessor], config: Optional[PipelineConfig] = None, output_dir: Optional[Path] = None, ): """Initialize the Pipeline. Args: processors: List of processor instances to use config: Optional PipelineConfig for pipeline settings output_dir: Optional output directory for content Raises: ValueError: If a processor's STEP_NAME is not in config.steps """ self.processors: Dict[str, BaseProcessor] = {p.STEP_NAME: p for p in processors} self.config = config or PipelineConfig() self.settings = get_settings() self.output_dir = output_dir or Path(self.settings.intermediate_path) # Validate processors - each processor's STEP_NAME must be in config.steps for step_name in self.processors: if step_name not in self.config.steps: raise ValueError(f"Unknown processor: {step_name}") def _parse_url(self, url: str) -> Optional[ParsedLink]: """Parse URL using LinkParser. Args: url: URL to parse Returns: ParsedLink with extracted information, or None if URL is not recognized """ if not url: return None parser = LinkParser() parsed = parser.parse(url) if not parsed: logger.warning(f"Failed to parse URL: {url}") return None return parsed def _get_content_dir(self, parsed: ParsedLink) -> Path: """Get content directory from parsed URL. Creates the directory if it doesn't exist. Args: parsed: ParsedLink with platform and content_id Returns: Path to the content directory """ content_dir = self.output_dir / f"{parsed.platform}_{parsed.content_id}" content_dir.mkdir(parents=True, exist_ok=True) return content_dir async def run(self, url: str) -> PipelineOutput: """Run the pipeline for a given URL. Executes all configured processors in order: 1. Parse URL to get platform and content_id 2. Create content directory 3. Run each processor sequentially 4. Handle errors and state persistence 5. Return results Args: url: URL to process Returns: PipelineOutput with success status, results, and content_dir """ # 1. Parse URL parsed = self._parse_url(url) if not parsed: return PipelineOutput( success=False, error_message=f"Failed to parse URL: {url}", results={}, content_dir=None, ) # 2. Create content directory content_dir = self._get_content_dir(parsed) # 3. Initialize state manager from kb_workflow.pipeline.state import WorkflowStateManager state_manager = WorkflowStateManager(content_dir) # Get or create workflow state state = state_manager.get_or_create( content_id=parsed.content_id, platform=parsed.platform, source_url=url, ) # 4. Run processors in order results: Dict = {} for step_name in self.config.steps: if step_name not in self.processors: logger.warning(f"Step {step_name} not found in processors, skipping") continue # Check if step should be skipped if self.config.skip_completed and state.is_step_completed(step_name): # Ask user whether to re-run the completed step should_rerun = self._ask_rerun_step(step_name, content_dir) if should_rerun: logger.info(f"Re-running completed step: {step_name}") # Clear the step's state record so it can be re-executed if step_name in state.steps: del state.steps[step_name] state_manager.save(state) else: logger.info(f"Skipping completed step: {step_name}") # Scan existing files and populate results for skipped steps self._populate_skipped_step_results(step_name, results, content_dir, parsed) continue # Skip ASR for text-based content types text_content_types = {"article", "answer", "zhuanlan", "question"} if step_name == "asr" and parsed.content_type in text_content_types: logger.info(f"Skipping ASR for text content type: {parsed.content_type}") continue # B站视频：跳过 analyze 步骤，使用 srt 步骤替代 is_bilibili_video = ( parsed.platform == "bilibili" and parsed.content_type == "video" ) if step_name == "analyze" and is_bilibili_video: logger.info(f"Skipping analyze for Bilibili video, using SRT processor instead") continue # B站文章/知乎等：跳过 srt 步骤 if step_name == "srt" and not is_bilibili_video: logger.info(f"Skipping SRT for non-Bilibili-video content: {parsed.platform}/{parsed.content_type}") continue processor = self.processors[step_name] # Mark step as started state.mark_step_started(step_name) state_manager.save(state) try: # Create step input based on step type from kb_workflow.pipeline.schemas import ( ParseInput, DownloadInput, ASRInput, AnalyzeInput, ArchiveInput ) if step_name == "parse": step_input = ParseInput( url=url, content_id=parsed.content_id, platform=parsed.platform, ) elif step_name == "download": step_input = DownloadInput( content_id=parsed.content_id, platform=parsed.platform, content_dir=str(content_dir), content_type=parsed.content_type, source_url=url, ) elif step_name == "asr": # ASR needs video path from download step download_result = results.get("download", {}) video_path = download_result.get("content_file", "") step_input = ASRInput( content_id=parsed.content_id, platform=parsed.platform, video_path=video_path, content_type=parsed.content_type, content_dir=str(content_dir), ) elif step_name == "analyze": # Analyze needs metadata and content files download_result = results.get("download", {}) asr_result = results.get("asr", {}) metadata_file = download_result.get("metadata_file", "") # Use subtitle JSON if available, otherwise content file content_file = asr_result.get("output_files", {}).get("json", "") or download_result.get("content_file", "") step_input = AnalyzeInput( content_id=parsed.content_id, platform=parsed.platform, metadata_file=metadata_file, content_file=content_file, content_type=parsed.content_type, output_file=str(content_dir / "analysis.md"), ) elif step_name == "srt": # SRT needs srt file from ASR step and metadata download_result = results.get("download", {}) asr_result = results.get("asr", {}) metadata_file = download_result.get("metadata_file", "") # SRT file from ASR output srt_file = asr_result.get("output_files", {}).get("srt", "") step_input = SRTInput( content_id=parsed.content_id, platform=parsed.platform, srt_file=srt_file, metadata_file=metadata_file, output_file=str(content_dir / "analysis.md"), content_type=parsed.content_type, ) elif step_name == "archive": # Archive needs analysis file from analyze step (or srt step for Bilibili video) analyze_result = results.get("analyze", {}) or results.get("srt", {}) analysis_file = analyze_result.get("analysis_file", "") or analyze_result.get("output_file", "") step_input = ArchiveInput( content_id=parsed.content_id, platform=parsed.platform, content_dir=str(content_dir), analysis_file=analysis_file, content_type=parsed.content_type, ) else: # Generic step input from kb_workflow.pipeline.schemas import StepInput step_input = StepInput( content_id=parsed.content_id, platform=parsed.platform, ) # Run processor step_output = await processor.process(step_input) # Record result results[step_name] = { "success": step_output.success, "error_message": step_output.error_message, } if step_output.success: # Add output fields to results for field_name in ["metadata_file", "content_file", "title", "author", "analysis_file", "summary", "output_files", "archive_file", "skipped", "output_file", "paragraphs", "tags", "corrected"]: if hasattr(step_output, field_name): results[step_name][field_name] = getattr(step_output, field_name) state.mark_step_completed(step_name) else: state.mark_step_failed(step_name, step_output.error_message or "Unknown error") # Stop pipeline on error state_manager.save(state) return PipelineOutput( success=False, error_message=f"Step {step_name} failed: {step_output.error_message}", results=results, content_dir=content_dir, ) except Exception as e: logger.error(f"Error in step {step_name}: {e}") state.mark_step_failed(step_name, str(e)) state_manager.save(state) return PipelineOutput( success=False, error_message=f"Step {step_name} raised exception: {e}", results=results, content_dir=content_dir, ) # Save state after each successful step state_manager.save(state) # 5. Return success return PipelineOutput( success=True, error_message=None, results=results, content_dir=content_dir, ) def _ask_rerun_step(self, step_name: str, content_dir: Path) -> bool: """Ask user whether to re-run a completed step. Args: step_name: Name of the step to potentially re-run content_dir: Content directory path Returns: True if user wants to re-run, False to skip """ try: from rich.console import Console from rich.prompt import Confirm console = Console() return Confirm.ask( f"Step '{step_name}' is already completed. Re-run it?", default=False ) except ImportError: # Fallback to basic input if rich is not available response = input(f"Step '{step_name}' is already completed. Re-run it? [y/N] ") return response.lower() in ('y', 'yes') def _populate_skipped_step_results( self, step_name: str, results: Dict, content_dir: Path, parsed: ParsedLink ) -> None: """Populate results for skipped steps by scanning existing files. This ensures downstream steps can find the outputs they need, even when a step is skipped. Args: step_name: Name of the skipped step results: Results dictionary to populate content_dir: Content directory path parsed: Parsed link with platform and content_id """ if step_name == "asr": # Scan for subtitle files generated by ASR # Try both new naming convention (with _subtitle suffix) and old convention (without suffix) base_name_new = f"{parsed.platform}_{parsed.content_id}_subtitle" base_name_old = f"{parsed.platform}_{parsed.content_id}" json_path = content_dir / f"{base_name_new}.json" if not json_path.exists(): json_path = content_dir / f"{base_name_old}.json" srt_path = content_dir / f"{base_name_new}.srt" if not srt_path.exists(): srt_path = content_dir / f"{base_name_old}.srt" output_files = {} if json_path.exists(): output_files["json"] = str(json_path) if srt_path.exists(): output_files["srt"] = str(srt_path) if output_files: results["asr"] = { "success": True, "error_message": None, "output_files": output_files, } logger.info(f"Found existing ASR outputs: {output_files}") else: logger.warning(f"No existing ASR outputs found in {content_dir}") results["asr"] = { "success": False, "error_message": "No existing subtitle files found", "output_files": {}, } elif step_name == "download": # Scan for metadata file metadata_path = content_dir / "metadata.json" if metadata_path.exists(): results["download"] = { "success": True, "error_message": None, "metadata_file": str(metadata_path), } logger.info(f"Found existing download metadata: {metadata_path}") elif step_name == "analyze": # Scan for analysis output analysis_path = content_dir / "analysis.md" if analysis_path.exists(): results["analyze"] = { "success": True, "error_message": None, "analysis_file": str(analysis_path), } logger.info(f"Found existing analysis: {analysis_path}") elif step_name == "srt": # Scan for SRT output (same file as analyze - analysis.md) analysis_path = content_dir / "analysis.md" if analysis_path.exists(): results["srt"] = { "success": True, "error_message": None, "output_file": str(analysis_path), } logger.info(f"Found existing SRT output: {analysis_path}")