text_splitter.py

"""
文本分块器：把长文档切成小块（chunk），每块单独做 embedding 和检索
"""

from document_loader import Document


def split_text(text: str, chunk_size: int = 500, chunk_overlap: int = 100) -> list[str]:
    """
    把一段文本按固定大小切分，相邻块之间有重叠。

    参数：
        text:          要切分的文本（一整个字符串）
        chunk_size:    每块的最大字符数（默认 500）
        chunk_overlap: 相邻块重叠的字符数（默认 100）

    返回：
        字符串列表，每个元素是一个 chunk

    相当于 C++:
        vector<string> split_text(string text, int chunk_size, int chunk_overlap)
    """
    # 空文本直接返回空列表
    if not text.strip():
        return []

    # 步进 = chunk_size - overlap
    # 相当于 C: int step = chunk_size - chunk_overlap;
    step = chunk_size - chunk_overlap

    chunks = []  # vector<string> chunks;

    # range(start, stop, step) 生成等差数列
    # 相当于 C: for (int i = 0; i < len; i += step)
    for i in range(0, len(text), step):
        # text[i : i+chunk_size] 是切片，取第 i 到第 i+chunk_size-1 个字符
        # 相当于 C++: text.substr(i, chunk_size)
        # Python 切片不会越界，超出部分自动截断
        chunk = text[i : i + chunk_size]

        # 跳过纯空白的块（strip() 去掉首尾空白后判断）
        if chunk.strip():
            chunks.append(chunk)

    return chunks


def split_documents(
    documents: list[Document],
    chunk_size: int = 500,
    chunk_overlap: int = 100,
) -> list[Document]:
    """
    把 Document 列表中的每个文档切分成多个小块 Document。

    参数：
        documents:     document_loader 返回的 Document 列表
        chunk_size:    每块最大字符数
        chunk_overlap: 重叠字符数

    返回：
        新的 Document 列表，每个 Document 是一个 chunk
        metadata 里额外记录 chunk_index（第几块）和 source_filepath（原文件路径）
    """
    all_chunks = []  # vector<Document> all_chunks;

    for doc in documents:
        # 调用上面的 split_text 把文本切成字符串列表
        text_chunks = split_text(doc.content, chunk_size, chunk_overlap)

        # enumerate() 同时拿到下标和值
        # 相当于 C: for (int i = 0; i < text_chunks.size(); i++)
        #           string chunk = text_chunks[i];
        for i, chunk in enumerate(text_chunks):
            # 创建新的 Document，metadata 继承原文件信息 + 新增 chunk 信息
            chunk_doc = Document(
                content=chunk,
                metadata={
                    # **doc.metadata 是"解包"，把原 dict 的所有键值对复制过来
                    # 相当于 C++: new_map.insert(old_map.begin(), old_map.end())
                    **doc.metadata,
                    "chunk_index": i,             # 这是第几块（从 0 开始）
                    "chunk_total": len(text_chunks),  # 这个文档一共被切成几块
                },
            )
            all_chunks.append(chunk_doc)

    print(f"分块完成：{len(documents)} 个文档 → {len(all_chunks)} 个块")
    return all_chunks


# ============ 测试代码 ============
if __name__ == "__main__":
    from document_loader import load_documents

    # 第一步：加载文档
    docs = load_documents("F:/AI_Program/enterprise-kb")

    # 第二步：分块
    chunks = split_documents(docs, chunk_size=500, chunk_overlap=100)

    # 打印结果看看
    for chunk in chunks:
        meta = chunk.metadata
        print(f"[{meta['filename']}] 块 {meta['chunk_index']+1}/{meta['chunk_total']}")
        print(f"  长度: {len(chunk.content)} 字符")
        print(f"  预览: {chunk.content[:80]}...")
        print()