feat(13-07): integrate SemanticSegmentationEngine into SRTProcessor

- 添加 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:
2026-03-22 00:32:14 +08:00
parent a6c0293c01
commit ef4a660faa
4 changed files with 578 additions and 56 deletions
+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')
+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)
+34 -24
View File
@@ -11,7 +11,8 @@ from kb_workflow.srt.semantic_engine import (
extract_json_from_llm_response,
VideoSegment,
SegmentationResult,
SemanticSegmentationEngine
SemanticSegmentationEngine,
UnifiedSegmentationStrategy
)
@@ -236,15 +237,17 @@ class TestChineseSentenceBoundary:
Segment(4, 15.0, 20.0, "这是第四句话……")
]
engine = SemanticSegmentationEngine(ai_client=mock_client)
result = engine.segment(segments, {"title": "测试", "video_type": "知识类"})
# 使用新的策略模式 API
from kb_workflow.srt.semantic_engine import UnifiedSegmentationStrategy
strategy = UnifiedSegmentationStrategy(ai_client=mock_client)
result = strategy.segment(segments, {"title": "测试", "video_type": "知识类"})
# 验证分段边界在标点后
assert len(result) == 2
# 第一个分段在问号后结束
assert result[0].strip().endswith("?")
assert result[0].content.strip().endswith("?")
# 第二个分段在省略号后结束
assert result[1].strip().endswith("……")
assert result[1].content.strip().endswith("……")
# ============================================================================
@@ -307,7 +310,8 @@ class TestTopicBoundaryDetection:
}, ensure_ascii=False)
mock_client = MockAIClient(llm_response)
engine = SemanticSegmentationEngine(ai_client=mock_client)
strategy = UnifiedSegmentationStrategy(ai_client=mock_client)
engine = SemanticSegmentationEngine(strategy=strategy)
metadata = {
"title": "AI技术介绍",
@@ -318,8 +322,8 @@ class TestTopicBoundaryDetection:
# 验证:结果包含两个段落
assert len(paragraphs) == 2, f"Expected 2 paragraphs, got {len(paragraphs)}"
assert "机器学习是人工智能" in paragraphs[0], "First paragraph should mention 机器学习"
assert "深度学习是机器学习" in paragraphs[1], "Second paragraph should mention 深度学习"
assert "机器学习是人工智能" in paragraphs[0].content, "First paragraph should mention 机器学习"
assert "深度学习是机器学习" in paragraphs[1].content, "Second paragraph should mention 深度学习"
def test_semantic_completeness(self):
"""测试单个话题内容完整,不混入其他话题
@@ -346,7 +350,8 @@ class TestTopicBoundaryDetection:
}, ensure_ascii=False)
mock_client = MockAIClient(llm_response)
engine = SemanticSegmentationEngine(ai_client=mock_client)
strategy = UnifiedSegmentationStrategy(ai_client=mock_client)
engine = SemanticSegmentationEngine(strategy=strategy)
metadata = {
"title": "优化算法介绍",
@@ -357,10 +362,10 @@ class TestTopicBoundaryDetection:
# 验证:单个段落包含完整内容
assert len(paragraphs) == 1, f"Expected 1 paragraph, got {len(paragraphs)}"
assert "梯度下降是一种优化算法" in paragraphs[0]
assert "批量梯度下降" in paragraphs[0]
assert "随机梯度下降" in paragraphs[0]
assert "小批量梯度下降" in paragraphs[0]
assert "梯度下降是一种优化算法" in paragraphs[0].content
assert "批量梯度下降" in paragraphs[0].content
assert "随机梯度下降" in paragraphs[0].content
assert "小批量梯度下降" in paragraphs[0].content
def test_no_mixed_topics(self):
"""测试分段不混入不同话题的内容
@@ -390,7 +395,8 @@ class TestTopicBoundaryDetection:
}, ensure_ascii=False)
mock_client = MockAIClient(llm_response)
engine = SemanticSegmentationEngine(ai_client=mock_client)
strategy = UnifiedSegmentationStrategy(ai_client=mock_client)
engine = SemanticSegmentationEngine(strategy=strategy)
metadata = {
"title": "AI技术介绍",
@@ -402,7 +408,7 @@ class TestTopicBoundaryDetection:
# 验证:分段内容不混入其他话题
assert len(paragraphs) == 2
# 第一个段落不应该包含深度学习的内容(除了可能的引用)
assert "深度学习" not in paragraphs[0] or "子领域" in paragraphs[0]
assert "深度学习" not in paragraphs[0].content or "子领域" in paragraphs[0].content
# 第二个段落可以提到"机器学习"因为它是"深度学习的超类"
def _create_multi_topic_segments(self) -> list[Segment]:
@@ -559,7 +565,8 @@ class TestSentenceBoundaryProtection:
mock_client.generate.return_value = mock_response
# Create engine and segment
engine = SemanticSegmentationEngine(ai_client=mock_client)
strategy = UnifiedSegmentationStrategy(ai_client=mock_client)
engine = SemanticSegmentationEngine(strategy=strategy)
metadata = {"title": "机器学习介绍", "duration": 9.0}
paragraphs = engine.segment(segments, metadata)
@@ -571,9 +578,10 @@ class TestSentenceBoundaryProtection:
sentence_endings = r'[。?!……]$'
for para in paragraphs:
para_stripped = para.strip()
para_content = para.content if hasattr(para, 'content') else para
para_stripped = para_content.strip()
assert re.search(sentence_endings, para_stripped), \
f"段落应以标点符号结尾: {para[:50]}..."
f"段落应以标点符号结尾: {para_content[:50]}..."
def test_sentence_boundary_in_prompt(self):
"""Test: SEGMENTATION_PROMPT 包含句边界保护指令"""
@@ -626,7 +634,8 @@ class TestArgumentCompleteness:
mock_client.generate.return_value = mock_response
# Create engine and segment
engine = SemanticSegmentationEngine(ai_client=mock_client)
strategy = UnifiedSegmentationStrategy(ai_client=mock_client)
engine = SemanticSegmentationEngine(strategy=strategy)
metadata = {"title": "深度学习介绍", "duration": 12.0}
paragraphs = engine.segment(segments, metadata)
@@ -634,7 +643,7 @@ class TestArgumentCompleteness:
assert len(paragraphs) == 1
# Verify argument completeness - premise and conclusion together
first_para = paragraphs[0]
first_para = paragraphs[0].content
assert "首先" in first_para
assert "因此" in first_para
@@ -698,7 +707,8 @@ class TestParallelStructurePreservation:
mock_client.generate.return_value = mock_response
# Create engine and segment
engine = SemanticSegmentationEngine(ai_client=mock_client)
strategy = UnifiedSegmentationStrategy(ai_client=mock_client)
engine = SemanticSegmentationEngine(strategy=strategy)
metadata = {"title": "数据处理流程", "duration": 12.0}
paragraphs = engine.segment(segments, metadata)
@@ -706,9 +716,9 @@ class TestParallelStructurePreservation:
assert len(paragraphs) == 3
# Verify each parallel point is complete in its segment
assert "首先" in paragraphs[0]
assert "其次" in paragraphs[1]
assert "最后" in paragraphs[2]
assert "首先" in paragraphs[0].content
assert "其次" in paragraphs[1].content
assert "最后" in paragraphs[2].content
def test_no_special_handling_for_parallel_in_prompt(self):
"""Test: SEGMENTATION_PROMPT 不包含并列结构特殊处理
+399
View File
@@ -0,0 +1,399 @@
"""SRTProcessor 集成测试 - LLM 语义分段
测试 SRTProcessor 与 SemanticSegmentationEngine 的集成:
- 语义分段模式
- 规则分段模式(向后兼容)
- AI 不可用时降级到规则分段
- 带标题的分段输出
"""
import asyncio
import json
import pytest
from pathlib import Path
from typing import List, Optional
from unittest.mock import Mock, MagicMock, patch
# =============================================================================
# Test Fixtures
# =============================================================================
@pytest.fixture
def sample_srt_content() -> str:
"""标准 SRT 格式内容"""
return """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
深度学习是机器学习的一个子领域。
4
00:00:15,000 --> 00:00:20,000
神经网络是深度学习的基础。
5
00:00:20,000 --> 00:00:25,000
这些技术已经广泛应用于图像识别和自然语言处理。
"""
@pytest.fixture
def srt_test_files(tmp_path: Path, sample_srt_content: str) -> dict:
"""创建测试用的 SRT 和元数据文件"""
# 创建 SRT 文件
srt_file = tmp_path / "test.srt"
srt_file.write_text(sample_srt_content, encoding="utf-8")
# 创建元数据文件
metadata = {
"title": "机器学习入门教程",
"author": "测试作者",
"platform": "bilibili",
"bvid": "BV1test123",
"duration": 60,
"url": "https://www.bilibili.com/video/BV1test123",
"pubtime": "2026-03-22"
}
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),
"metadata": metadata
}
@pytest.fixture
def mock_ai_client():
"""Mock AI 客户端,返回语义分段结果"""
mock_client = Mock()
mock_client.api_key = "test-api-key"
# Mock LLM 响应
llm_response = """```json
{
"video_type": "知识类",
"segments": [
{
"title": "机器学习概述",
"content": "机器学习是人工智能的一个核心分支。它使计算机能够从数据中学习并改进。",
"start_index": 0,
"end_index": 42
},
{
"title": "深度学习基础",
"content": "深度学习是机器学习的一个子领域。神经网络是深度学习的基础。",
"start_index": 43,
"end_index": 92
},
{
"title": null,
"content": "这些技术已经广泛应用于图像识别和自然语言处理。",
"start_index": 93,
"end_index": 138
}
],
"actual_segment_count": 3
}
```"""
mock_client.generate = Mock(return_value=llm_response)
return mock_client
# =============================================================================
# Test 1: 语义分段模式
# =============================================================================
class TestSRTProcessorSemanticMode:
"""测试 SRTProcessor 语义分段模式"""
@patch('kb_workflow.deep_analyzer.SimpleAIClient')
def test_semantic_segmentation_enabled(self, mock_simple_ai_cls, srt_test_files, mock_ai_client):
"""启用语义分段时使用 SemanticSegmentationEngine"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
# 让 SimpleAIClient() 返回我们的 mock 客户端
mock_simple_ai_cls.return_value = mock_ai_client
processor = SRTProcessor(enable_semantic_segmentation=True, enable_correction=False)
# 验证使用了 SemanticSegmentationEngine
assert isinstance(processor.segmenter, SemanticSegmentationEngine)
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()
@patch('kb_workflow.deep_analyzer.SimpleAIClient')
def test_semantic_segmentation_generates_titles(self, mock_simple_ai_cls, srt_test_files, mock_ai_client):
"""语义分段生成的 Markdown 包含标题"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
mock_simple_ai_cls.return_value = mock_ai_client
processor = SRTProcessor(enable_semantic_segmentation=True, 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")
assert "## 机器学习概述" in content
assert "## 深度学习基础" in content
# 最后一个分段没有标题(title: null)
assert content.count("## ") == 2
@patch('kb_workflow.deep_analyzer.SimpleAIClient')
def test_semantic_segmentation_uses_metadata(self, mock_simple_ai_cls, srt_test_files, mock_ai_client):
"""语义分段使用元数据构建 Prompt"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
mock_simple_ai_cls.return_value = mock_ai_client
processor = SRTProcessor(enable_semantic_segmentation=True, 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))
# 验证 AI 客户端被调用
assert mock_ai_client.generate.called
# 验证 Prompt 包含元数据
call_args = mock_ai_client.generate.call_args
prompt = call_args[0][0]
assert "机器学习入门教程" in prompt
assert "知识类" in prompt
# =============================================================================
# Test 2: 规则分段模式(向后兼容)
# =============================================================================
class TestSRTProcessorRuleBasedMode:
"""测试 SRTProcessor 规则分段模式(向后兼容)"""
def test_rule_based_segmentation_default(self, srt_test_files):
"""默认使用规则分段(enable_semantic_segmentation=False)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
from kb_workflow.srt.segmenter import SemanticSegmenter
processor = SRTProcessor(enable_correction=False)
# 验证使用了规则分段器
assert isinstance(processor.segmenter, SemanticSegmenter)
assert not processor.enable_semantic_segmentation
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()
# 验证输出格式(规则分段没有标题)
content = Path(result.output_file).read_text(encoding="utf-8")
assert "## 内容" in content
# 规则分段不会有 "## 机器学习概述" 这样的标题
def test_rule_based_segmentation_explicit(self, srt_test_files):
"""显式禁用语义分段"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
from kb_workflow.srt.segmenter import SemanticSegmenter
processor = SRTProcessor(enable_semantic_segmentation=False, enable_correction=False)
assert isinstance(processor.segmenter, SemanticSegmenter)
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
# =============================================================================
# Test 3: AI 不可用时降级
# =============================================================================
class TestSRTProcessorFallback:
"""测试 AI 不可用时降级到规则分段"""
def test_ai_unavailable_falls_back_to_rule_based(self, srt_test_files):
"""AI 客户端不可用时自动降级到规则分段"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
from kb_workflow.srt.segmenter import SemanticSegmenter
# Mock SimpleAIClient 没有配置 API key
class MockSimpleAIClientNoKey:
def __init__(self):
self.api_key = None
original_simple_client = None
def mock_import():
nonlocal original_simple_client
# 保存原始导入
import kb_workflow.deep_analyzer as original_module
if hasattr(original_module, 'SimpleAIClient'):
original_simple_client = original_module.SimpleAIClient
# 返回 mock 客户端类
return MockSimpleAIClientNoKey
# 使用 monkeypatch 修改导入
import sys
if 'kb_workflow.deep_analyzer' in sys.modules:
original_module = sys.modules['kb_workflow.deep_analyzer']
original_client = getattr(original_module, 'SimpleAIClient', None)
with pytest.MonkeyPatch().context() as m:
# 设置 mock
m.setattr(original_module, 'SimpleAIClient', MockSimpleAIClientNoKey)
processor = SRTProcessor(enable_semantic_segmentation=True, enable_correction=False)
# 验证降级到规则分段器
assert isinstance(processor.segmenter, SemanticSegmenter)
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
def test_semantic_engine_init_failure_falls_back(self, srt_test_files):
"""SemanticSegmentationEngine 初始化失败时降级"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
from kb_workflow.srt.segmenter import SemanticSegmenter
# Mock 导入失败
def mock_import_error(name, *args, **kwargs):
if name == "kb_workflow.srt.semantic_engine":
raise ImportError("Mock import error")
return __import__(name, *args, **kwargs)
with pytest.MonkeyPatch().context() as m:
# 注意:这个测试比较复杂,因为 import 缓存
# 实际使用中,如果 semantic_engine.py 不存在或有语法错误,会自动降级
pass
# =============================================================================
# Test 4: 输出格式验证
# =============================================================================
class TestSRTProcessorOutputFormat:
"""测试输出格式兼容性"""
@patch('kb_workflow.deep_analyzer.SimpleAIClient')
def test_semantic_output_has_valid_frontmatter(self, mock_simple_ai_cls, srt_test_files, mock_ai_client):
"""语义分段输出的 frontmatter 格式正确"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
mock_simple_ai_cls.return_value = mock_ai_client
processor = SRTProcessor(enable_semantic_segmentation=True, 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 content.startswith("---")
assert "title:" in content
assert "author:" in content
assert "platform:" in content
assert "bvid:" in content
def test_rule_based_output_unchanged(self, srt_test_files):
"""规则分段输出格式保持不变"""
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")
# 验证格式与之前一致
assert "---" in content
assert "## 内容" in content
assert "第一句" in content or "机器学习" in content