"""测试语义分段引擎""" 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