diff --git a/kb_workflow/pipeline/processors/srt.py b/kb_workflow/pipeline/processors/srt.py new file mode 100644 index 0000000..2723e38 --- /dev/null +++ b/kb_workflow/pipeline/processors/srt.py @@ -0,0 +1,333 @@ +"""SRTProcessor - B站字幕处理处理器 + +将 B站 SRT 字幕文件转换为简化的 Markdown 文档。 + +处理流程: +1. 解析 SRT 文件(调用 parser.parse_srt_file) +2. 语义分段(调用 segmenter.segment) +3. ASR 纠错(调用 corrector.correct,可选) +4. 生成 Markdown 文件 + +简化版:不做深度观点分析,只保留字幕内容和元数据。 +""" + +from __future__ import annotations + +import json +import logging +from datetime import datetime +from pathlib import Path +from typing import Dict, Any, List, Optional + +from kb_workflow.pipeline.config import get_settings +from kb_workflow.pipeline.processors.base import BaseProcessor +from kb_workflow.pipeline.schemas import SRTInput, SRTOutput +from kb_workflow.pipeline.state import WorkflowStateManager +from kb_workflow.srt.parser import parse_srt_file +from kb_workflow.srt.segmenter import SemanticSegmenter + +logger = logging.getLogger(__name__) + + +class SRTProcessor(BaseProcessor): + """SRT 字幕处理器 + + 将 SRT 字幕文件转换为简化的 Markdown 文档。 + 继承 BaseProcessor,遵循 Pipeline 处理器模式。 + + Attributes: + STEP_NAME: 步骤名称 "srt" + settings: KB 配置 + state_manager: 状态管理器 + segmenter: 语义分段器 + corrector: ASR 纠错器(可选) + """ + + STEP_NAME = "srt" + + def __init__( + self, + state_manager: Optional[WorkflowStateManager] = None, + enable_correction: bool = True + ): + """初始化 SRT 处理器 + + Args: + state_manager: 状态管理器 + enable_correction: 是否启用 ASR 纠错(默认 True) + """ + super().__init__(state_manager) + self.enable_correction = enable_correction + self.segmenter = SemanticSegmenter() + self.corrector = None + + # 延迟初始化纠错器 + if enable_correction: + try: + from kb_workflow.srt.corrector import ASRCorrector + from kb_workflow.deep_analyzer import SimpleAIClient + + # 尝试创建 AI 客户端 + ai_client = SimpleAIClient() + if ai_client.api_key: + self.corrector = ASRCorrector(ai_client=ai_client) + logger.info("ASR corrector initialized") + else: + logger.info("AI API key not configured, ASR correction disabled") + except Exception as e: + logger.warning(f"Failed to initialize ASR corrector: {e}") + + async def process(self, input_data: SRTInput) -> SRTOutput: + """处理 SRT 文件 + + Args: + input_data: SRT 输入数据 + + Returns: + SRT 处理结果 + + Raises: + Exception: 处理失败时抛出异常 + """ + try: + # 记录开始状态 + if self.state_manager: + state = self.state_manager.load() + if state: + state.mark_step_started(self.STEP_NAME) + self.state_manager.save(state) + + # Step 1: 解析 SRT 文件 + logger.info(f"Parsing SRT file: {input_data.srt_file}") + segments = parse_srt_file(Path(input_data.srt_file)) + logger.info(f"Parsed {len(segments)} segments") + + # Step 2: 语义分段 + logger.info("Segmenting subtitles into paragraphs") + paragraphs = self.segmenter.segment(segments) + logger.info(f"Generated {len(paragraphs)} paragraphs") + + # Step 3: ASR 纠错(可选) + corrected = False + if self.corrector and paragraphs: + logger.info("Applying ASR correction") + corrected_paragraphs = [] + for para in paragraphs: + corrected_para = self.corrector.correct(para) + corrected_paragraphs.append(corrected_para) + if corrected_para != para: + corrected = True + paragraphs = corrected_paragraphs + logger.info(f"ASR correction applied: {corrected}") + + # Step 4: 加载元数据 + metadata = self._load_metadata(input_data.metadata_file) + + # Step 5: 生成 Markdown + logger.info(f"Generating Markdown: {input_data.output_file}") + tags = self._extract_tags(paragraphs) + self._generate_markdown( + output_file=Path(input_data.output_file), + metadata=metadata, + paragraphs=paragraphs, + tags=tags + ) + + # 记录完成状态 + if self.state_manager: + state = self.state_manager.load() + if state: + summary = f"paragraphs={len(paragraphs)}, corrected={corrected}" + state.mark_step_completed(self.STEP_NAME, summary) + self.state_manager.save(state) + + return SRTOutput( + success=True, + output_file=str(input_data.output_file), + paragraphs=paragraphs, + tags=tags, + corrected=corrected + ) + + except Exception as e: + logger.error(f"SRT processing failed: {e}") + # 记录失败状态 + if self.state_manager: + state = self.state_manager.load() + if state: + state.mark_step_failed(self.STEP_NAME, str(e)) + self.state_manager.save(state) + raise + + def _load_metadata(self, metadata_file: str) -> Dict[str, Any]: + """加载元数据文件 + + Args: + metadata_file: 元数据文件路径 + + Returns: + 元数据字典 + """ + path = Path(metadata_file) + if not path.exists(): + logger.warning(f"Metadata file not found: {metadata_file}") + return {} + + with open(path, 'r', encoding='utf-8') as f: + return json.load(f) + + def _extract_tags(self, paragraphs: List[str]) -> List[str]: + """从段落中提取标签 + + 简单实现:不使用 LLM,只返回默认标签。 + 未来可以增强为 LLM 提取。 + + Args: + paragraphs: 段落列表 + + Returns: + 标签列表 + """ + # 简单实现:返回空列表 + # TODO: 可以使用 LLM 提取标签 + return [] + + def _generate_markdown( + self, + output_file: Path, + metadata: Dict[str, Any], + paragraphs: List[str], + tags: List[str] + ) -> None: + """生成 Markdown 文件 + + Args: + output_file: 输出文件路径 + metadata: 元数据 + paragraphs: 段落列表 + tags: 标签列表 + """ + # 确保输出目录存在 + output_file.parent.mkdir(parents=True, exist_ok=True) + + # 生成 frontmatter + frontmatter = self._generate_frontmatter(metadata, tags) + + # 构建内容 + content_parts = [frontmatter, ""] + + # 添加摘要(可选) + if paragraphs: + content_parts.append("## 摘要") + content_parts.append("") + # 使用第一段作为简单摘要 + first_para = paragraphs[0][:200] if len(paragraphs[0]) > 200 else paragraphs[0] + content_parts.append(first_para + ("..." if len(paragraphs[0]) > 200 else "")) + content_parts.append("") + + # 添加内容 + content_parts.append("## 内容") + content_parts.append("") + for para in paragraphs: + content_parts.append(para) + content_parts.append("") + + # 写入文件 + output_file.write_text('\n'.join(content_parts), encoding='utf-8') + logger.info(f"Markdown file generated: {output_file}") + + def _generate_frontmatter(self, metadata: Dict[str, Any], tags: List[str]) -> str: + """生成 YAML frontmatter + + Args: + metadata: 元数据 + tags: 标签列表 + + Returns: + YAML frontmatter 字符串 + """ + lines = ["---"] + + # 标题 + title = metadata.get('title', '未命名') + lines.append(f'title: "{self._escape_yaml(title)}"') + + # 日期 + pubtime = metadata.get('pubtime', '') or metadata.get('published_at', '') + if pubtime: + date_str = str(pubtime)[:10] + else: + date_str = datetime.now().strftime('%Y-%m-%d') + lines.append(f'date: {date_str}') + + # 作者 + author = metadata.get('author', '未知') + lines.append(f'author: {self._escape_yaml(author)}') + + # 平台 + platform = metadata.get('platform', 'bilibili') + lines.append(f'platform: {platform}') + + # 时长(视频) + duration = metadata.get('duration', 0) + if duration: + duration_str = self._format_duration(duration) + lines.append(f'duration: {duration_str}') + + # BVID + bvid = metadata.get('bvid', '') + if bvid: + lines.append(f'bvid: {bvid}') + + # URL + url = metadata.get('url', '') + if url: + lines.append(f'url: {url}') + + # 标签 + if tags: + tags_str = ', '.join(f'"{t}"' for t in tags) + lines.append(f'tags: [{tags_str}]') + + lines.append("---") + return '\n'.join(lines) + + def _escape_yaml(self, text: str) -> str: + """转义 YAML 特殊字符 + + Args: + text: 原始文本 + + Returns: + 转义后的文本 + """ + # 移除零宽度字符 + invisible_chars = ['\u200b', '\u200c', '\u200d', '\ufeff', '\u00ad'] + for char in invisible_chars: + text = text.replace(char, '') + + # 移除换行符 + text = text.replace('\n', ' ').strip() + + # 如果包含特殊字符,用引号包裹 + if any(c in text for c in [':', '"', "'", '\n', '#']): + text = text.replace('"', '\\"') + return text + + return text + + def _format_duration(self, seconds: int) -> str: + """格式化时长 + + Args: + seconds: 秒数 + + Returns: + 格式化的时长字符串 (HH:MM:SS 或 MM:SS) + """ + h = seconds // 3600 + m = (seconds % 3600) // 60 + s = seconds % 60 + if h > 0: + return f"{h}:{m:02d}:{s:02d}" + return f"{m}:{s:02d}" diff --git a/kb_workflow/pipeline/schemas.py b/kb_workflow/pipeline/schemas.py index 83a3e02..1f5922a 100644 --- a/kb_workflow/pipeline/schemas.py +++ b/kb_workflow/pipeline/schemas.py @@ -444,3 +444,42 @@ class ArchiveOutput(StepOutput): archive_file: str = Field(default="", description="Path to archived file") skipped: bool = Field(default=False, description="Whether archiving was skipped") platform: str = Field(default="", description="Platform name") + + +# ============================================================================= +# SRT Schemas (Phase 12) +# ============================================================================= + + +class SRTInput(StepInput): + """SRT 步骤输入 + + 用于 SRTProcessor 处理 B站字幕。 + + Attributes: + srt_file: SRT 字幕文件路径 + metadata_file: 元数据 JSON 文件路径 + output_file: 输出 Markdown 文件路径 + content_type: 内容类型(默认 video) + """ + srt_file: str = Field(..., min_length=1, description="Path to SRT subtitle file") + metadata_file: str = Field(..., min_length=1, description="Path to metadata JSON file") + output_file: str = Field(..., min_length=1, description="Path to output Markdown file") + content_type: str = Field(default="video", description="Content type") + + +class SRTOutput(StepOutput): + """SRT 步骤输出 + + 继承 StepOutput 的基础字段,添加 SRT 处理特定字段。 + + Attributes: + output_file: 生成的 Markdown 文件路径 + paragraphs: 生成的段落列表 + tags: 提取的标签列表 + corrected: 是否进行了 ASR 纠错 + """ + output_file: str = Field(..., description="Path to generated Markdown file") + paragraphs: List[str] = Field(default_factory=list, description="Generated paragraphs") + tags: List[str] = Field(default_factory=list, description="Extracted tags") + corrected: bool = Field(default=False, description="Whether ASR correction was applied") diff --git a/tests/test_srt_processor.py b/tests/test_srt_processor.py index cbdcadd..c731e0b 100644 --- a/tests/test_srt_processor.py +++ b/tests/test_srt_processor.py @@ -7,6 +7,9 @@ Tests for: - SRTProcessor (pipeline processor) """ +import asyncio +import json + import pytest from pathlib import Path from dataclasses import dataclass @@ -340,3 +343,150 @@ class TestASRCorrector: result = corrector.correct("") assert result == "" + + +# ============================================================================= +# Test 4: SRTProcessor +# ============================================================================= + + +@pytest.fixture +def srt_test_files(tmp_path: Path): + """创建测试用的 SRT 和元数据文件""" + # 创建 SRT 文件 + srt_content = """1 +00:00:00,000 --> 00:00:05,000 +第一句字幕内容 + +2 +00:00:05,000 --> 00:00:10,000 +第二句字幕内容 + +3 +00:00:10,000 --> 00:00:15,000 +第三句字幕内容 +""" + srt_file = tmp_path / "test.srt" + srt_file.write_text(srt_content, encoding="utf-8") + + # 创建元数据文件 + metadata = { + "title": "测试视频标题", + "author": "测试作者", + "platform": "bilibili", + "bvid": "BV1test123", + "duration": 60, + "url": "https://www.bilibili.com/video/BV1test123", + "pubtime": "2026-03-19" + } + metadata_file = tmp_path / "metadata.json" + metadata_file.write_text(json.dumps(metadata, ensure_ascii=False), encoding="utf-8") + + # 输出文件路径 + output_file = tmp_path / "output.md" + + return { + "srt_file": str(srt_file), + "metadata_file": str(metadata_file), + "output_file": str(output_file) + } + + +class TestSRTProcessor: + """测试 SRTProcessor 处理器""" + + def test_srt_processor_generates_markdown(self, srt_test_files): + """SRTProcessor 能处理 SRT 文件生成 Markdown""" + from kb_workflow.pipeline.processors.srt import SRTProcessor + from kb_workflow.pipeline.schemas import SRTInput + + processor = SRTProcessor(enable_correction=False) + + input_data = SRTInput( + content_id="BV1test123", + platform="bilibili", + srt_file=srt_test_files["srt_file"], + metadata_file=srt_test_files["metadata_file"], + output_file=srt_test_files["output_file"] + ) + + # 运行处理 + result = asyncio.run(processor.process(input_data)) + + assert result.success + assert Path(result.output_file).exists() + + # 验证 Markdown 内容 + content = Path(result.output_file).read_text(encoding="utf-8") + assert "---" in content # frontmatter + assert "title:" in content + assert "测试视频标题" in content + assert "## 内容" in content + assert "第一句字幕内容" in content + + def test_srt_processor_with_yaml_frontmatter(self, srt_test_files): + """输出包含 YAML frontmatter(元数据和标签)""" + from kb_workflow.pipeline.processors.srt import SRTProcessor + from kb_workflow.pipeline.schemas import SRTInput + + processor = SRTProcessor(enable_correction=False) + + input_data = SRTInput( + content_id="BV1test123", + platform="bilibili", + srt_file=srt_test_files["srt_file"], + metadata_file=srt_test_files["metadata_file"], + output_file=srt_test_files["output_file"] + ) + + result = asyncio.run(processor.process(input_data)) + + content = Path(result.output_file).read_text(encoding="utf-8") + + # 验证 frontmatter 字段 + assert "title:" in content + assert "date:" in content + assert "author:" in content + assert "platform:" in content + assert "bvid:" in content + assert "url:" in content + + def test_srt_processor_follows_base_processor_pattern(self): + """处理器遵循 BaseProcessor 模式""" + from kb_workflow.pipeline.processors.srt import SRTProcessor + from kb_workflow.pipeline.processors.base import BaseProcessor + + # SRTProcessor 应该继承 BaseProcessor + assert issubclass(SRTProcessor, BaseProcessor) + + # 应该有 STEP_NAME + assert hasattr(SRTProcessor, 'STEP_NAME') + assert SRTProcessor.STEP_NAME == "srt" + + def test_srt_processor_handles_missing_metadata(self, tmp_path: Path): + """元数据文件缺失时能正常处理""" + from kb_workflow.pipeline.processors.srt import SRTProcessor + from kb_workflow.pipeline.schemas import SRTInput + + # 创建 SRT 文件 + srt_content = "1\n00:00:00,000 --> 00:00:05,000\n测试内容\n" + srt_file = tmp_path / "test.srt" + srt_file.write_text(srt_content, encoding="utf-8") + + # 不存在的元数据文件 + output_file = tmp_path / "output.md" + + processor = SRTProcessor(enable_correction=False) + + input_data = SRTInput( + content_id="test123", + platform="bilibili", + srt_file=str(srt_file), + metadata_file=str(tmp_path / "nonexistent.json"), + output_file=str(output_file) + ) + + result = asyncio.run(processor.process(input_data)) + + assert result.success + assert Path(result.output_file).exists()