starmen1265
bb321d9cd4
- Add enable_cache parameter to SRTProcessor.__init__() - Create SemanticCache instance when semantic segmentation is enabled - Pass cache to UnifiedSegmentationStrategy through SemanticSegmentationEngine - Add logging for cache status (enabled/disabled) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
430 lines
15 KiB
Python
430 lines
15 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, Union
|
||
|
||
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,
|
||
enable_semantic_segmentation: bool = False,
|
||
enable_cache: bool = True
|
||
):
|
||
"""初始化 SRT 处理器
|
||
|
||
Args:
|
||
state_manager: 状态管理器
|
||
enable_correction: 是否启用 ASR 纠错（默认 True）
|
||
enable_semantic_segmentation: 是否启用 LLM 语义分段（默认 False）
|
||
enable_cache: 是否启用缓存（默认 True，仅在启用语义分段时有效）
|
||
"""
|
||
super().__init__(state_manager)
|
||
self.enable_correction = enable_correction
|
||
self.enable_semantic_segmentation = enable_semantic_segmentation
|
||
self.segmenter = None
|
||
self.corrector = None
|
||
|
||
# 初始化分段器
|
||
if enable_semantic_segmentation:
|
||
try:
|
||
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
|
||
from kb_workflow.srt.semantic_engine import UnifiedSegmentationStrategy
|
||
from kb_workflow.deep_analyzer import SimpleAIClient
|
||
|
||
# 创建 AI 客户端
|
||
ai_client = SimpleAIClient()
|
||
if ai_client.api_key:
|
||
# 初始化缓存
|
||
cache = None
|
||
if enable_cache:
|
||
from kb_workflow.srt.cache import SemanticCache
|
||
cache = SemanticCache()
|
||
|
||
# 创建统一策略，传入 AI 客户端和缓存
|
||
strategy = UnifiedSegmentationStrategy(
|
||
ai_client=ai_client,
|
||
cache=cache
|
||
)
|
||
self.segmenter = SemanticSegmentationEngine(strategy=strategy)
|
||
logger.info(f"Semantic segmentation enabled (cache={'enabled' if cache else 'disabled'})")
|
||
else:
|
||
logger.info("AI API key not configured, using rule-based segmenter")
|
||
self.segmenter = SemanticSegmenter()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to initialize semantic engine: {e}, using rule-based segmenter")
|
||
self.segmenter = SemanticSegmenter()
|
||
else:
|
||
self.segmenter = SemanticSegmenter()
|
||
|
||
# 延迟初始化纠错器
|
||
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: 加载元数据（需要在分段之前加载，因为 LLM 分段需要元数据）
|
||
metadata = self._load_metadata(input_data.metadata_file)
|
||
|
||
# Step 3: 语义分段
|
||
logger.info("Segmenting subtitles into paragraphs")
|
||
if self.enable_semantic_segmentation:
|
||
# 使用 LLM 语义分段
|
||
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
|
||
if isinstance(self.segmenter, SemanticSegmentationEngine):
|
||
# 返回 VideoSegment 列表（包含标题）
|
||
paragraphs = self.segmenter.segment(segments, metadata=metadata)
|
||
else:
|
||
# 规则分段返回纯文本列表
|
||
paragraphs = self.segmenter.segment(segments)
|
||
else:
|
||
# 规则分段返回纯文本列表
|
||
paragraphs = self.segmenter.segment(segments)
|
||
logger.info(f"Generated {len(paragraphs)} paragraphs")
|
||
|
||
# Step 4: ASR 纠错（可选）
|
||
corrected = False
|
||
if self.corrector and paragraphs:
|
||
logger.info("Applying ASR correction")
|
||
corrected_paragraphs = []
|
||
for para in paragraphs:
|
||
# 处理 VideoSegment 或纯文本
|
||
if hasattr(para, 'content'):
|
||
# VideoSegment 对象
|
||
content = para.content
|
||
corrected_content = self.corrector.correct(content)
|
||
# 创建新的 VideoSegment 对象（假设不可变）
|
||
from kb_workflow.srt.semantic_engine import VideoSegment
|
||
corrected_para = VideoSegment(
|
||
title=para.title,
|
||
content=corrected_content,
|
||
start_index=para.start_index,
|
||
end_index=para.end_index
|
||
)
|
||
else:
|
||
# 纯文本
|
||
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 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
|
||
)
|
||
|
||
# 转换 VideoSegment 为纯文本（用于 SRTOutput）
|
||
if paragraphs and hasattr(paragraphs[0], 'content'):
|
||
paragraphs_text = [para.content if hasattr(para, 'content') else para for para in paragraphs]
|
||
else:
|
||
paragraphs_text = paragraphs
|
||
|
||
# 记录完成状态
|
||
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_text,
|
||
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:
|
||
raw_metadata = json.load(f)
|
||
|
||
# 扁平化嵌套结构（与 AnalyzeProcessor._normalize_metadata 一致）
|
||
if not raw_metadata:
|
||
return {}
|
||
|
||
# 如果已经是扁平结构（顶层有 title）
|
||
if 'title' in raw_metadata:
|
||
return raw_metadata
|
||
|
||
# 处理嵌套结构 {metadata: {...}, stats: {...}}
|
||
if 'metadata' in raw_metadata:
|
||
return raw_metadata['metadata'].copy()
|
||
|
||
return raw_metadata
|
||
|
||
def _extract_tags(self, paragraphs: Union[List[str], List[Any]]) -> List[str]:
|
||
"""从段落中提取标签
|
||
|
||
简单实现：不使用 LLM，只返回默认标签。
|
||
未来可以增强为 LLM 提取。
|
||
|
||
Args:
|
||
paragraphs: 段落列表（可以是纯文本或 VideoSegment 对象）
|
||
|
||
Returns:
|
||
标签列表
|
||
"""
|
||
# 简单实现：返回空列表
|
||
# TODO: 可以使用 LLM 提取标签
|
||
return []
|
||
|
||
def _generate_markdown(
|
||
self,
|
||
output_file: Path,
|
||
metadata: Dict[str, Any],
|
||
paragraphs: Union[List[str], List[Any]], # List[str] 或 List[VideoSegment]
|
||
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, ""]
|
||
|
||
# 检查是否是 VideoSegment 类型（LLM 分段结果）
|
||
from kb_workflow.srt.semantic_engine import VideoSegment
|
||
is_semantic_segmentation = (
|
||
paragraphs and
|
||
hasattr(paragraphs[0], 'title') and
|
||
hasattr(paragraphs[0], 'content')
|
||
)
|
||
|
||
if is_semantic_segmentation:
|
||
# LLM 语义分段：包含标题和内容
|
||
for seg in paragraphs:
|
||
if seg.title:
|
||
content_parts.append(f"## {seg.title}")
|
||
content_parts.append("")
|
||
content_parts.append(seg.content)
|
||
content_parts.append("")
|
||
else:
|
||
# 规则分段：纯文本
|
||
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 已经是字符串格式（如 "05:42"），直接使用
|
||
# 否则（整数秒数）调用 _format_duration 转换
|
||
if isinstance(duration, str):
|
||
duration_str = duration
|
||
else:
|
||
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}"
|