"""Chunking - 文档分块策略 提供两种分块策略： - TextChunker: 按字符数分块，带重叠 - StructuralChunker: 按文档结构（标题/段落）分块，适用于 Markdown/HTML """ from __future__ import annotations import logging import re import uuid from dataclasses import dataclass, field from typing import TypeAlias from agentkit.memory.base import MetadataDict logger = logging.getLogger(__name__) # 分块元数据：source_doc/position/char_count/chunking_strategy/heading/heading_level # — 全部为原始标量（str/int）。 ChunkMetadata: TypeAlias = MetadataDict # _split_by_headings 返回的节段结构。 SectionInfo: TypeAlias = dict[str, str | int] @dataclass class Chunk: """文档分块""" chunk_id: str content: str metadata: ChunkMetadata = field(default_factory=dict) def __post_init__(self) -> None: if "source_doc" not in self.metadata: self.metadata["source_doc"] = "" if "position" not in self.metadata: self.metadata["position"] = 0 def to_dict(self) -> dict[str, object]: return { "chunk_id": self.chunk_id, "content": self.content, "metadata": self.metadata, } class TextChunker: """按字符数分块，带重叠 适用于纯文本文档，按固定字符数切分，相邻块之间有重叠区域。 """ def __init__( self, chunk_size: int = 1000, chunk_overlap: int = 200, separator: str = "\n\n", ): """ Args: chunk_size: 每个块的最大字符数 chunk_overlap: 相邻块之间的重叠字符数 separator: 优先分割符 """ if chunk_overlap >= chunk_size: raise ValueError( f"chunk_overlap ({chunk_overlap}) must be less than chunk_size ({chunk_size})" ) self._chunk_size = chunk_size self._chunk_overlap = chunk_overlap self._separator = separator def chunk( self, text: str, source_doc_id: str = "", metadata: ChunkMetadata | None = None, ) -> list[Chunk]: """将文本分块 Args: text: 待分块文本 source_doc_id: 源文档 ID metadata: 附加元数据 Returns: Chunk 列表 """ if not text.strip(): return [] # 先尝试按分隔符分割 segments = self._split_by_separator(text) # 合并小段，切分大段 chunks_text = self._merge_and_split(segments) base_meta = dict(metadata or {}) base_meta["source_doc"] = source_doc_id base_meta["chunking_strategy"] = "text" chunks = [] for i, chunk_text in enumerate(chunks_text): chunk_meta = dict(base_meta) chunk_meta["position"] = i chunk_meta["char_count"] = len(chunk_text) chunks.append( Chunk( chunk_id=str(uuid.uuid4()), content=chunk_text, metadata=chunk_meta, ) ) return chunks def _split_by_separator(self, text: str) -> list[str]: """按分隔符分割文本""" segments = text.split(self._separator) # 过滤空段 return [s.strip() for s in segments if s.strip()] def _merge_and_split(self, segments: list[str]) -> list[str]: """合并小段，切分大段""" result: list[str] = [] current: list[str] = [] current_len = 0 for segment in segments: seg_len = len(segment) # 如果单个段超过 chunk_size，需要进一步切分 if seg_len > self._chunk_size: # 先把当前累积的段输出 if current: result.append(self._separator.join(current)) current = [] current_len = 0 # 切分大段 for sub in self._split_large_segment(segment): result.append(sub) continue # 如果加入当前段会超过 chunk_size，先输出当前累积 if current_len + seg_len + len(self._separator) > self._chunk_size and current: result.append(self._separator.join(current)) # 保留重叠部分 overlap_text = self._separator.join(current) overlap_start = max(0, len(overlap_text) - self._chunk_overlap) overlap_segments = self._get_overlap_segments( overlap_text[overlap_start:], segments ) current = overlap_segments current_len = sum(len(s) for s in current) + len(self._separator) * max( 0, len(current) - 1 ) current.append(segment) current_len += seg_len + len(self._separator) if current: result.append(self._separator.join(current)) return result def _split_large_segment(self, segment: str) -> list[str]: """切分超大段""" result = [] start = 0 while start < len(segment): end = start + self._chunk_size # 尝试在句子边界切分 if end < len(segment): # 查找最近的句子结束符 for sep in ["。", ".", "！", "!", "？", "?", "\n"]: last_sep = segment.rfind(sep, start + self._chunk_size // 2, end) if last_sep > start: end = last_sep + len(sep) break result.append(segment[start:end].strip()) start = end - self._chunk_overlap if start <= 0 and end >= len(segment): break if start < 0: start = 0 return [r for r in result if r] def _get_overlap_segments(self, overlap_text: str, segments: list[str]) -> list[str]: """从重叠文本中提取完整段""" # 简化实现：将重叠文本作为一个段 if overlap_text.strip(): return [overlap_text.strip()] return [] class StructuralChunker: """按文档结构分块 适用于 Markdown 和 HTML 等有标题结构的文档。 按标题层级分块，每个标题下的内容作为一个块。 如果某个块超过 chunk_size，则回退到 TextChunker 继续切分。 """ def __init__( self, chunk_size: int = 1000, chunk_overlap: int = 200, heading_levels: int = 3, ): """ Args: chunk_size: 每个块的最大字符数 chunk_overlap: 回退 TextChunker 时的重叠字符数 heading_levels: 识别的标题层级数（1-6 对应 # 到 ######） """ self._chunk_size = chunk_size self._chunk_overlap = chunk_overlap self._heading_levels = min(max(heading_levels, 1), 6) self._text_chunker = TextChunker( chunk_size=chunk_size, chunk_overlap=chunk_overlap, ) def chunk( self, text: str, source_doc_id: str = "", metadata: ChunkMetadata | None = None, ) -> list[Chunk]: """将文本按结构分块 Args: text: 待分块文本（Markdown 格式） source_doc_id: 源文档 ID metadata: 附加元数据 Returns: Chunk 列表 """ if not text.strip(): return [] sections = self._split_by_headings(text) base_meta = dict(metadata or {}) base_meta["source_doc"] = source_doc_id base_meta["chunking_strategy"] = "structural" chunks = [] position = 0 for section in sections: heading = section["heading"] content = section["content"] level = section["level"] if not content.strip(): continue # 如果内容超过 chunk_size，使用 TextChunker 继续切分 if len(content) > self._chunk_size: sub_chunks = self._text_chunker.chunk( content, source_doc_id=source_doc_id, metadata=metadata, ) for sub in sub_chunks: sub.metadata["position"] = position sub.metadata["heading"] = heading sub.metadata["heading_level"] = level sub.metadata["chunking_strategy"] = "structural" position += 1 chunks.append(sub) else: chunk_meta = dict(base_meta) chunk_meta["position"] = position chunk_meta["heading"] = heading chunk_meta["heading_level"] = level chunk_meta["char_count"] = len(content) chunks.append( Chunk( chunk_id=str(uuid.uuid4()), content=content, metadata=chunk_meta, ) ) position += 1 return chunks def _split_by_headings(self, text: str) -> list[SectionInfo]: """按标题分割 Markdown 文本 Returns: 列表，每项包含 heading, content, level """ lines = text.split("\n") sections: list[SectionInfo] = [] current_heading = "" current_level = 0 current_lines: list[str] = [] heading_pattern = re.compile(r"^(#{1," + str(self._heading_levels) + r"})\s+(.+)$") for line in lines: match = heading_pattern.match(line) if match: # 保存当前节 if current_lines: content = "\n".join(current_lines).strip() if content: sections.append( { "heading": current_heading, "content": content, "level": current_level, } ) # 开始新节 current_heading = match.group(2).strip() current_level = len(match.group(1)) current_lines = [line] else: current_lines.append(line) # 保存最后一节 if current_lines: content = "\n".join(current_lines).strip() if content: sections.append( { "heading": current_heading, "content": content, "level": current_level, } ) # 如果没有标题结构，整体作为一个块 if not sections: sections.append( { "heading": "", "content": text.strip(), "level": 0, } ) return sections