RAG 检索增强生成实战
构建企业级 RAG 系统,结合向量数据库实现知识库问答与文档检索。
RAG (Retrieval-Augmented Generation) 是将大模型与外部知识库结合的关键技术,可以让 LLM 基于私有数据给出准确回答。本文将详细介绍 RAG 系统的构建方法。
RAG 基础原理
工作流程
┌─────────────────────────────────────────────────────────────┐
│ RAG 系统架构 │
├─────────────────────────────────────────────────────────────┤
│ │
│ 用户提问 ──────┐ │
│ │ │
│ ▼ │
│ ┌──────────────┐ │
│ │ Embedding │ 将问题转为向量 │
│ └──────┬───────┘ │
│ │ │
│ ▼ │
│ ┌──────────────┐ ┌──────────────┐ │
│ │ 向量检索 │ ◄─── │ 向量数据库 │ │
│ └──────┬───────┘ │ (知识库) │ │
│ │ └──────────────┘ │
│ ▼ │
│ ┌──────────────┐ │
│ │ 上下文组装 │ 检索结果 + 原始问题 │
│ └──────┬───────┘ │
│ │ │
│ ▼ │
│ ┌──────────────┐ │
│ │ LLM │ 生成最终回答 │
│ └──────┬───────┘ │
│ │ │
│ ▼ │
│ 答案 │
│ │
└─────────────────────────────────────────────────────────────┘
核心组件
| 组件 | 功能 | 常用工具 |
|---|---|---|
| Embedding 模型 | 文本向量化 | OpenAI, BGE, Jina |
| 向量数据库 | 存储和检索向量 | Chroma, Milvus, Pinecone |
| 文档处理 | 切分和预处理 | LangChain, LlamaIndex |
| LLM | 生成回答 | GPT-4, Claude, Qwen |
快速搭建 RAG
使用 LangChain
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA
# 1. 加载文档
loader = PyPDFLoader("knowledge.pdf")
documents = loader.load()
# 2. 文档切分
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
separators=["\n\n", "\n", "。", "!", "?", ";", " "]
)
splits = text_splitter.split_documents(documents)
# 3. 创建向量存储
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(
documents=splits,
embedding=embeddings,
persist_directory="./chroma_db"
)
# 4. 创建 RAG 链
llm = ChatOpenAI(model="gpt-4")
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=vectorstore.as_retriever(search_kwargs={"k": 3})
)
# 5. 提问
result = qa_chain.invoke("文档中提到了什么技术?")
print(result["result"])
使用 LlamaIndex
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.llms.openai import OpenAI
# 加载文档
documents = SimpleDirectoryReader("./docs").load_data()
# 创建索引
index = VectorStoreIndex.from_documents(documents)
# 创建查询引擎
query_engine = index.as_query_engine(
llm=OpenAI(model="gpt-4"),
similarity_top_k=3
)
# 查询
response = query_engine.query("项目的主要功能是什么?")
print(response)
向量数据库选型
主流方案对比
| 数据库 | 类型 | 特点 | 适用场景 |
|---|---|---|---|
| Chroma | 嵌入式 | 轻量、易用 | 开发测试 |
| Milvus | 分布式 | 高性能、可扩展 | 生产环境 |
| Pinecone | 云服务 | 免运维 | 快速上线 |
| Weaviate | 自托管/云 | 多模态支持 | 复杂场景 |
| Qdrant | 自托管/云 | 高性能 | 大规模部署 |
Chroma 使用
import chromadb
from chromadb.utils import embedding_functions
# 创建客户端
client = chromadb.PersistentClient(path="./chroma_db")
# 使用 OpenAI Embedding
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
api_key="your-api-key",
model_name="text-embedding-3-small"
)
# 创建集合
collection = client.create_collection(
name="documents",
embedding_function=openai_ef
)
# 添加文档
collection.add(
documents=["文档内容1", "文档内容2", "文档内容3"],
metadatas=[{"source": "doc1"}, {"source": "doc2"}, {"source": "doc3"}],
ids=["id1", "id2", "id3"]
)
# 检索
results = collection.query(
query_texts=["查询内容"],
n_results=3
)
print(results)
Milvus 使用
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType
# 连接 Milvus
connections.connect("default", host="localhost", port="19530")
# 定义 Schema
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1536)
]
schema = CollectionSchema(fields, "文档集合")
# 创建集合
collection = Collection("documents", schema)
# 创建索引
index_params = {
"metric_type": "COSINE",
"index_type": "IVF_FLAT",
"params": {"nlist": 1024}
}
collection.create_index("embedding", index_params)
# 插入数据
import numpy as np
embeddings = np.random.random((100, 1536)).tolist()
texts = [f"文档{i}" for i in range(100)]
collection.insert([texts, embeddings])
collection.load()
# 检索
search_params = {"metric_type": "COSINE", "params": {"nprobe": 10}}
results = collection.search(
data=[query_embedding],
anns_field="embedding",
param=search_params,
limit=5,
output_fields=["text"]
)
Embedding 模型选择
主流模型对比
| 模型 | 厂商 | 维度 | 中文支持 | 成本 |
|---|---|---|---|---|
| text-embedding-3-small | OpenAI | 1536 | 良好 | $0.02/1M |
| text-embedding-3-large | OpenAI | 3072 | 良好 | $0.13/1M |
| BGE-large-zh | BAAI | 1024 | 最佳 | 免费 |
| Jina-embeddings-v2 | Jina AI | 768 | 良好 | 免费 |
本地 Embedding
from sentence_transformers import SentenceTransformer
# 加载本地模型
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
# 生成 Embedding
texts = ["第一段文本", "第二段文本"]
embeddings = model.encode(texts)
print(embeddings.shape) # (2, 1024)
文档处理策略
切分策略
from langchain.text_splitter import (
RecursiveCharacterTextSplitter,
MarkdownTextSplitter,
PythonCodeTextSplitter
)
# 通用文本切分
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50,
length_function=len,
separators=["\n\n", "\n", "。", "!", "?", " ", ""]
)
# Markdown 文档
md_splitter = MarkdownTextSplitter(
chunk_size=500,
chunk_overlap=50
)
# 代码文档
code_splitter = PythonCodeTextSplitter(
chunk_size=1000,
chunk_overlap=100
)
元数据提取
from langchain_community.document_loaders import PyPDFLoader
from langchain.schema import Document
def load_with_metadata(file_path: str) -> list[Document]:
loader = PyPDFLoader(file_path)
pages = loader.load()
documents = []
for page in pages:
# 添加元数据
page.metadata.update({
"file_name": file_path.split("/")[-1],
"page_number": page.metadata.get("page", 0),
"file_type": "pdf",
"created_at": "2025-01-01"
})
documents.append(page)
return documents
高级检索策略
混合检索
from langchain.retrievers import BM25Retriever, EnsembleRetriever
# 向量检索
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
# BM25 关键词检索
bm25_retriever = BM25Retriever.from_documents(documents)
bm25_retriever.k = 5
# 混合检索
ensemble_retriever = EnsembleRetriever(
retrievers=[vector_retriever, bm25_retriever],
weights=[0.6, 0.4] # 向量检索权重更高
)
# 使用混合检索
docs = ensemble_retriever.get_relevant_documents("查询内容")
重排序 (Reranking)
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CohereRerank
# 初始检索
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
# Cohere 重排序
compressor = CohereRerank(model="rerank-multilingual-v2.0", top_n=3)
# 压缩检索器
compression_retriever = ContextualCompressionRetriever(
base_compressor=compressor,
base_retriever=base_retriever
)
docs = compression_retriever.get_relevant_documents("查询内容")
自查询检索
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo
# 定义元数据字段
metadata_field_info = [
AttributeInfo(name="category", description="文档类别", type="string"),
AttributeInfo(name="date", description="发布日期", type="string"),
AttributeInfo(name="author", description="作者", type="string"),
]
# 自查询检索器
retriever = SelfQueryRetriever.from_llm(
llm=ChatOpenAI(model="gpt-4"),
vectorstore=vectorstore,
document_contents="技术文档",
metadata_field_info=metadata_field_info,
)
# 自动解析查询条件
docs = retriever.get_relevant_documents(
"找出 2024 年张三写的关于机器学习的文章"
)
生产级 RAG 架构
完整系统架构
from typing import List, Dict, Any
from dataclasses import dataclass
@dataclass
class RAGConfig:
embedding_model: str = "text-embedding-3-small"
llm_model: str = "gpt-4"
chunk_size: int = 500
chunk_overlap: int = 50
top_k: int = 5
rerank: bool = True
class ProductionRAG:
def __init__(self, config: RAGConfig):
self.config = config
self._init_components()
def _init_components(self):
"""初始化组件"""
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import Milvus
self.embeddings = OpenAIEmbeddings(model=self.config.embedding_model)
self.llm = ChatOpenAI(model=self.config.llm_model)
self.vectorstore = Milvus(
embedding_function=self.embeddings,
connection_args={"host": "localhost", "port": "19530"}
)
def ingest(self, documents: List[str], metadata: List[Dict] = None):
"""文档入库"""
from langchain.text_splitter import RecursiveCharacterTextSplitter
splitter = RecursiveCharacterTextSplitter(
chunk_size=self.config.chunk_size,
chunk_overlap=self.config.chunk_overlap
)
chunks = []
for i, doc in enumerate(documents):
splits = splitter.split_text(doc)
for j, chunk in enumerate(splits):
chunks.append({
"text": chunk,
"metadata": {
**(metadata[i] if metadata else {}),
"chunk_index": j
}
})
self.vectorstore.add_texts(
texts=[c["text"] for c in chunks],
metadatas=[c["metadata"] for c in chunks]
)
def query(self, question: str) -> Dict[str, Any]:
"""查询"""
# 1. 检索
docs = self.vectorstore.similarity_search(
question,
k=self.config.top_k
)
# 2. 重排序 (可选)
if self.config.rerank:
docs = self._rerank(question, docs)
# 3. 生成回答
context = "\n\n".join([doc.page_content for doc in docs])
prompt = f"""基于以下上下文回答问题。
上下文:
{context}
问题:{question}
回答:"""
response = self.llm.invoke(prompt)
return {
"answer": response.content,
"sources": [doc.metadata for doc in docs]
}
def _rerank(self, query: str, docs: List) -> List:
"""重排序文档"""
# 使用 LLM 或专用重排序模型
return docs[:3] # 简化实现
# 使用
rag = ProductionRAG(RAGConfig())
rag.ingest(["文档1内容", "文档2内容"], [{"source": "doc1"}, {"source": "doc2"}])
result = rag.query("文档中提到了什么?")
print(result["answer"])
评估与优化
评估指标
from ragas import evaluate
from ragas.metrics import (
faithfulness,
answer_relevancy,
context_precision,
context_recall
)
# 准备评估数据
eval_data = {
"question": ["问题1", "问题2"],
"answer": ["回答1", "回答2"],
"contexts": [["上下文1"], ["上下文2"]],
"ground_truths": [["标准答案1"], ["标准答案2"]]
}
# 评估
result = evaluate(
eval_data,
metrics=[faithfulness, answer_relevancy, context_precision, context_recall]
)
print(result)
总结
构建高质量 RAG 系统的关键:
- 文档处理: 合理的切分策略和元数据管理
- Embedding 选择: 中文场景推荐 BGE,英文用 OpenAI
- 检索策略: 混合检索 + 重排序提升准确率
- 向量数据库: 开发用 Chroma,生产用 Milvus
- 持续优化: 基于评估指标迭代改进
下一篇将介绍 Function Calling 与 Tool Use 的实现方法。