""" 文本切分器，用于将文档切分成块。 """ from enum import Enum from typing import List, Optional, Tuple, Dict, Any from dataclasses import dataclass, field from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter, TextSplitter from langchain_experimental.text_splitter import SemanticChunker class SplitterType(str, Enum): RECURSIVE = "recursive" SEMANTIC = "semantic" PARENT_CHILD = "parent_child" @dataclass class RecursiveSplitterConfig: chunk_size: int = 500 chunk_overlap: int = 50 separators: List[str] = field(default_factory=lambda: ["\n\n", "\n", "。", "！", "？", " ", ""]) keep_separator: bool = True strip_whitespace: bool = True @dataclass class SemanticSplitterConfig: embeddings: Any buffer_size: int = 1 add_start_index: bool = False breakpoint_threshold_type: str = "percentile" breakpoint_threshold_amount: float = 0.6 # 非 None，切分更积极 number_of_chunks: Optional[int] = None sentence_split_regex: str = r"(?<=[。！？；.!?;])" # 中文友好 min_chunk_size: int = 100 @dataclass class ParentChildSplitterConfig: embeddings: Any # 语义切分（用于父块） semantic_threshold_type: str = "percentile" semantic_threshold_amount: float = 0.6 semantic_buffer_size: int = 1 semantic_min_chunk_size: int = 100 # 子块（递归字符切分） child_chunk_size: int = 400 child_chunk_overlap: int = 50 # ---------- 适配器 ---------- class SemanticChunkerAdapter(TextSplitter): def __init__(self, config: SemanticSplitterConfig, **kwargs): super().__init__(**kwargs) self._config = config self._chunker = SemanticChunker( embeddings=config.embeddings, buffer_size=config.buffer_size, add_start_index=config.add_start_index, breakpoint_threshold_type=config.breakpoint_threshold_type, breakpoint_threshold_amount=config.breakpoint_threshold_amount, number_of_chunks=config.number_of_chunks, sentence_split_regex=config.sentence_split_regex, min_chunk_size=config.min_chunk_size, ) def split_text(self, text: str) -> List[str]: return self._chunker.split_text(text) def split_documents(self, documents: List[Document]) -> List[Document]: result = [] for doc in documents: chunks = self.split_text(doc.page_content) for i, chunk in enumerate(chunks): result.append(Document( page_content=chunk, metadata={**doc.metadata, "chunk_index": i} )) return result # ---------- 工厂函数 ---------- def get_splitter(splitter_type: SplitterType, **kwargs) -> TextSplitter: if splitter_type == SplitterType.RECURSIVE: config = RecursiveSplitterConfig( chunk_size=kwargs.get("chunk_size", 500), chunk_overlap=kwargs.get("chunk_overlap", 50), separators=kwargs.get("separators", ["\n\n", "\n", "。", "！", "？", " ", ""]), ) return RecursiveCharacterTextSplitter( chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap, separators=config.separators, keep_separator=config.keep_separator, strip_whitespace=config.strip_whitespace, ) elif splitter_type == SplitterType.SEMANTIC: embeddings = kwargs.get("embeddings") if embeddings is None: raise ValueError("语义切分器需要提供 'embeddings' 参数") if "config" in kwargs and isinstance(kwargs["config"], SemanticSplitterConfig): config = kwargs["config"] else: config = SemanticSplitterConfig( embeddings=embeddings, buffer_size=kwargs.get("buffer_size", 1), breakpoint_threshold_type=kwargs.get("breakpoint_threshold_type", "percentile"), breakpoint_threshold_amount=kwargs.get("breakpoint_threshold_amount", 0.6), number_of_chunks=kwargs.get("number_of_chunks"), min_chunk_size=kwargs.get("min_chunk_size", 100), ) return SemanticChunkerAdapter(config) elif splitter_type == SplitterType.PARENT_CHILD: raise ValueError("父子切分器应通过 ParentChildSplitter 直接创建") else: raise ValueError(f"不支持的切分器类型: {splitter_type}") # ---------- 父子切分器 ---------- class ParentChildSplitter: """ 切分流程： 1. 语义切分 → 父块 2. 递归字符切分 → 子块 """ def __init__(self, config: ParentChildSplitterConfig): self.config = config # 语义切分（父块） semantic_config = SemanticSplitterConfig( embeddings=config.embeddings, buffer_size=config.semantic_buffer_size, breakpoint_threshold_type=config.semantic_threshold_type, breakpoint_threshold_amount=config.semantic_threshold_amount, min_chunk_size=config.semantic_min_chunk_size, ) self.semantic_splitter = SemanticChunkerAdapter(semantic_config) # 递归字符切分（子块，大小由 child_chunk_size 控制） self.recursive_splitter = RecursiveCharacterTextSplitter( chunk_size=config.child_chunk_size, chunk_overlap=config.child_chunk_overlap, separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""] ) self.parent_to_children: Dict[str, List[str]] = {} self.child_to_parent: Dict[str, str] = {} def split_documents(self, documents: List[Document]) -> Tuple[List[Document], List[Document]]: parent_chunks = [] child_chunks = [] for doc in documents: # Step 1: 语义切分（父块） semantic_blocks = self.semantic_splitter.split_text(doc.page_content) for p_idx, semantic_block in enumerate(semantic_blocks): parent_id = f"parent_{len(parent_chunks)}" parent_doc = Document( page_content=semantic_block, metadata={**doc.metadata, "id": parent_id, "chunk_index": p_idx} ) parent_chunks.append(parent_doc) # Step 2: 递归字符切分（子块） sub_chunks = self.recursive_splitter.split_text(semantic_block) for c_idx, sub_chunk in enumerate(sub_chunks): child_id = f"child_{len(child_chunks)}" child_doc = Document( page_content=sub_chunk, metadata={**doc.metadata, "id": child_id, "parent_id": parent_id, "child_index": c_idx} ) child_chunks.append(child_doc) self.child_to_parent[child_id] = parent_id if parent_id not in self.parent_to_children: self.parent_to_children[parent_id] = [] self.parent_to_children[parent_id].append(child_id) return parent_chunks, child_chunks def get_parent_for_child(self, child_id: str) -> Optional[str]: return self.child_to_parent.get(child_id) def get_children_for_parent(self, parent_id: str) -> List[str]: return self.parent_to_children.get(parent_id, [])