diff --git a/kb_workflow/srt/semantic_engine.py b/kb_workflow/srt/semantic_engine.py new file mode 100644 index 0000000..4fc756a --- /dev/null +++ b/kb_workflow/srt/semantic_engine.py @@ -0,0 +1,323 @@ +"""语义分段引擎 + +使用 LLM 实现智能视频字幕语义分段。 + +功能： +- 视频类型识别（知识类/故事类/访谈类/新闻类/vlog 等） +- 语义边界检测 +- 内容完整性保护 +- 主题标题生成 +""" + +from __future__ import annotations + +import json +import logging +import re +from typing import List, Optional, Protocol + +from pydantic import BaseModel, Field, field_validator, ValidationError + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# AI Client Protocol +# ============================================================================ + +class AIClientProtocol(Protocol): + """AI 客户端协议（与 ASRCorrector 保持一致）""" + + def generate(self, prompt: str, max_tokens: int = 16000, + disable_reasoning: bool = False) -> str: + """生成文本""" + ... + + +# ============================================================================ +# Pydantic Models +# ============================================================================ + +class VideoSegment(BaseModel): + """视频分段 + + Attributes: + title: 分段标题（可选） + content: 分段内容（必填） + start_index: 起始字幕索引 + end_index: 结束字幕索引 + """ + title: Optional[str] = Field( + None, + description="分段标题，简洁描述型（5-15字）" + ) + content: str = Field( + ..., + description="分段内容，合并后的字幕文本" + ) + start_index: int = Field( + ..., + ge=1, + description="起始字幕索引（从1开始）" + ) + end_index: int = Field( + ..., + ge=1, + description="结束字幕索引" + ) + + @field_validator('title') + @classmethod + def validate_title_length(cls, v: Optional[str]) -> Optional[str]: + """验证标题长度（5-15字）""" + if v is not None and len(v) > 15: + raise ValueError(f"标题长度不能超过15字，当前为{len(v)}字") + return v + + model_config = { + "json_schema_extra": { + "examples": [ + { + "title": "机器学习概述", + "content": "机器学习是人工智能的一个分支...", + "start_index": 1, + "end_index": 10 + } + ] + } + } + + +class SegmentationResult(BaseModel): + """分段结果 + + Attributes: + video_type: 视频类型（知识类/故事类/访谈类/新闻类/vlog 等） + segments: 分段列表 + actual_segment_count: 实际分段数量 + """ + video_type: str = Field( + ..., + description="识别出的视频类型" + ) + segments: List[VideoSegment] = Field( + default_factory=list, + description="分段列表" + ) + actual_segment_count: int = Field( + ..., + ge=0, + description="实际分段数量" + ) + + model_config = { + "json_schema_extra": { + "examples": [ + { + "video_type": "知识类", + "segments": [ + { + "title": "机器学习概述", + "content": "...", + "start_index": 1, + "end_index": 10 + } + ], + "actual_segment_count": 5 + } + ] + } + } + + +# ============================================================================ +# Video Type Detection +# ============================================================================ + +# 视频类型关键词映射 +VIDEO_TYPE_KEYWORDS = { + "知识类": [ + "教程", "课程", "教学", "入门", "基础", "进阶", "实战", + "编程", "开发", "技术", "学习", "教育", "培训", "讲座" + ], + "故事类": [ + "影视", "电影", "电视剧", "动漫", "动画", "小说", "故事", + "剧情", "短片", "纪录片" + ], + "访谈类": [ + "访谈", "对话", "采访", "专访", "聊天", "脱口秀", "播客" + ], + "新闻类": [ + "新闻", "报道", "资讯", "时事", "热点", "快讯" + ], + "vlog": [ + "vlog", "日常", "生活", "旅行", "美食", "探店", + "分享", "体验" + ] +} + +# 默认视频类型 +DEFAULT_VIDEO_TYPE = "知识类" + + +# ============================================================================ +# Semantic Segmentation Engine +# ============================================================================ + +class SemanticSegmentationEngine: + """语义分段引擎 + + 使用 LLM 实现智能视频字幕语义分段。 + + 功能： + - 识别视频类型（从元数据推断） + - 构建包含类型信息的 Prompt + - 调用 LLM 进行语义分段 + - 验证和返回分段结果 + + Attributes: + ai_client: AI 客户端（可选，暂未实现完整 LLM 调用） + """ + + def __init__(self, ai_client: Optional[object] = None): + """初始化语义分段引擎 + + Args: + ai_client: AI 客户端实例（预留，后续计划实现） + """ + self.ai_client = ai_client + + def _infer_video_type(self, metadata: dict) -> str: + """从元数据推断视频类型 + + Args: + metadata: 视频元数据，包含 title, category, tags 等字段 + + Returns: + 推断出的视频类型，默认为"知识类" + + Examples: + >>> engine = SemanticSegmentationEngine() + >>> engine._infer_video_type({"title": "Python教程"}) + '知识类' + >>> engine._infer_video_type({"category": "影视"}) + '故事类' + """ + if not metadata: + return DEFAULT_VIDEO_TYPE + + # 提取所有文本字段 + title = metadata.get("title", "") or "" + category = metadata.get("category", "") or "" + tags = metadata.get("tags", []) or [] + description = metadata.get("description", "") or "" + + # 合并所有文本用于匹配 + all_text = " ".join([title, category, description] + list(tags)) + + # 如果没有有效文本，返回默认类型 + if not all_text or all_text.isspace(): + return DEFAULT_VIDEO_TYPE + + # 按优先级匹配类型关键词 + for video_type, keywords in VIDEO_TYPE_KEYWORDS.items(): + for keyword in keywords: + if keyword in all_text: + logger.debug( + f"Detected video type '{video_type}' " + f"from keyword '{keyword}'" + ) + return video_type + + # 未匹配到任何类型，返回默认 + logger.debug("No specific video type detected, using default") + return DEFAULT_VIDEO_TYPE + + def _build_prompt(self, content: str, metadata: dict) -> str: + """构建语义分段 Prompt + + Args: + content: 字幕内容 + metadata: 视频元数据 + + Returns: + 构建好的 Prompt 字符串 + + Examples: + >>> engine = SemanticSegmentationEngine() + >>> prompt = engine._build_prompt("字幕内容", {"title": "教程"}) + >>> "视频类型" in prompt or "知识类" in prompt + True + """ + # 推断视频类型 + video_type = self._infer_video_type(metadata) + title = metadata.get("title", "") or "" + + # 构建 Prompt + total_chars = len(content) + prompt = f"""你是一个专业的视频内容分析师。请将以下字幕内容按语义分段。 + +【视频信息】 +类型：{video_type} +标题：{title} +总字数：{total_chars}字 + +【字幕内容】 +{content} + +【分段要求】 +1. 建议将内容分为 5-8 个段落（可根据语义完整性灵活调整） +2. 每个段落约 1000-2000 字（为保持语义完整性可以突破此范围） +3. 在话题转换处创建新分段，不要混入不同话题的内容 +4. **重要**：保持句子完整性，分段边界应放在自然停顿处（句号。、问号？、感叹号！、省略号……等标点符号后） + - 基于语义停顿判断，不仅依赖标点符号 + - 保持内容的连贯性和可读性 +5. 保持论点完整性，不要在前提-结论结构中间截断 +6. 为主题/章节类型的分段生成简洁标题（5-15字），论点/观点类型分段不加标题 + +【输出格式】 +必须返回以下 JSON 格式（不要添加任何额外文本）： +```json +{{ + "video_type": "{video_type}", + "segments": [ + {{ + "title": "分段标题（如果此分段不需要标题则设为 null）", + "content": "分段内容", + "start_index": 0, + "end_index": 123 + }} + ], + "actual_segment_count": 6 +}} +``` + +【注意事项】 +- 优先保证语义完整性，不要为了凑数而强行截断 +- 标题应简洁准确，如「机器学习概述」「梯度下降原理」 +- 如果某个分段无法生成合适的标题（如过渡性内容），将 title 设为 null""" + return prompt + + def segment( + self, + subtitles: List[dict], + metadata: dict + ) -> SegmentationResult: + """执行语义分段（占位方法，后续计划实现） + + Args: + subtitles: 字幕列表，每个元素包含 index, start, end, text + metadata: 视频元数据 + + Returns: + 分段结果 + """ + # 暂时返回默认结果 + video_type = self._infer_video_type(metadata) + + # TODO: 在后续计划中实现完整的 LLM 调用 + return SegmentationResult( + video_type=video_type, + segments=[], + actual_segment_count=0 + ) diff --git a/tests/test_semantic_engine.py b/tests/test_semantic_engine.py index 928dcf3..322682b 100644 --- a/tests/test_semantic_engine.py +++ b/tests/test_semantic_engine.py @@ -1,6 +1,9 @@ """测试语义分段引擎""" import pytest +import json +from unittest.mock import Mock +from pydantic import ValidationError from kb_workflow.srt.semantic_engine import ( VideoSegment, @@ -127,3 +130,384 @@ class TestPydanticModels: 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 + + +@pytest.mark.skip(reason="属于计划 13-05（中文句边界检测），不在计划 13-01 范围内") +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