文档分块

演示 RAG 中将长文档切分为语义块的核心策略：固定长度分块（带重叠）、按句子分块、 递归分块。分块质量直接影响检索召回率。 对应教程：第6卷《RAG》第1章 分块与索引。

文件路径: examples/volume-06-rag\chunking.py

---

源代码

"""文本分块（Chunking）示例。

演示 RAG 中将长文档切分为语义块的核心策略：固定长度分块（带重叠）、按句子分块、
递归分块。分块质量直接影响检索召回率。

对应教程：第6卷《RAG》第1章 分块与索引。
"""

from __future__ import annotations

import re
from dataclasses import dataclass


@dataclass(frozen=True)
class Chunk:
    """一个文本块。

    Attributes:
        text: 块文本。
        start: 在原文中的起始字符偏移。
        end: 结束字符偏移。
        index: 块序号。
    """

    text: str
    start: int
    end: int
    index: int


def split_by_chars(text: str, chunk_size: int = 100, overlap: int = 20) -> list[Chunk]:
    """按固定字符数分块（带重叠）。

    Args:
        text: 原始文本。
        chunk_size: 每块字符数。
        overlap: 相邻块重叠字符数。

    Returns:
        ``Chunk`` 列表。

    Raises:
        ValueError: 参数非法时。
    """
    if chunk_size <= 0:
        raise ValueError("chunk_size 必须为正")
    if overlap < 0 or overlap >= chunk_size:
        raise ValueError("overlap 必须 >= 0 且 < chunk_size")
    chunks: list[Chunk] = []
    if not text:
        return chunks
    step = chunk_size - overlap
    idx = 0
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(Chunk(text=text[start:end], start=start, end=end, index=idx))
        idx += 1
        if end >= len(text):
            break
        start += step
    return chunks


# 中英文句末标点
_SENTENCE_END_RE = re.compile(r"(?<=[。！？!?\.])\s*")


def split_sentences(text: str) -> list[str]:
    """按句子切分（兼容中英文标点）。

    Args:
        text: 原始文本。

    Returns:
        句子列表（保留标点，去除纯空白项）。
    """
    if not text:
        return []
    parts = _SENTENCE_END_RE.split(text)
    return [s.strip() for s in parts if s.strip()]


def split_by_sentences(
    text: str,
    max_chars: int = 120,
    overlap_sentences: int = 0,
) -> list[Chunk]:
    """按句子聚合分块，每块不超过 ``max_chars`` 字符。

    Args:
        text: 原始文本。
        max_chars: 每块最大字符数。
        overlap_sentences: 相邻块重叠的句子数。

    Returns:
        ``Chunk`` 列表。
    """
    if max_chars <= 0:
        raise ValueError("max_chars 必须为正")
    sentences = split_sentences(text)
    if not sentences:
        return []
    chunks: list[Chunk] = []
    current: list[str] = []
    current_len = 0
    idx = 0

    def flush(start: int) -> int:
        nonlocal idx
        if not current:
            return start
        chunk_text = "".join(current)
        end = start + len(chunk_text)
        chunks.append(Chunk(text=chunk_text, start=start, end=end, index=idx))
        idx += 1
        return end

    pos = 0
    for sent in sentences:
        if current and current_len + len(sent) > max_chars:
            pos = flush(pos)
            # 处理重叠：保留最后 overlap_sentences 句
            if overlap_sentences > 0:
                tail = current[-overlap_sentences:]
                current = list(tail)
                current_len = sum(len(s) for s in current)
                pos -= current_len
            else:
                current = []
                current_len = 0
        current.append(sent)
        current_len += len(sent)
    flush(pos)
    return chunks


def split_recursive(text: str, separators: list[str] | None = None) -> list[str]:
    """递归分块：依次尝试多种分隔符，尽量在更高级别切分。

    Args:
        text: 原始文本。
        separators: 分隔符优先级列表，默认 ``["\n\n", "\n", "。", " "]``。

    Returns:
        非空文本块列表。
    """
    if separators is None:
        separators = ["\n\n", "\n", "。", " "]
    if not text:
        return []
    if not separators:
        return [text] if text.strip() else []

    sep = separators[0]
    rest = separators[1:]
    if sep == "":
        return [text]

    parts = text.split(sep)
    result: list[str] = []
    for part in parts:
        if not part.strip():
            continue
        # 若该部分已经足够短，直接保留；否则用下一级分隔符继续切
        if len(part) <= 80 or not rest:
            result.append(part.strip())
        else:
            result.extend(split_recursive(part, rest))
    return result


def main() -> None:
    """演示三种分块策略。"""
    text = (
        "AI Agent 是能感知环境、自主决策并执行任务的智能体。\n"
        "它由大语言模型驱动，结合工具调用与记忆机制。\n\n"
        "RAG 通过检索增强生成，缓解幻觉。"
    )
    print("== 按字符分块 ==")
    for c in split_by_chars(text, chunk_size=30, overlap=10):
        print(f"  [{c.index}] {c.text!r} ({c.start}-{c.end})")

    print("\n== 按句子分块 ==")
    for c in split_by_sentences(text, max_chars=40):
        print(f"  [{c.index}] {c.text!r}")

    print("\n== 递归分块 ==")
    for i, c in enumerate(split_recursive(text)):
        print(f"  [{i}] {c!r}")


if __name__ == "__main__":
    main()