starmen1265
759503f48a
更新 SEGMENTATION_PROMPT，增强句边界检测： - 列举所有中文句末标点（。？！……） - 强调「自然停顿处」和「语义停顿」 - 说明「不仅依赖标点符号」 添加测试验证句边界检测效果： - test_chinese_sentence_boundary_punctuations - test_prompt_emphasizes_semantic_pause - test_prompt_not_only_punctuation - test_sentence_boundary_in_segmentation Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
516 lines
18 KiB
Python
516 lines
18 KiB
Python
"""测试语义分段引擎"""
|
||
from typing import Protocol
|
||
|
||
import pytest
|
||
import json
|
||
from unittest.mock import Mock
|
||
from pydantic import ValidationError
|
||
|
||
from kb_workflow.srt.parser import Segment
|
||
from kb_workflow.srt.semantic_engine import (
|
||
extract_json_from_llm_response,
|
||
VideoSegment,
|
||
SegmentationResult,
|
||
SemanticSegmentationEngine
|
||
)
|
||
|
||
|
||
class TestVideoTypeDetection:
|
||
"""测试视频类型识别功能"""
|
||
|
||
def test_video_type_detection_from_title(self):
|
||
"""测试从标题推断视频类型"""
|
||
# 知识类 - 从标题推断
|
||
metadata_edu = {"title": "Python编程教程", "category": "教育"}
|
||
engine = SemanticSegmentationEngine()
|
||
video_type = engine._infer_video_type(metadata_edu)
|
||
assert video_type == "知识类"
|
||
|
||
# 知识类 - 包含"课程"
|
||
metadata_course = {"title": "机器学习课程", "category": ""}
|
||
video_type = engine._infer_video_type(metadata_course)
|
||
assert video_type == "知识类"
|
||
|
||
def test_video_type_detection_from_category(self):
|
||
"""测试从分类推断视频类型"""
|
||
# 故事类 - 从分类推断
|
||
metadata_movie = {"title": "星际穿越", "category": "影视"}
|
||
engine = SemanticSegmentationEngine()
|
||
video_type = engine._infer_video_type(metadata_movie)
|
||
assert video_type == "故事类"
|
||
|
||
# 访谈类 - 从标签推断
|
||
metadata_interview = {"title": "名人访谈", "tags": ["访谈", "对话"]}
|
||
video_type = engine._infer_video_type(metadata_interview)
|
||
assert video_type == "访谈类"
|
||
|
||
def test_video_type_fallback_to_default(self):
|
||
"""测试元数据缺失时使用默认类型"""
|
||
# 空元数据
|
||
engine = SemanticSegmentationEngine()
|
||
video_type = engine._infer_video_type({})
|
||
assert video_type == "知识类" # 默认类型
|
||
|
||
# 无效元数据
|
||
video_type = engine._infer_video_type({"title": "", "category": ""})
|
||
assert video_type == "知识类"
|
||
|
||
|
||
class TestVideoTypeInPrompt:
|
||
"""测试视频类型在 Prompt 中的使用"""
|
||
|
||
def test_video_type_in_prompt(self):
|
||
"""测试 Prompt 包含视频类型信息"""
|
||
engine = SemanticSegmentationEngine()
|
||
metadata = {"title": "机器学习教程", "category": "教育"}
|
||
|
||
prompt = engine._build_prompt("测试内容", metadata)
|
||
|
||
assert "视频类型" in prompt or "知识类" in prompt
|
||
assert "机器学习教程" in prompt
|
||
|
||
def test_prompt_format(self):
|
||
"""测试 Prompt 格式正确"""
|
||
engine = SemanticSegmentationEngine()
|
||
content = "这是测试内容"
|
||
metadata = {"title": "测试视频"}
|
||
|
||
prompt = engine._build_prompt(content, metadata)
|
||
|
||
assert "请分析以下视频字幕" in prompt or "字幕内容" in prompt
|
||
assert content in prompt
|
||
|
||
|
||
class TestPydanticModels:
|
||
"""测试 Pydantic 模型验证"""
|
||
|
||
def test_video_segment_model(self):
|
||
"""测试 VideoSegment 模型"""
|
||
# 创建完整分段
|
||
segment = VideoSegment(
|
||
title="机器学习概述",
|
||
content="机器学习是人工智能的一个分支",
|
||
start_index=1,
|
||
end_index=5
|
||
)
|
||
assert segment.title == "机器学习概述"
|
||
assert segment.content == "机器学习是人工智能的一个分支"
|
||
assert segment.start_index == 1
|
||
assert segment.end_index == 5
|
||
|
||
# 创建无标题分段（title 可选）
|
||
segment_no_title = VideoSegment(
|
||
content="测试内容",
|
||
start_index=1,
|
||
end_index=3
|
||
)
|
||
assert segment_no_title.title is None
|
||
|
||
def test_segmentation_result_model(self):
|
||
"""测试 SegmentationResult 模型"""
|
||
segments = [
|
||
VideoSegment(title="第一章", content="内容1", start_index=1, end_index=5),
|
||
VideoSegment(title="第二章", content="内容2", start_index=6, end_index=10)
|
||
]
|
||
|
||
result = SegmentationResult(
|
||
video_type="知识类",
|
||
segments=segments,
|
||
actual_segment_count=2
|
||
)
|
||
|
||
assert result.video_type == "知识类"
|
||
assert len(result.segments) == 2
|
||
assert result.actual_segment_count == 2
|
||
|
||
def test_model_validation(self):
|
||
"""测试模型字段验证"""
|
||
# 缺少必填字段应该抛出异常
|
||
with pytest.raises(Exception): # Pydantic ValidationError
|
||
VideoSegment(
|
||
title="测试",
|
||
# 缺少 content
|
||
start_index=1,
|
||
end_index=3
|
||
)
|
||
|
||
def test_title_length_validation_fails(self):
|
||
"""测试标题长度验证（不超过 15 字）"""
|
||
with pytest.raises(ValidationError) as exc_info:
|
||
VideoSegment(
|
||
title="这是一个超过十五个字符限制的非常非常长的标题",
|
||
content="内容...",
|
||
start_index=1,
|
||
end_index=100
|
||
)
|
||
assert "标题长度" in str(exc_info.value) or "15" in str(exc_info.value)
|
||
|
||
def test_title_exactly_15_chars_passes(self):
|
||
"""测试标题正好 15 字时应该通过验证"""
|
||
# 使用明确的15字字符串
|
||
title_15 = "一二三四五六七八九十十一二十三"
|
||
segment = VideoSegment(
|
||
title=title_15,
|
||
content="内容...",
|
||
start_index=1,
|
||
end_index=100
|
||
)
|
||
assert len(segment.title) == 15
|
||
assert segment.title == title_15
|
||
|
||
|
||
class TestChineseSentenceBoundary:
|
||
"""测试中文句边界检测功能（Plan 13-05）"""
|
||
|
||
def test_chinese_sentence_boundary_punctuations(self):
|
||
"""Test 1: Prompt 列举所有中文句末标点（。？！……）"""
|
||
engine = SemanticSegmentationEngine()
|
||
|
||
# 构建一个测试 prompt
|
||
content = "测试内容"
|
||
metadata = {"title": "测试视频", "video_type": "知识类"}
|
||
prompt = engine._build_prompt(content, metadata)
|
||
|
||
# 验证 Prompt 包含所有中文句末标点
|
||
assert "。" in prompt or "句号" in prompt
|
||
assert "？" in prompt or "问号" in prompt
|
||
assert "！" in prompt or "感叹号" in prompt
|
||
assert "……" in prompt or "省略号" in prompt
|
||
|
||
def test_prompt_emphasizes_semantic_pause(self):
|
||
"""Test 2: Prompt 强调「自然停顿处」和「语义停顿」"""
|
||
engine = SemanticSegmentationEngine()
|
||
|
||
content = "测试内容"
|
||
metadata = {"title": "测试视频", "video_type": "知识类"}
|
||
prompt = engine._build_prompt(content, metadata)
|
||
|
||
# 验证 Prompt 强调语义停顿
|
||
assert "自然停顿" in prompt or "语义停顿" in prompt
|
||
|
||
def test_prompt_not_only_punctuation(self):
|
||
"""Test 3: Prompt 说明「不仅依赖标点符号」"""
|
||
engine = SemanticSegmentationEngine()
|
||
|
||
content = "测试内容"
|
||
metadata = {"title": "测试视频", "video_type": "知识类"}
|
||
prompt = engine._build_prompt(content, metadata)
|
||
|
||
# 验证 Prompt 说明不仅依赖标点
|
||
assert "不仅" in prompt or "不仅依赖" in prompt or "不仅限于" in prompt
|
||
|
||
def test_sentence_boundary_in_segmentation(self):
|
||
"""Test 4: 分段边界在标点符号后"""
|
||
# 创建 mock AI 客户端
|
||
mock_client = Mock()
|
||
mock_client.api_key = "test_key"
|
||
|
||
# 模拟包含正确句边界的分段响应
|
||
mock_response = json.dumps({
|
||
"video_type": "知识类",
|
||
"segments": [
|
||
{
|
||
"title": "第一段",
|
||
"content": "这是第一句话。这是第二句话？",
|
||
"start_index": 0,
|
||
"end_index": 20
|
||
},
|
||
{
|
||
"title": "第二段",
|
||
"content": "这是第三句话！这是第四句话……",
|
||
"start_index": 20,
|
||
"end_index": 40
|
||
}
|
||
],
|
||
"actual_segment_count": 2
|
||
}, ensure_ascii=False)
|
||
|
||
mock_client.generate.return_value = mock_response
|
||
|
||
# 创建测试数据
|
||
from kb_workflow.srt.parser import Segment
|
||
segments = [
|
||
Segment(1, 0.0, 5.0, "这是第一句话。"),
|
||
Segment(2, 5.0, 10.0, "这是第二句话？"),
|
||
Segment(3, 10.0, 15.0, "这是第三句话！"),
|
||
Segment(4, 15.0, 20.0, "这是第四句话……")
|
||
]
|
||
|
||
engine = SemanticSegmentationEngine(ai_client=mock_client)
|
||
result = engine.segment(segments, {"title": "测试", "video_type": "知识类"})
|
||
|
||
# 验证分段边界在标点后
|
||
assert len(result) == 2
|
||
# 第一个分段在问号后结束
|
||
assert result[0].strip().endswith("？")
|
||
# 第二个分段在省略号后结束
|
||
assert result[1].strip().endswith("……")
|
||
|
||
|
||
# ============================================================================
|
||
# 计划 13-02: 语义边界检测测试
|
||
# ============================================================================
|
||
|
||
class AIClientProtocol(Protocol):
|
||
"""AI 客户端协议"""
|
||
|
||
def generate(self, prompt: str, max_tokens: int = 16000, disable_reasoning: bool = False) -> str:
|
||
"""生成文本"""
|
||
...
|
||
|
||
|
||
class MockAIClient:
|
||
"""模拟 AI 客户端"""
|
||
|
||
def __init__(self, response: str):
|
||
self.response = response
|
||
self.api_key = "mock_key"
|
||
|
||
def generate(self, prompt: str, max_tokens: int = 16000, disable_reasoning: bool = False) -> str:
|
||
"""返回预设的响应"""
|
||
return self.response
|
||
|
||
|
||
class TestTopicBoundaryDetection:
|
||
"""测试话题转换点检测"""
|
||
|
||
def test_topic_boundary_detection(self):
|
||
"""测试 LLM 正确识别话题转换点
|
||
|
||
场景：一段包含两个不同主题的字幕
|
||
- 主题1：机器学习介绍（0-500字）
|
||
- 主题2：深度学习介绍（500-1000字）
|
||
|
||
预期：LLM 在话题转换处创建分段边界
|
||
"""
|
||
# 准备测试数据
|
||
segments = self._create_multi_topic_segments()
|
||
|
||
# 模拟 LLM 响应：正确识别两个主题
|
||
llm_response = json.dumps({
|
||
"video_type": "知识类",
|
||
"segments": [
|
||
{
|
||
"title": "机器学习概述",
|
||
"content": "机器学习是人工智能的一个分支。它使计算机能够从数据中学习。",
|
||
"start_index": 0,
|
||
"end_index": 35
|
||
},
|
||
{
|
||
"title": "深度学习基础",
|
||
"content": "深度学习是机器学习的子领域，使用神经网络模拟人脑。",
|
||
"start_index": 36,
|
||
"end_index": 71
|
||
}
|
||
],
|
||
"actual_segment_count": 2
|
||
}, ensure_ascii=False)
|
||
|
||
mock_client = MockAIClient(llm_response)
|
||
engine = SemanticSegmentationEngine(ai_client=mock_client)
|
||
|
||
metadata = {
|
||
"title": "AI技术介绍",
|
||
"video_type": "知识类"
|
||
}
|
||
|
||
paragraphs = engine.segment(segments, metadata)
|
||
|
||
# 验证：结果包含两个段落
|
||
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 深度学习"
|
||
|
||
def test_semantic_completeness(self):
|
||
"""测试单个话题内容完整，不混入其他话题
|
||
|
||
场景：LLM 应该保持每个话题的完整性
|
||
- 不应该在话题中间截断
|
||
- 不应该混入不同话题的内容
|
||
"""
|
||
# 准备测试数据
|
||
segments = self._create_single_topic_segments()
|
||
|
||
# 模拟 LLM 响应：单个完整话题
|
||
llm_response = json.dumps({
|
||
"video_type": "知识类",
|
||
"segments": [
|
||
{
|
||
"title": "梯度下降算法",
|
||
"content": "梯度下降是一种优化算法。它通过迭代调整参数来最小化损失函数。梯度下降有三种变体：批量梯度下降、随机梯度下降和小批量梯度下降。",
|
||
"start_index": 0,
|
||
"end_index": 105
|
||
}
|
||
],
|
||
"actual_segment_count": 1
|
||
}, ensure_ascii=False)
|
||
|
||
mock_client = MockAIClient(llm_response)
|
||
engine = SemanticSegmentationEngine(ai_client=mock_client)
|
||
|
||
metadata = {
|
||
"title": "优化算法介绍",
|
||
"video_type": "知识类"
|
||
}
|
||
|
||
paragraphs = engine.segment(segments, metadata)
|
||
|
||
# 验证：单个段落包含完整内容
|
||
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]
|
||
|
||
def test_no_mixed_topics(self):
|
||
"""测试分段不混入不同话题的内容
|
||
|
||
验证每个分段只包含一个话题的内容
|
||
"""
|
||
segments = self._create_multi_topic_segments()
|
||
|
||
# 模拟 LLM 响应：确保分段不重叠
|
||
llm_response = json.dumps({
|
||
"video_type": "知识类",
|
||
"segments": [
|
||
{
|
||
"title": "机器学习概述",
|
||
"content": "机器学习是人工智能的一个分支。",
|
||
"start_index": 0,
|
||
"end_index": 20
|
||
},
|
||
{
|
||
"title": "深度学习基础",
|
||
"content": "深度学习是机器学习的子领域。",
|
||
"start_index": 21,
|
||
"end_index": 40
|
||
}
|
||
],
|
||
"actual_segment_count": 2
|
||
}, ensure_ascii=False)
|
||
|
||
mock_client = MockAIClient(llm_response)
|
||
engine = SemanticSegmentationEngine(ai_client=mock_client)
|
||
|
||
metadata = {
|
||
"title": "AI技术介绍",
|
||
"video_type": "知识类"
|
||
}
|
||
|
||
paragraphs = engine.segment(segments, metadata)
|
||
|
||
# 验证：分段内容不混入其他话题
|
||
assert len(paragraphs) == 2
|
||
# 第一个段落不应该包含深度学习的内容（除了可能的引用）
|
||
assert "深度学习" not in paragraphs[0] or "子领域" in paragraphs[0]
|
||
# 第二个段落可以提到"机器学习"因为它是"深度学习的超类"
|
||
|
||
def _create_multi_topic_segments(self) -> list[Segment]:
|
||
"""创建包含多个话题的字幕分段"""
|
||
return [
|
||
Segment(1, 0.0, 2.0, "机器学习是人工智能的一个分支。"),
|
||
Segment(2, 2.0, 4.0, "它使计算机能够从数据中学习。"),
|
||
Segment(3, 4.0, 6.0, "深度学习是机器学习的子领域，"),
|
||
Segment(4, 6.0, 8.0, "使用神经网络模拟人脑。"),
|
||
]
|
||
|
||
def _create_single_topic_segments(self) -> list[Segment]:
|
||
"""创建包含单个话题的字幕分段"""
|
||
return [
|
||
Segment(1, 0.0, 2.0, "梯度下降是一种优化算法。"),
|
||
Segment(2, 2.0, 4.0, "它通过迭代调整参数来最小化损失函数。"),
|
||
Segment(3, 4.0, 6.0, "梯度下降有三种变体："),
|
||
Segment(4, 6.0, 8.0, "批量梯度下降、随机梯度下降和小批量梯度下降。"),
|
||
]
|
||
|
||
|
||
class TestJSONExtraction:
|
||
"""测试 JSON 提取功能"""
|
||
|
||
def test_extract_json_from_code_block(self):
|
||
"""测试从 ```json ``` 代码块中提取 JSON"""
|
||
response = '''这是分析结果：
|
||
|
||
```json
|
||
{
|
||
"video_type": "知识类",
|
||
"segments": []
|
||
}
|
||
```
|
||
|
||
希望这个结果有帮助。'''
|
||
|
||
result = extract_json_from_llm_response(response)
|
||
assert result["video_type"] == "知识类"
|
||
assert "segments" in result
|
||
|
||
def test_extract_json_from_plain_text(self):
|
||
"""测试从纯 JSON 文本中提取"""
|
||
response = '{"video_type": "知识类", "segments": []}'
|
||
|
||
result = extract_json_from_llm_response(response)
|
||
assert result["video_type"] == "知识类"
|
||
|
||
def test_extract_json_from_mixed_text(self):
|
||
"""测试从混合文本中提取 JSON"""
|
||
response = '分析如下：{"video_type": "知识类", "segments": []} 以上是结果。'
|
||
|
||
result = extract_json_from_llm_response(response)
|
||
assert result["video_type"] == "知识类"
|
||
|
||
def test_extract_json_invalid_response(self):
|
||
"""测试无效响应抛出异常"""
|
||
with pytest.raises(ValueError, match="未找到 JSON 内容|无法从 LLM 响应中提取有效 JSON"):
|
||
extract_json_from_llm_response("这是一段普通文本，没有 JSON")
|
||
|
||
|
||
class TestPydanticValidationExtended:
|
||
"""测试 Pydantic 验证功能（扩展）"""
|
||
|
||
def test_validate_segment_overlap(self):
|
||
"""测试分段不重叠"""
|
||
# 不重叠的分段
|
||
data = {
|
||
"video_type": "知识类",
|
||
"segments": [
|
||
{
|
||
"title": "第一段",
|
||
"content": "第一段内容",
|
||
"start_index": 0,
|
||
"end_index": 10
|
||
},
|
||
{
|
||
"title": "第二段",
|
||
"content": "第二段内容",
|
||
"start_index": 10,
|
||
"end_index": 20
|
||
}
|
||
],
|
||
"actual_segment_count": 2
|
||
}
|
||
|
||
result = SegmentationResult(**data)
|
||
assert len(result.segments) == 2
|
||
# 验证不重叠
|
||
for i in range(len(result.segments) - 1):
|
||
assert result.segments[i].end_index <= result.segments[i + 1].start_index, \
|
||
f"Segment {i} overlaps with segment {i+1}"
|
||
|
||
def test_allow_null_title(self):
|
||
"""测试允许标题为 null"""
|
||
data = {
|
||
"video_type": "知识类",
|
||
"segments": [
|
||
{
|
||
"title": None,
|
||
"content": "这是一段过渡性内容，不需要标题",
|
||
"start_index": 0,
|
||
"end_index": 20
|
||
}
|
||
],
|
||
"actual_segment_count": 1
|
||
}
|
||
|
||
result = SegmentationResult(**data)
|
||
assert result.segments[0].title is None
|