334 lines
10 KiB
Python
334 lines
10 KiB
Python
"""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}"
|