353 lines
11 KiB
Python
353 lines
11 KiB
Python
"""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
|