code/04-rag/chunking.py

"""几种切分文档的办法，用 httpx 的文档比较它们切出来的块。

在 AI-Course 目录下运行：python code/04-rag/chunking.py
"""
import re
from pathlib import Path

DOCS = Path("data/httpx-docs")


def load_docs():
    return {str(p.relative_to(DOCS)): p.read_text() for p in sorted(DOCS.rglob("*.md")) if p.name != "LICENSE.md"}


def split_fixed(text, size=800, overlap=100):
    """办法一：每 size 个字符切一刀，相邻两块重叠 overlap 个字符。"""
    chunks, start = [], 0
    while start < len(text):
        chunks.append(text[start:start + size])
        start += size - overlap
    return chunks


def split_by_heading_naive(text):
    """办法二（有问题的版本）：凡是以 # 开头的行都当成标题切开。"""
    parts = re.split(r"(?m)^(?=#{1,3} )", text)
    return [p.strip() for p in parts if p.strip()]


def split_by_heading(text):
    """办法二（修正版）：同样按标题切，但跳过代码块里以 # 开头的注释行。"""
    chunks, current, in_code = [], [], False
    for line in text.splitlines():
        if line.startswith("```"):
            in_code = not in_code
        if not in_code and re.match(r"#{1,3} ", line) and current:
            chunks.append("\n".join(current).strip())
            current = []
        current.append(line)
    if current:
        chunks.append("\n".join(current).strip())
    return [c for c in chunks if c]


def split_by_heading_capped(text, max_size=1500):
    """办法三：先按标题切；太长的块再按空行（段落）切开，并在每一小块前面补上所属的标题。"""
    chunks = []
    for section in split_by_heading(text):
        if len(section) <= max_size:
            chunks.append(section)
            continue
        title = section.splitlines()[0] if section.startswith("#") else ""
        current = ""
        for para in section.split("\n\n"):
            if current and len(current) + len(para) > max_size:
                chunks.append(current.strip())
                current = title + "\n\n" if title else ""
            current += para + "\n\n"
        if current.strip():
            chunks.append(current.strip())
    return chunks


def stats(name, chunks):
    sizes = [len(c) for c in chunks]
    print(f"{name}：{len(chunks)} 块，平均 {sum(sizes) // len(sizes)} 字符，最短 {min(sizes)}，最长 {max(sizes)}")


if __name__ == "__main__":
    docs = load_docs()
    splitters = [
        ("固定长度", split_fixed),
        ("按标题（有问题）", split_by_heading_naive),
        ("按标题（修正）", split_by_heading),
        ("按标题+限长", split_by_heading_capped),
    ]
    for name, splitter in splitters:
        stats(name, [c for text in docs.values() for c in splitter(text)])

    text = docs["advanced/timeouts.md"]
    print("\n固定长度切法，timeouts.md 的第 2 块：")
    print(split_fixed(text)[1])

    print("\n有问题的按标题切法，timeouts.md 的各块开头：")
    for c in split_by_heading_naive(text):
        print(f"  [{len(c):4d} 字符] {c.splitlines()[0]}")
    print("\n修正后的按标题切法，timeouts.md 的各块开头：")
    for c in split_by_heading(text):
        print(f"  [{len(c):4d} 字符] {c.splitlines()[0]}")