RAPTOR 递归文档树优化策略
这是什么
RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval) 是一种递归式的文档检索增强技术，通过构建层次化的文档树结构来优化 RAG 系统。
传统 RAG 系统将文档切分成扁平的小块进行检索，容易丢失全局上下文。RAPTOR 通过递归聚类和摘要生成，构建多层文档树：
- 叶子节点：原始文档块
- 中间节点：由聚类生成的摘要（Level 1, Level 2, Level 3...）
这样在检索时，既能找到细粒度的具体内容，也能找到粗粒度的主题摘要。
有什么用
- 多层次检索：同时支持细节级别（叶子节点）和主题级别（摘要节点）的检索
- 长文档理解：适合处理书籍、技术手册、论文等长文档
- 全局视角：高层摘要提供整体理解，避免信息碎片化
- 语义聚合：自动发现和聚合相似内容，减少重复信息
示例代码
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("你的问题")
流程图
原始文档 (流浪地球.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() | 递归构建文档树 |
