自动同步: 2026-03-24 02:10

2026-03-24 02:10:04 +08:00
parent c5da76cd04
commit 01885232c3
70 changed files with 2348 additions and 263 deletions
@@ -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,70 @@ class SRTProcessor(BaseProcessor):
    def __init__(
        self,
        state_manager: Optional[WorkflowStateManager] = None,
-        enable_correction: bool = True
+        enable_correction: bool = True,
+        enable_semantic_segmentation: Optional[bool] = None,
+        enable_cache: bool = True
    ):
        """初始化 SRT 处理器

        Args:
            state_manager: 状态管理器
            enable_correction: 是否启用 ASR 纠错（默认 True）
+            enable_semantic_segmentation: 是否启用 LLM 语义分段（None 时从配置读取）
+            enable_cache: 是否启用缓存（默认 True，仅在启用语义分段时有效）
        """
        super().__init__(state_manager)
+
+        # 如果未明确指定，从配置读取
+        if enable_semantic_segmentation is None:
+            settings = get_settings()
+            enable_semantic_segmentation = settings.enable_semantic_segmentation
+            logger.info(
+                f"SRTProcessor using config: "
+                f"semantic_segmentation={'enabled' if enable_semantic_segmentation else 'disabled'}"
+            )
+        elif enable_semantic_segmentation:
+            logger.info("SRTProcessor using explicit parameter: semantic_segmentation=enabled")
+        else:
+            logger.info("SRTProcessor using explicit parameter: semantic_segmentation=disabled")
+
        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:
+                    # 初始化缓存
+                    cache = None
+                    if enable_cache:
+                        from kb_workflow.srt.cache import SemanticCache
+                        cache = SemanticCache()
+
+                    # 创建统一策略，传入 AI 客户端和缓存
+                    strategy = UnifiedSegmentationStrategy(
+                        ai_client=ai_client,
+                        cache=cache
+                    )
+                    self.segmenter = SemanticSegmentationEngine(strategy=strategy)
+                    logger.info(f"Semantic segmentation enabled (cache={'enabled' if cache else 'disabled'})")
+                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 +153,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 +210,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 +227,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 +273,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 +293,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 +313,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')