deepdata: core product code (backend, frontend, runtime, deepflow module)

runtime/deepdata-python/app/module/rag/infra/document/splitter.py

@@ -0,0 +1,148 @@
+"""
+文档分块器
+
+包含 DocumentSplitter 基类、工厂和 LangChain 实现。
+根据 ProcessType 创建对应的分块策略。
+"""
+from __future__ import annotations
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any, List, Optional
+
+from langchain_text_splitters import (
+    CharacterTextSplitter,
+    RecursiveCharacterTextSplitter,
+)
+
+from app.module.rag.infra.document.types import DocumentChunk
+from app.module.rag.schema.enums import ProcessType
+
+
+# 各 ProcessType 对应的分隔符配置（优先保持较大语义块）
+SEPARATORS_BY_PROCESS_TYPE = {
+    ProcessType.PARAGRAPH_CHUNK: ["\n\n", "\n", " ", ""],
+    ProcessType.SENTENCE_CHUNK: ["\n\n", "\n", "。", "！", "？", ". ", "! ", "? ", " ", ""],
+    ProcessType.DEFAULT_CHUNK: ["\n\n", "\n", " ", ""],
+    ProcessType.CUSTOM_SEPARATOR_CHUNK: None,
+}
+
+
+class DocumentSplitter(ABC):
+    """文档分块器基类
+
+    所有具体的分块器都需要继承此类并实现 split 方法。
+    """
+
+    def __init__(self, chunk_size: int = 500, overlap_size: int = 50):
+        """初始化分块器
+
+        Args:
+            chunk_size: 分块大小
+            overlap_size: 重叠大小
+        """
+        self.chunk_size = chunk_size
+        self.overlap_size = overlap_size
+
+    @abstractmethod
+    async def split(self, text: str, **metadata: Any) -> List[DocumentChunk]:
+        """分割文档
+
+        Args:
+            text: 文档文本
+            **metadata: 额外的元数据
+
+        Returns:
+            分块列表
+        """
+        pass
+
+    def _create_chunk(self, text: str, chunk_index: int, **metadata: Any) -> DocumentChunk:
+        """创建文档分块"""
+        chunk_metadata = {"chunk_index": chunk_index, **metadata}
+        return DocumentChunk(text=text, metadata=chunk_metadata)
+
+
+class LangChainDocumentSplitter(DocumentSplitter):
+    """基于 LangChain 的分块器实现
+
+    根据 ProcessType 选择 RecursiveCharacterTextSplitter 或 CharacterTextSplitter。
+    """
+
+    def __init__(
+        self,
+        process_type: ProcessType,
+        chunk_size: int = 500,
+        overlap_size: int = 50,
+        delimiter: Optional[str] = None,
+    ):
+        super().__init__(chunk_size=chunk_size, overlap_size=overlap_size)
+        self._process_type = process_type
+        self._delimiter = delimiter or "\n\n"
+        self._splitter = self._create_splitter()
+
+    def _create_splitter(self) -> RecursiveCharacterTextSplitter | CharacterTextSplitter:
+        """创建 LangChain 分块器实例"""
+        if self._process_type == ProcessType.LENGTH_CHUNK:
+            return CharacterTextSplitter(
+                chunk_size=self.chunk_size,
+                chunk_overlap=self.overlap_size,
+                length_function=len,
+            )
+
+        separators = SEPARATORS_BY_PROCESS_TYPE.get(self._process_type)
+        if self._process_type == ProcessType.CUSTOM_SEPARATOR_CHUNK:
+            separators = [self._delimiter, "\n", " ", ""]
+        if separators is None:
+            separators = ["\n\n", "\n", " ", ""]
+
+        return RecursiveCharacterTextSplitter(
+            chunk_size=self.chunk_size,
+            chunk_overlap=self.overlap_size,
+            separators=separators,
+            length_function=len,
+        )
+
+    async def split(self, text: str, **metadata: Any) -> List[DocumentChunk]:
+        """分割文档（异步）"""
+        if not text or not text.strip():
+            return []
+
+        texts = await asyncio.to_thread(self._splitter.split_text, text)
+        return [
+            DocumentChunk(text=t, metadata={**metadata, "chunk_index": i})
+            for i, t in enumerate(texts)
+        ]
+
+
+class DocumentSplitterFactory:
+    """文档分块器工厂
+
+    根据处理类型创建对应的分块器实例。
+    """
+
+    @classmethod
+    def create_splitter(
+        cls,
+        process_type: ProcessType,
+        chunk_size: int = 500,
+        overlap_size: int = 50,
+        delimiter: Optional[str] = None,
+    ) -> DocumentSplitter:
+        """创建分块器
+
+        Args:
+            process_type: 处理类型
+            chunk_size: 分块大小
+            overlap_size: 重叠大小
+            delimiter: 自定义分隔符（仅用于 CUSTOM_SEPARATOR_CHUNK）
+
+        Returns:
+            DocumentSplitter 实例
+        """
+        return LangChainDocumentSplitter(
+            process_type=process_type,
+            chunk_size=chunk_size,
+            overlap_size=overlap_size,
+            delimiter=delimiter,
+        )