§1.3.12

用 50 行 Python 实现 Naive RAG？

手写练习
  • —用 sentence-transformers + FAISS + OpenAI SDK 搓一个最小 RAG

核心概念

检索增强生成（Retrieval-Augmented Generation, RAG）是一种将大型语言模型（LLM）的参数化知识与外部非参数化知识库相结合的架构。其核心思想是，在生成回答之前，首先从一个大规模的文档语料库（如维基百科、公司内部文档）中检索出与用户问题相关的上下文信息。然后，将这些检索到的信息与原始问题一起打包成一个更丰富的提示（Prompt），并将其输入给 LLM，从而生成更准确、更具事实依据、更不易产生幻觉的回答。

原理与推导

Naive RAG 的工作流程可以分为两个核心阶段：索引（Indexing） 和 检索与生成（Retrieval & Generation）。

1. 索引阶段 (Offline)

此阶段的目标是建立一个可供快速检索的知识库。

  • 分块 (Chunking): 将原始文档分割成更小的、语义完整的文本块（chunks）。
  • 嵌入 (Embedding): 使用一个预训练的文本嵌入模型（如 Sentence-Transformers）将每个文本块转换成一个高维向量。这个向量 viv_i 代表了文本块 cic_i 的语义信息。 vi=Encoder(ci)where vi∈Rdv_i = \text{Encoder}(c_i) \quad \text{where } v_i \in \mathbb{R}^d 其中 dd 是嵌入向量的维度。
  • 建库 (Indexing): 将所有文本块的向量 viv_i 存入一个向量数据库（如 FAISS）。向量数据库会构建一个高效的索引结构（如 Flat, HNSW），以便于后续进行快速的相似度搜索。

2. 检索与生成阶段 (Online)

此阶段在用户提出问题时实时触发。

  • 查询嵌入 (Query Embedding): 使用与索引阶段相同的嵌入模型，将用户的查询 qq 转换为查询向量 vqv_q。 vq=Encoder(q)v_q = \text{Encoder}(q)
  • 向量检索 (Vector Retrieval): 在向量数据库中，计算查询向量 vqv_q 与所有已索引的文档块向量 viv_i 之间的相似度。最常用的相似度度量是余弦相似度（Cosine Similarity）或 L2 距离。然后，找出与 vqv_q 最相似的 Top-K 个文档块。 Similarity(vq,vi)=vq⋅vi∥vq∥∥vi∥\text{Similarity}(v_q, v_i) = \frac{v_q \cdot v_i}{\|v_q\| \|v_i\|} 检索操作可以表示为： {c1,c2,...,cK}=arg top-kci∈CorpusSimilarity(Encoder(q),Encoder(ci))\{c_1, c_2, ..., c_K\} = \text{arg top-k}_{c_i \in \text{Corpus}} \text{Similarity}(\text{Encoder}(q), \text{Encoder}(c_i))
  • 上下文增强 (Context Augmentation): 将检索到的 Top-K 个文档块拼接成一个连续的文本字符串，作为“上下文”（Context）。
  • 生成 (Generation): 构建一个结构化的提示，将上下文和原始问题结合起来，然后将其喂给一个大型语言模型（如 GPT 系列）。
    text
     1Prompt Template:
     2"请根据以下上下文信息来回答用户的问题。如果上下文中没有相关信息，请说你不知道。
     3

     4上下文:
     5{retrieved_chunks}
     6

     7问题:
     8{user_query}
     9

    10回答:"
    LLM 基于这个增强后的提示生成最终答案。这个过程可以形式化地表示为，模型在给定查询 qq 和检索到的上下文 CretrievedC_{retrieved} 的条件下，生成答案 AA 的概率： P(A∣q,Cretrieved)P(A | q, C_{retrieved})

算法复杂度

  • 索引阶段:
    • 嵌入: 假设有 NN 个文档块，平均长度为 LL，嵌入模型的复杂度近似为 O(N⋅L)O(N \cdot L)。
    • FAISS IndexFlatL2 (暴力搜索) 建库: O(N⋅d)O(N \cdot d)，其中 dd 是嵌入维度。
  • 推理阶段:
    • 查询嵌入: O(Lq)O(L_q)，其中 LqL_q 是查询长度。
    • FAISS IndexFlatL2 检索: O(N⋅d)O(N \cdot d)。对于大型数据库，会使用 HNSW 等近似最近邻（ANN）算法，其检索复杂度可降至近似 O(log⁡N)O(\log N)。
    • LLM 生成: 复杂度与生成文本的长度和模型架构相关，通常是性能瓶颈。

代码实现

以下代码使用 sentence-transformers 进行嵌入，faiss-cpu 进行向量检索，openai SDK 调用 LLM，完整实现了 Naive RAG 的核心逻辑。

python
 1import numpy as np
 2import faiss
 3import openai
 4from sentence_transformers import SentenceTransformer
 5

 6# --- 0. 环境设置 ---
 7# 需要提前设置 OpenAI API 密钥，例如: export OPENAI_API_KEY="your_key_here"
 8# 或者在代码中直接设置: openai.api_key = "your_key_here"
 9# 确保已安装库: pip install numpy faiss-cpu openai sentence-transformers
10try:
11    client = openai.OpenAI()
12except openai.OpenAIError:
13    print("OpenAI API 密钥未设置，请确保已配置环境变量 OPENAI_API_KEY")
14    exit()
15

16# --- 1. 索引阶段 (Indexing) ---
17

18# 知识库：一些关于太阳系的事实
19knowledge_base = [
20    "水星是太阳系中最小的行星，也是最靠近太阳的行星。",
21    "金星是太阳系中最热的行星，其表面温度可达462摄氏度。",
22    "地球是我们目前所知唯一存在生命的行星。",
23    "火星因其表面的氧化铁而呈现红色，被称为“红色星球”。",
24    "木星是太阳系中最大的行星，是一个巨大的气体巨星。",
25    "土星以其壮观的行星环而闻名，这些环主要由冰粒组成。",
26    "天王星是一颗冰巨星，它以奇特的侧躺方式绕太阳公转。",
27    "海王星是距离太阳最远的行星，气候极端恶劣。"
28]
29

30# 加载嵌入模型
31# 这是一个轻量级但效果不错的多语言嵌入模型
32embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
33

34# 将知识库中的每个文档转换为向量
35# 为什么这样做：将文本的语义信息编码为计算机可以处理的数字向量
36document_embeddings = embedding_model.encode(knowledge_base)
37

38# 获取向量维度
39d = document_embeddings.shape[1]
40

41# 创建 FAISS 索引
42# 为什么用 IndexFlatL2：这是一个简单的基于 L2 距离（欧氏距离）的暴力检索引擎。
43# 对于小型知识库，它足够快且 100% 准确。
44index = faiss.IndexFlatL2(d)
45

46# 将文档向量添加到索引中
47# 为什么这样做：将所有知识库向量加载到检索引擎中，以便后续进行快速搜索。
48index.add(document_embeddings)
49

50# --- 2. 检索与生成函数 (RAG Function) ---
51

52def naive_rag(query: str, k: int = 2):
53    """
54    一个简单的 RAG 实现，结合了检索和生成两个阶段。
55    """
56    # --- 检索阶段 (Retrieval) ---
57    
58    # 1. 将用户查询转换为向量
59    query_embedding = embedding_model.encode([query])
60    
61    # 2. 在 FAISS 索引中搜索最相似的 k 个文档
62    # index.search 返回两个数组：D (distances) 和 I (indices)
63    # 为什么这样做：这是 RAG 的核心，从知识库中找到与问题最相关的信息。
64    distances, indices = index.search(query_embedding, k)
65    
66    # 3. 根据索引获取相关的文档内容
67    retrieved_chunks = [knowledge_base[i] for i in indices[0]]
68    
69    # --- 生成阶段 (Generation) ---
70    
71    # 4. 构建增强的提示 (Augmented Prompt)
72    context = "\n".join(retrieved_chunks)
73    prompt = f"""
74    请根据以下上下文信息来回答用户的问题。如果上下文中没有足够信息，请回答“根据现有知识，我无法回答该问题”。
75

76    上下文:
77    {context}
78

79    问题: {query}
80

81    回答:
82    """
83    
84    # 5. 调用 LLM 进行生成
85    # 为什么这样做：利用 LLM 强大的语言理解和生成能力，基于我们提供的上下文来合成一个自然、准确的答案。
86    response = client.chat.completions.create(
87        model="gpt-3.5-turbo",
88        messages=[
89            {"role": "system", "content": "你是一个有用的问答助手。"},
90            {"role": "user", "content": prompt}
91        ]
92    )
93    
94    return response.choices[0].message.content
95

96# --- 3. 示例调用 ---
97if __name__ == "__main__":
98    query1 = "太阳系中最大的行星是哪个？"
99    answer1 = naive_rag(query1)
100    print(f"问题: {query1}\n答案: {answer1}\n")
101

102    query2 = "哪个行星有漂亮的环？"
103    answer2 = naive_rag(query2)
104    print(f"问题: {query2}\n答案: {answer2}\n")
105

106    query3 = "苹果公司的创始人是谁？"
107    answer3 = naive_rag(query3)
108    print(f"问题: {query3}\n答案: {answer3}\n")

工程实践

  • 使用场景: Naive RAG 是构建企业级智能问答系统、文档摘要、客服机器人等应用最快速、最常见的起点。它适用于任何需要模型基于特定、可控的知识源进行回答的场景。
  • 超参数选择:
    • chunk_size 和 chunk_overlap: 这是文本分块的关键。chunk_size 通常在 256 到 1024 个 token 之间。较小的块能提供更精确的匹配，但可能丢失全局上下文；较大的块保留更多上下文，但可能引入噪声。chunk_overlap (如 10-20% 的 chunk_size) 能确保块之间的语义连续性。
    • top_k: 通常选择 2 到 5。太小可能错过关键信息，太大可能引入不相关的上下文，干扰 LLM 的判断，并可能超出模型的上下文窗口限制。
    • embedding_model: 模型的选择是效果和成本的权衡。all-MiniLM-L6-v2 速度快、资源占用小，适合快速原型。对于更高要求，可选用 bge-large-zh-v1.5 (中文) 或 text-embedding-3-large (OpenAI) 等更强大的模型。
  • 性能/显存/吞吐:
    • 检索性能: 对于百万级以上的文档库，IndexFlatL2 的暴力搜索（O(N⋅d)O(N \cdot d)）会变得非常慢。生产环境必须使用近似最近邻（ANN）索引，如 FAISS 的 IndexHNSWFlat 或 IndexIVFPQ。HNSW 提供了速度和精度的良好平衡，而 IVFPQ 通过量化技术极大压缩了内存占用。
    • 显存: 嵌入向量会占用大量内存（文档数 * 维度 * 4字节）。一个百万文档、768维的库就需要约 3GB 内存。使用 float16 或 FAISS 的标量/乘积量化（Scalar/Product Quantization）可以显著降低内存占用。
    • 吞吐: RAG 的瓶颈通常在 LLM 的生成。可以通过批量处理查询、使用更快的推理引擎（如 vLLM, TensorRT-LLM）或使用更小的 LLM 来提高吞吐。

常见误区与边界情况

  • 误区1: RAG 能完全杜绝幻觉。
    • 纠正: RAG 极大减少了幻觉，但不能完全杜绝。LLM 仍可能忽略提供的上下文、错误地综合信息，或在上下文信息不足时退回到其内部的参数化知识。
  • 误区2: 检索到的就是最相关的。
    • 纠正: 语义相似不完全等同于“回答问题所需”。用户的提问方式、嵌入模型的局限性都可能导致检索失败，即最相关的信息没有排在 Top-K 中。这被称为“大海捞针”（Needle-in-a-Haystack）问题。
  • 边界情况1: 知识库外的问题 (Out-of-Domain)。
    • 表现: 当用户提问的问题在知识库中完全不存在时（如代码中的“苹果公司创始人”），系统会检索到一些不相关的文档。
    • 处理: 一个设计良好的 RAG 系统应该能识别这种情况。关键在于 Prompt 设计，明确指示 LLM 在上下文不足时承认“不知道”，而不是强行编造答案。
  • 边界情况2: "Lost in the Middle" 现象。
    • 表现: 研究表明，当大量上下文被提供给 LLM 时，模型更关注上下文的开头和结尾部分，而中间部分的信息容易被忽略。
    • 处理: 优化 top_k 的大小，避免提供过多无关上下文。更高级的方法包括使用重排（Re-ranking）模型，在检索后用一个更强大的模型对 Top-K 结果进行重排序，将最关键的文档放在开头或结尾。
  • 常见面试追问:
    • "如何评估一个 RAG 系统的好坏？"
      • 回答要点: 评估分为两个层面。1) 组件评估: 单独评估检索器的召回率（Recall）和精确率（Precision）；2) 端到端评估: 使用 RAGAS 等框架，评估生成答案的 Faithfulness (忠实度，是否基于上下文)、Answer Relevancy (答案相关性)、Context Precision/Recall (上下文相关度)。最终，人工评估（Human Evaluation）是黄金标准。
    • "如果知识库需要频繁更新，你如何设计系统？"
      • 回答要点: 设计一个增量索引（Incremental Indexing）流程。可以每天或每小时运行一个批处理任务，只对新增或更新的文档进行嵌入和索引更新，而不是每次都重建整个索引库。需要管理好文档 ID 与向量索引的映射关系。
    • "Naive RAG 有什么缺点？如何改进？"
      • 回答要点: Naive RAG 的缺点包括：单次检索可能不准；检索和生成是解耦的。改进方向包括：1) 多轮检索/迭代式RAG: 进行多轮“检索-生成-反思”循环；2) 查询重写(Query Rewriting): 使用 LLM 优化或分解用户的原始查询，以获得更好的检索结果；3) 混合检索(Hybrid Search): 结合向量检索和传统的关键词检索（如 BM25），提高召回率。
相关题目