RAPTOR 递归文档树优化策略

这是什么

RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval) 是一种递归式的文档检索增强技术，通过构建层次化的文档树结构来优化 RAG 系统。

传统 RAG 系统将文档切分成扁平的小块进行检索，容易丢失全局上下文。RAPTOR 通过递归聚类和摘要生成，构建多层文档树：

叶子节点：原始文档块
中间节点：由聚类生成的摘要（Level 1, Level 2, Level 3...）

这样在检索时，既能找到细粒度的具体内容，也能找到粗粒度的主题摘要。

有什么用

多层次检索：同时支持细节级别（叶子节点）和主题级别（摘要节点）的检索
长文档理解：适合处理书籍、技术手册、论文等长文档
全局视角：高层摘要提供整体理解，避免信息碎片化
语义聚合：自动发现和聚合相似内容，减少重复信息

示例代码

python 复制代码

import dotenv
import numpy as np
import pandas as pd
import umap
import weaviate
from langchain_community.document_loaders import UnstructuredFileLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_weaviate import WeaviateVectorStore
from sklearn.mixture import GaussianMixture
from weaviate.auth import AuthApiKey

dotenv.load_dotenv()

# 1. 初始化组件
RANDOM_SEED = 224
embd = HuggingFaceEmbeddings(
    model_name="thenlper/gte-small",
    encode_kwargs={"normalize_embeddings": True},
)
model = ChatOpenAI(model="gpt-3.5-turbo-16k", temperature=0)
db = WeaviateVectorStore(
    client=weaviate.connect_to_wcs(
        cluster_url="your-cluster-url",
        auth_credentials=AuthApiKey("your-api-key"),
    ),
    index_name="RaptorRAG",
    text_key="text",
    embedding=embd,
)

# 2. 加载和分割文档
loaders = [
    UnstructuredFileLoader("./文档1.txt"),
    UnstructuredFileLoader("./文档2.txt"),
]
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=0,
)
docs = []
for loader in loaders:
    docs.extend(loader.load_and_split(text_splitter))

# 3. 递归构建文档树（核心函数）
def recursive_embed_cluster_summarize(
    texts: list[str], level: int = 1, n_levels: int = 3,
) -> dict[int, tuple[pd.DataFrame, pd.DataFrame]]:
    """
    递归地嵌入、聚类和总结文本，构建文档树
    """
    results = {}

    # 嵌入、聚类、总结当前层级
    df_clusters, df_summary = embed_cluster_summarize_texts(texts, level)
    results[level] = (df_clusters, df_summary)

    # 递归处理下一层
    unique_clusters = df_summary["cluster"].nunique()
    if level < n_levels and unique_clusters > 1:
        new_texts = df_summary["summaries"].tolist()
        next_level_results = recursive_embed_cluster_summarize(
            new_texts, level + 1, n_levels
        )
        results.update(next_level_results)

    return results

# 4. 构建文档树
leaf_texts = [doc.page_content for doc in docs]
results = recursive_embed_cluster_summarize(leaf_texts, level=1, n_levels=3)

# 5. 合并所有层次的文本
all_texts = leaf_texts.copy()
for level in sorted(results.keys()):
    summaries = results[level][1]["summaries"].tolist()
    all_texts.extend(summaries)

# 6. 存入向量数据库
db.add_texts(all_texts)

# 7. 检索（可以检索任意层次的内容）
retriever = db.as_retriever(search_type="mmr")
search_docs = retriever.invoke("你的问题")

流程图

scss 复制代码

原始文档 (流浪地球.txt, 电商产品数据.txt, 项目API文档.md)
    │
    ▼
┌───────────────────┐
│  文本分割          │
│  (chunk_size=500) │
└─────────┬─────────┘
          │
          ▼
    原始文档块
    (Leaf Nodes)
          │
          ▼
┌───────────────────┐
│  文本嵌入          │
│  (Embeddings)     │
└─────────┬─────────┘
          │
    ┌─────┴─────┐
    ▼           ▼
┌───────┐  ┌───────────┐
│ UMAP  │  │ 聚类分析   │
│ 降维  │  │ (GMM+BIC) │
└───┬───┘  └─────┬─────┘
    │            │
    └─────┬──────┘
          ▼
    ┌──────────────────┐
    │  全局聚类         │
    │  (宏观主题)       │
    └────────┬─────────┘
             │
             ▼
    ┌──────────────────┐
    │  局部聚类         │
    │  (细化主题)       │
    └────────┬─────────┘
             │
             ▼
    ┌──────────────────┐
    │  LLM 摘要生成    │
    │  (每个聚类生成   │
    │   一个摘要)       │
    └────────┬─────────┘
             │
             ▼
       Level 1 摘要
             │
             ▼
       [重复上述流程]
             │
             ▼
       Level 2 摘要
             │
             ▼
       Level 3 摘要
             │
    ┌────────┴────────┐
    ▼                 ▼
Leaf Nodes      Summary Nodes
(原始块)          (各级摘要)
    │                 │
    └────────┬────────┘
             ▼
      向量化并存入 DB
             │
             ▼
       MMR/相似度检索
             │
             ▼
        返回多层次结果

核心函数说明

函数	功能
`global_cluster_embeddings()`	UMAP 全局降维
`local_cluster_embeddings()`	UMAP 局部降维
`get_optimal_clusters()`	使用 BIC 确定最优聚类数
`gmm_cluster()`	高斯混合模型聚类
`embed_cluster_summarize_texts()`	嵌入+聚类+LLM 摘要
`recursive_embed_cluster_summarize()`	递归构建文档树