别再死磕分块了！分层RAG才是长文档召回的终极答案

1

pip install langchain langchain-openai umap-learn scikit-learn numpy tiktoken

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168

import numpy as np
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from umap import UMAP
from sklearn.mixture import GaussianMixture
from langchain_core.documents import Document

# ========== 1. 基础配置 ==========
LLM_MODEL = "gpt-4o-mini"
EMBEDDING_MODEL = "text-embedding-3-small"
# 叶子节点分块大小
CHUNK_SIZE = 500
CHUNK_OVERLAP = 50
# 降维目标维度（聚类用推荐5-15维，2-3维仅用于可视化）
REDUCE_DIM = 10
# 递归终止条件：当节点数小于等于该值时，停止递归，生成根节点
STOP_CLUSTER_NUM = 2

# 初始化模型
embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)
llm = ChatOpenAI(model=LLM_MODEL, temperature=0)


# ========== 2. 工具函数封装 ==========
def split_text_to_chunks(raw_text: str) -> list[Document]:
    """第一步：文本分块，生成叶子节点"""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", "。", " ", ""]
    )
    return splitter.create_documents([raw_text])


def get_embeddings(docs: list[Document]) -> np.ndarray:
    """生成向量数组"""
    texts = [doc.page_content for doc in docs]
    return np.array(embeddings.embed_documents(texts))


def reduce_dimensions(vectors: np.ndarray, target_dim: int = REDUCE_DIM) -> np.ndarray:
    """
    降维预处理：默认用UMAP，可替换为PCA
    想换PCA的话，把下面UMAP换成PCA即可：
    from sklearn.decomposition import PCA
    reducer = PCA(n_components=target_dim, random_state=42)
    """
    reducer = UMAP(
        n_components=target_dim,
        n_neighbors=15,  # 平衡局部/全局语义，数值越大越关注全局结构
        min_dist=0.1,    # 簇的紧凑程度，数值越小簇越紧凑
        random_state=42,
        metric="cosine"  # 用余弦距离，适配语义向量
    )
    return reducer.fit_transform(vectors)


def gmm_clustering(low_dim_vectors: np.ndarray) -> list[int]:
    """GMM软聚类，返回每个向量所属的簇标签"""
    # 自动寻找最优聚类数（BIC准则）
    max_clusters = min(len(low_dim_vectors) // 3, 20)
    best_bic = float("inf")
    best_n = 2
    for n in range(2, max_clusters + 1):
        gmm = GaussianMixture(n_components=n, random_state=42, covariance_type="full")
        gmm.fit(low_dim_vectors)
        bic = gmm.bic(low_dim_vectors)
        if bic < best_bic:
            best_bic = bic
            best_n = n
    
    # 用最优簇数做最终聚类
    gmm = GaussianMixture(n_components=best_n, random_state=42)
    labels = gmm.fit_predict(low_dim_vectors)
    return labels.tolist()


def generate_cluster_summary(cluster_docs: list[Document]) -> str:
    """调用LLM生成单个簇的摘要，作为上一层节点"""
    combined_text = "\n\n".join([doc.page_content for doc in cluster_docs])
    prompt = f"""请将以下多段文本合并，生成一段精简的主题摘要，保留核心信息，不要丢失关键概念：
{combined_text}
"""
    return llm.invoke(prompt).content


# ========== 3. 递归构建RAPTOR语义树主函数 ==========
def build_raptor_tree(docs: list[Document], current_level: int = 0) -> dict:
    """
    自底向上递归构建语义树
    返回结构：{level: 层数, nodes: 该层所有节点文档, children: 子层节点}
    """
    print(f"正在构建第 {current_level} 层，当前节点数：{len(docs)}")
    
    # 递归终止条件：节点数足够少，生成根节点
    if len(docs) <= STOP_CLUSTER_NUM:
        root_summary = generate_cluster_summary(docs)
        root_doc = Document(page_content=root_summary, metadata={"level": current_level, "is_root": True})
        return {
            "level": current_level,
            "nodes": [root_doc],
            "children": None
        }
    
    # 1. 向量化
    vectors = get_embeddings(docs)
    
    # 2. UMAP降维（想换算法就在这里改）
    low_dim_vectors = reduce_dimensions(vectors)
    
    # 3. GMM聚类
    cluster_labels = gmm_clustering(low_dim_vectors)
    
    # 4. 按簇分组，生成摘要
    cluster_groups = {}
    for doc, label in zip(docs, cluster_labels):
        cluster_groups.setdefault(label, []).append(doc)
    
    upper_level_docs = []
    for cluster_id, cluster_docs in cluster_groups.items():
        summary = generate_cluster_summary(cluster_docs)
        upper_doc = Document(
            page_content=summary,
            metadata={
                "level": current_level + 1,
                "cluster_id": cluster_id,
                "child_count": len(cluster_docs)
            }
        )
        upper_level_docs.append(upper_doc)
    
    # 递归构建上一层
    upper_tree = build_raptor_tree(upper_level_docs, current_level + 1)
    
    return {
        "level": current_level,
        "nodes": docs,
        "children": upper_tree
    }


# ========== 4. 运行示例 ==========
if __name__ == "__main__":
    # 示例文本（替换为你的长文档/PDF提取文本）
    sample_text = """
    RAG（检索增强生成）是一种结合信息检索和大语言模型的技术框架。
    基础RAG分为文档分块、向量化存储、检索匹配、大模型生成四个核心环节。
    MMR（最大边际相关性）是一种重排算法，用于平衡检索结果的相关性和多样性，避免返回重复内容。
    RRF（倒数排名融合）是一种多路检索结果融合算法，常用于混合关键词检索和向量检索的结果排序。
    MultiVector多表征检索会为同一段文档生成摘要、假设问题等多个向量，提升召回率。
    RAPTOR递归语义树通过分层聚类和摘要构建多层索引，解决长文档跨章节召回的问题。
    向量数据库是RAG的核心存储组件，常见的有Chroma、Milvus、Pinecone等。
    自查询检索（SelfQuery）可以让大模型自动提取元数据过滤条件，实现结构化的精准检索。
    上下文压缩检索可以对召回的文档做二次精简，只保留和问题相关的片段，节省Token。
    """
    
    # 第一步：生成分块叶子节点
    leaf_docs = split_text_to_chunks(sample_text)
    print(f"叶子节点数量：{len(leaf_docs)}")
    
    # 第二步：构建RAPTOR语义树
    raptor_tree = build_raptor_tree(leaf_docs)
    
    # 输出结果
    print("\n=== RAPTOR树构建完成 ===")
    print(f"总层数：{raptor_tree['children']['level'] + 1}")
    root_node = raptor_tree['children']['nodes'][0]
    print(f"根节点摘要：{root_node.page_content}")