13.5 RAG Pipeline 代码实战
生活类比:RAG 就像一个开卷考试的学生。不背书(不靠模型自身知识),而是带着参考书(知识库),考试时先翻书找到相关段落(检索),再结合题目用自己的话回答(生成)。开卷考试的好处是不会"编答案"(减少幻觉),但前提是你得会翻书(检索准确)。
本节从零搭建一个完整的 RAG Pipeline,覆盖文档加载、分块、向量化、检索、重排序、生成的全链路,最后对接流式输出。
13.5.1 RAG Pipeline 全貌
一个完整的 RAG Pipeline 包含两个阶段:
阶段一:索引构建(离线,一次性或定期执行)
原始文档 → 加载 → 分块 → 向量化 → 存入向量数据库
阶段二:查询回答(在线,每次用户提问执行)
用户问题 → 向量化 → 向量检索 → [可选]重排序 → 组装 Prompt → LLM 生成 → 返回回答每个环节都有坑,下面逐步实现。
13.5.2 完整 RAG Pipeline 实现
python
import os
from typing import List
from langchain.text_splitter import RecursiveCharacterTextSplitter # 文本分块器
from langchain_community.document_loaders import PyPDFLoader, TextLoader # 文档加载器
from langchain_community.vectorstores import Chroma # 向量数据库
from langchain_openai import OpenAIEmbeddings, ChatOpenAI # 嵌入模型和聊天模型
from langchain.retrievers import ContextualCompressionRetriever # 压缩检索器
from langchain.retrievers.document_compressors import CrossEncoderReranker # 重排序器
from langchain_core.prompts import ChatPromptTemplate # Prompt 模板
from langchain_core.output_parsers import StrOutputParser # 输出解析器
from langchain_core.runnables import RunnablePassthrough # 透传组件
from langchain.schema import Document # 文档类型
class RAGPipeline:
"""完整的 RAG Pipeline——从文档到问答的一站式管道"""
def __init__(self, persist_dir: str = "./chroma_db"):
# 初始化嵌入模型:把文本转成向量
self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
# 初始化聊天模型:用 LLM 生成回答
self.llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
self.persist_dir = persist_dir # 向量数据库持久化目录
self.vectorstore = None # 向量存储实例
self.retriever = None # 检索器实例
def load_documents(self, file_paths: List[str]) -> List[Document]:
"""第一步:加载文档——支持 PDF 和纯文本"""
docs = []
for path in file_paths:
if path.endswith('.pdf'): # PDF 文件用 PyPDFLoader
loader = PyPDFLoader(path)
else: # 其他文件用 TextLoader
loader = TextLoader(path, encoding='utf-8')
docs.extend(loader.load()) # 加载并追加到文档列表
print(f"加载了 {len(docs)} 个文档片段")
return docs
def split_documents(self, docs: List[Document],
chunk_size: int = 1000,
chunk_overlap: int = 200) -> List[Document]:
"""
第二步:文档分块——把长文档切成小段
为什么要分块?因为:
1. 嵌入模型有输入长度限制
2. 检索时只取相关段落,不需要整篇文档
3. 分块越小检索越精确,但可能丢失上下文
chunk_overlap 的作用:相邻块有 200 字符重叠,
就像裁缝留缝头——拼接时不至于缺一段。
"""
splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size, # 每块最多 1000 字符
chunk_overlap=chunk_overlap, # 相邻块重叠 200 字符
# 分隔符优先级:先按段落分,再按句子分,最后按字符分
separators=["\n\n", "\n", "。", "!", "?", ";", ",", " ", ""]
)
chunks = splitter.split_documents(docs)
print(f"分成 {len(chunks)} 个 chunk")
return chunks
def build_index(self, chunks: List[Document]):
"""第三步:构建向量索引——把文本块转成向量并存储"""
self.vectorstore = Chroma.from_documents(
documents=chunks, # 文档块
embedding=self.embeddings, # 嵌入模型
persist_directory=self.persist_dir # 持久化路径
)
print(f"索引构建完成,存于 {self.persist_dir}")
def load_index(self):
"""加载已有的向量索引(重启时不需要重新构建)"""
self.vectorstore = Chroma(
persist_directory=self.persist_dir,
embedding_function=self.embeddings
)
def setup_retriever(self, k: int = 5, use_rerank: bool = False):
"""
第四步:设置检索器
k=5:检索 top-5 最相关的文档块
use_rerank:是否启用重排序
重排序的作用:向量检索快但粗,先捞 top-5,
再用 CrossEncoder 精排,选 top-3 最相关的。
就像招聘:先按简历关键词筛 5 份(粗筛),
再仔细看选 3 份(精筛)。
"""
base_retriever = self.vectorstore.as_retriever(
search_kwargs={"k": k} # 先检索 k 个候选
)
if use_rerank:
# 添加重排序器
compressor = CrossEncoderReranker(
model="BAAI/bge-reranker-base", # 开源重排序模型
top_n=3 # 精排后保留 top-3
)
self.retriever = ContextualCompressionRetriever(
base_compressor=compressor, # 重排序器
base_retriever=base_retriever # 底层检索器
)
else:
self.retriever = base_retriever # 不重排序,直接用向量检索结果
def build_chain(self):
"""第五步:构建 RAG 问答链——把检索、Prompt、LLM 串起来"""
template = """基于以下上下文回答问题。如果上下文不足以回答,请说明。
上下文:
{context}
问题:{question}
回答:"""
prompt = ChatPromptTemplate.from_template(template)
# 用 LCEL 语法串联管道(| 是管道操作符)
chain = (
# 第一步:组装输入数据
# retriever 自动检索上下文,RunnablePassthrough 透传 question
{"context": self.retriever, "question": RunnablePassthrough()}
| prompt # 第二步:填入 Prompt 模板
| self.llm # 第三步:调用 LLM
| StrOutputParser() # 第四步:解析输出为字符串
)
return chain
def query(self, question: str) -> dict:
"""第六步:查询——检索 + 生成"""
# 先检索相关文档块(用于展示来源)
docs = self.retriever.invoke(question)
# 构建问答链并执行
chain = self.build_chain()
answer = chain.invoke(question)
return {
"question": question, # 原始问题
"answer": answer, # LLM 生成的回答
"sources": [ # 引用来源
{"content": d.page_content[:200], "metadata": d.metadata}
for d in docs
]
}
async def query_stream(self, question: str):
"""流式查询——逐字返回结果,用户体验更好"""
docs = self.retriever.invoke(question) # 先检索
template = """基于以下上下文回答问题:
上下文:{context}
问题:{question}
回答:"""
prompt = ChatPromptTemplate.from_template(template)
chain = prompt | self.llm # 构建链(不需要输出解析器,直接取原始 chunk)
# 把检索到的文档块拼接成上下文文本
context = "\n\n".join([d.page_content for d in docs])
# 异步流式输出
async for chunk in chain.astream({"context": context, "question": question}):
if chunk.content: # 过滤空内容
yield chunk.content # yield 逐块返回13.5.3 使用示例
python
async def main():
pipeline = RAGPipeline()
# ===== 阶段一:索引构建 =====
# 1. 加载文档
docs = pipeline.load_documents(["docs/article1.txt", "docs/article2.pdf"])
# 2. 分块
chunks = pipeline.split_documents(docs)
# 3. 构建索引
pipeline.build_index(chunks)
# 4. 设置检索器(启用重排序)
pipeline.setup_retriever(k=5, use_rerank=True)
# ===== 阶段二:查询 =====
# 普通查询
result = pipeline.query("什么是 RAG?")
print(f"回答: {result['answer']}")
print(f"来源数: {len(result['sources'])}")
# 流式查询
async for chunk in pipeline.query_stream("RAG 的优点是什么?"):
print(chunk, end="", flush=True) # 逐字打印13.5.4 FastAPI + RAG:把 Pipeline 变成服务
RAG Pipeline 写好后,用 FastAPI 包装成 HTTP 服务,前端就能调用了:
python
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import json
app = FastAPI(title="RAG 问答服务")
pipeline = RAGPipeline()
@app.on_event("startup")
async def startup():
"""服务启动时加载索引"""
pipeline.load_index() # 加载已有索引
pipeline.setup_retriever(k=5, use_rerank=True) # 设置检索器
@app.post("/rag/query")
async def rag_query(request: dict):
"""普通查询接口——一次性返回完整回答"""
question = request["question"] # 从请求体取问题
result = pipeline.query(question) # 执行查询
return result # 返回 JSON 结果
@app.post("/rag/query/stream")
async def rag_query_stream(request: dict):
"""流式查询接口——SSE 逐字推送"""
question = request["question"]
async def event_generator():
"""SSE 事件生成器"""
async for chunk in pipeline.query_stream(question):
# SSE 格式:data: {json}\n\n
yield f"data: {json.dumps({'content': chunk}, ensure_ascii=False)}\n\n"
yield "data: [DONE]\n\n" # 结束标记
# 返回流式响应
return StreamingResponse(
event_generator(),
media_type="text/event-stream" # SSE 的 MIME 类型
)13.5.5 关键参数调优指南
| 参数 | 作用 | 推荐值 | 调优建议 |
|---|---|---|---|
chunk_size | 每个文档块的大小 | 500-1000 | 中文建议偏小,英文可偏大 |
chunk_overlap | 相邻块重叠量 | 50-200 | 为保证上下文不断裂,取 chunk_size 的 10-20% |
k(检索数) | 初始检索数量 | 5-10 | 太少漏检,太多引入噪声 |
top_n(重排序) | 精排后保留数 | 3-5 | 太少信息不足,太多浪费 token |
temperature | LLM 生成温度 | 0.3-0.7 | RAG 建议偏低,减少"创造" |
常见误区
- "分块越大检索越准":不一定。大块包含更多信息,但也引入更多噪声,检索时可能因为"部分相关"而得分虚高。小块更精确,但可能丢失上下文。
- "重排序是必须的":对于小规模知识库(<100 文档),向量检索已经够好,重排序的收益不明显。文档量大时才值得开启。
- "检索结果越多越好":检索 top-20 然后全塞给 LLM,既浪费 token 又稀释了相关信息。精排后保留 3-5 个最相关的通常效果最好。
- "Persist 目录可以随便放":Chroma 的持久化目录如果被多个进程同时写,可能损坏索引。每个 Pipeline 实例应该有独立的持久化目录。
- "流式输出和普通查询的结果不同":不会。流式只是把输出方式改成"逐字返回",最终拼起来和一次性生成的结果是一样的。
本节小结
| 步骤 | 实现方式 | 关键点 |
|---|---|---|
| 文档加载 | PyPDFLoader / TextLoader | 按文件类型选 Loader |
| 分块 | RecursiveCharacterTextSplitter | 中文分隔符优先级很重要 |
| 向量化 | OpenAIEmbeddings | 选对嵌入模型 |
| 存储 | Chroma | 支持持久化,重启不需重建 |
| 检索 | 相似度搜索 | k 值根据数据量调整 |
| 重排序 | CrossEncoderReranker | 文档量大时收益明显 |
| 生成 | ChatPromptTemplate + LLM | Prompt 要明确"基于上下文回答" |
| 流式 | astream + SSE | 用户体验显著提升 |
核心心法:RAG 不是"把文档塞给 LLM",而是"精准检索 + 约束生成"。检索做不好,生成再强也没用。
延伸阅读
- LangChain RAG Tutorial:https://python.langchain.com/docs/tutorials/rag/
- Chroma Vector Database:https://docs.trychroma.com
- LlamaIndex RAG Guide:https://docs.llamaindex.ai/en/stable/understanding/rag/