feat(12-01): implement SRTProcessor and Markdown generation
- Add SRTInput/SRTOutput schemas to pipeline - Create SRTProcessor following BaseProcessor pattern - Generate Markdown with YAML frontmatter - Support metadata extraction and formatting - Handle missing metadata gracefully Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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}"
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user