starmen1265/kb-skills
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(13-07): integrate SemanticSegmentationEngine into SRTProcessor

Browse Source 
- 添加 enable_semantic_segmentation 参数到 SRTProcessor.__init__
- 启用时使用 SemanticSegmentationEngine，禁用时使用 SemanticSegmenter
- AI 不可用时自动降级到规则分段
- 更新 _generate_markdown() 支持带标题的 VideoSegment 输出
- 调整步骤顺序：先加载元数据，再进行语义分段
- 更新 UnifiedSegmentationStrategy.segment() 返回 List[VideoSegment]
- 添加代理方法到 SemanticSegmentationEngine 以保持向后兼容
- 创建全面的集成测试（9 个测试用例）
- 更新现有语义引擎测试以适配新 API

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
starmen1265
2026-03-22 00:32:14 +08:00
parent a6c0293c01
commit ef4a660faa
4 changed files with 578 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kb_workflow/pipeline/processors/srt.py
+94 -19
View File
@@ -17,7 +17,7 @@ import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, List, Optional
from typing import Dict, Any, List, Optional, Union
from kb_workflow.pipeline.config import get_settings
from kb_workflow.pipeline.processors.base import BaseProcessor
@@ -48,19 +48,45 @@ class SRTProcessor(BaseProcessor):
def __init__(
self,
state_manager: Optional[WorkflowStateManager] = None,
enable_correction: bool = True
enable_correction: bool = True,
enable_semantic_segmentation: bool = False
):
"""初始化 SRT 处理器
Args:
state_manager: 状态管理器
enable_correction: 是否启用 ASR 纠错（默认 True）
enable_semantic_segmentation: 是否启用 LLM 语义分段（默认 False）
"""
super().__init__(state_manager)
self.enable_correction = enable_correction
self.segmenter = SemanticSegmenter()
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:
# 创建统一策略，传入 AI 客户端
strategy = UnifiedSegmentationStrategy(ai_client=ai_client)
self.segmenter = SemanticSegmentationEngine(strategy=strategy)
logger.info("Semantic segmentation enabled")
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:
@@ -102,27 +128,53 @@ class SRTProcessor(BaseProcessor):
segments = parse_srt_file(Path(input_data.srt_file))
logger.info(f"Parsed {len(segments)} segments")
# Step 2: 语义分段
# Step 2: 加载元数据（需要在分段之前加载，因为 LLM 分段需要元数据）
metadata = self._load_metadata(input_data.metadata_file)
# Step 3: 语义分段
logger.info("Segmenting subtitles into paragraphs")
paragraphs = self.segmenter.segment(segments)
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 3: ASR 纠错（可选）
# Step 4: 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)
# 处理 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 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)
@@ -133,6 +185,12 @@ class SRTProcessor(BaseProcessor):
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()
@@ -144,7 +202,7 @@ class SRTProcessor(BaseProcessor):
return SRTOutput(
success=True,
output_file=str(input_data.output_file),
paragraphs=paragraphs,
paragraphs=paragraphs_text,
tags=tags,
corrected=corrected
)
@@ -190,14 +248,14 @@ class SRTProcessor(BaseProcessor):
return raw_metadata
def _extract_tags(self, paragraphs: List[str]) -> List[str]:
def _extract_tags(self, paragraphs: Union[List[str], List[Any]]) -> List[str]:
"""从段落中提取标签
简单实现：不使用 LLM，只返回默认标签。
未来可以增强为 LLM 提取。
Args:
paragraphs: 段落列表
paragraphs: 段落列表（可以是纯文本或 VideoSegment 对象）
Returns:
标签列表
@@ -210,7 +268,7 @@ class SRTProcessor(BaseProcessor):
self,
output_file: Path,
metadata: Dict[str, Any],
paragraphs: List[str],
paragraphs: Union[List[str], List[Any]], # List[str] 或 List[VideoSegment]
tags: List[str]
) -> None:
"""生成 Markdown 文件
@@ -230,12 +288,29 @@ class SRTProcessor(BaseProcessor):
# 构建内容
content_parts = [frontmatter, ""]
# 直接添加内容，不生成摘要部分
content_parts.append("## 内容")
content_parts.append("")
for para in paragraphs:
content_parts.append(para)
# 检查是否是 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')
kb_workflow/srt/semantic_engine.py
+51 -13
View File
@@ -56,7 +56,7 @@ class SegmentationStrategy(Protocol):
self,
segments: List,
metadata: dict
) -> List[str]:
) -> List[Union[str, 'VideoSegment']]:
"""对字幕分段进行语义分段
Args:
@@ -64,7 +64,7 @@ class SegmentationStrategy(Protocol):
metadata: 视频元数据（title, duration 等）
Returns:
段落文本列表
分段列表（可以是纯文本列表或 VideoSegment 对象列表）
"""
...
@@ -117,7 +117,7 @@ class UnifiedSegmentationStrategy:
self,
segments: List,
metadata: dict
) -> List[str]:
) -> List['VideoSegment']:
"""对字幕分段进行语义分段
使用统一的 LLM Prompt 进行语义分段，不区分视频类型。
@@ -128,12 +128,17 @@ class UnifiedSegmentationStrategy:
metadata: 视频元数据（title, duration 等）
Returns:
段落文本列表
VideoSegment 对象列表（包含标题和内容）
"""
# AI 不可用，直接使用规则分段
if self.ai_client is None:
logger.info("AI client unavailable, using fallback segmenter")
return self.fallback_segmenter.segment(segments)
# 将规则分段结果转换为 VideoSegment 格式（无标题）
paragraphs = self.fallback_segmenter.segment(segments)
return [
VideoSegment(title=None, content=para, start_index=0, end_index=0)
for para in paragraphs
]
# 合并字幕文本
content = " ".join(seg.text for seg in segments)
@@ -158,15 +163,17 @@ class UnifiedSegmentationStrategy:
fallback_segments=segments
)
# 转换为段落文本
paragraphs = [seg.content for seg in video_segments]
logger.info(f"LLM semantic segmentation: {len(paragraphs)} paragraphs")
return paragraphs
logger.info(f"LLM semantic segmentation: {len(video_segments)} segments")
return video_segments
except Exception as e:
logger.error(f"LLM segmentation failed: {e}, using fallback")
return self.fallback_segmenter.segment(segments)
# 将规则分段结果转换为 VideoSegment 格式（无标题）
paragraphs = self.fallback_segmenter.segment(segments)
return [
VideoSegment(title=None, content=para, start_index=0, end_index=0)
for para in paragraphs
]
def _build_prompt(self, content: str, metadata: dict) -> str:
"""构建语义分段 Prompt
@@ -535,11 +542,42 @@ class SemanticSegmentationEngine:
f"{self.strategy.__class__.__name__}"
)
def _infer_video_type(self, metadata: dict) -> str:
"""从元数据推断视频类型（代理到策略）
Args:
metadata: 视频元数据
Returns:
推断出的视频类型
"""
# 如果策略有 _infer_video_type 方法，使用它
if hasattr(self.strategy, '_infer_video_type'):
return self.strategy._infer_video_type(metadata)
# 否则返回默认类型
return DEFAULT_VIDEO_TYPE
def _build_prompt(self, content: str, metadata: dict) -> str:
"""构建 Prompt（代理到策略）
Args:
content: 字幕内容
metadata: 视频元数据
Returns:
构建好的 Prompt
"""
# 如果策略有 _build_prompt 方法，使用它
if hasattr(self.strategy, '_build_prompt'):
return self.strategy._build_prompt(content, metadata)
# 否则返回基本 prompt
return f"Content: {content}\nMetadata: {metadata}"
def segment(
self,
segments: List,
metadata: dict
) -> List[str]:
) -> List['VideoSegment']:
"""对字幕分段进行语义分段
使用当前配置的策略进行分段。
@@ -549,6 +587,6 @@ class SemanticSegmentationEngine:
metadata: 视频元数据（title, duration 等）
Returns:
段落文本列表
VideoSegment 对象列表（包含标题和内容）
"""
return self.strategy.segment(segments, metadata)