AI工程师第六课 - RAG检索增强生成

学习代码记录仓库

第五课中使用OpenAI调用模型DeepSeek 代码,作为基准测试代码：

from openai import OpenAI
from dotenv import load_dotenv
import os

load_dotenv(override=True)

api_key = os.getenv("API_KEY")
base_url = os.getenv("BASE_URL")

client = OpenAI(api_key=api_key, base_url=base_url)

def main():
    prompt = "Hi"
    print(call_model([{"role": "user", "content": prompt}]))

def call_model(messages):
    response = client.chat.completions.create(
        model="deepseek-v4-pro",
        messages=messages,
        stream=False,
        reasoning_effort="high",
        extra_body={"thinking": {"type": "enabled"}}
    )
    return response.choices[0].message.content

if __name__ == "__main__":
    main()

RAG 是什么

LLM 不知道你的私有数据。

你问："公司最新的差旅报销标准是什么？"
LLM："抱歉，我的知识截止到某个日期，无法回答这个问题。"

它不知道你公司的内部文档。让它看过了，它就知道了。

RAG = Retrieval Augmented Generation = 检索增强生成

核心思路：先检索相关文档，再把文档和问题一起喂给 LLM。

不用 RAG：
  用户问题 → LLM → 回答（只能靠训练数据里的知识，可能编造）

用 RAG：
  用户问题 → 检索相关文档 → 文档 + 问题 → LLM → 回答（基于文档，有据可查）

RAG 五步流程

1. 加载（Load）     文档进来
2. 分块（Chunk）    切成小段
3. 向量化（Embed）   文字 → 数字向量
4. 存储（Store）    向量存数据库
5. 检索+生成（Retrieve+Generate） 搜相关片段，喂给 LLM

1. 加载文档

# 读一个文件
with open("README.md", "r") as f:
    document = f.read()
print(document[:200])  

2. 分块（Chunking）

LLM 有上下文窗口限制，不能一次塞整本书。把文档切成小段，每段几百字。

为什么不能切太碎？
  太碎 → 每段信息不完整，检索找不到

为什么不能切太大？
  太大 → 检索精度下降，LLM 窗口装不下

经验值：每块 500-1000 字，块与块之间有 100-200 字重叠

def chunk_text(text, chunk_size=500, overlap=100):
    """简单分块：按字符数切，块之间有重叠"""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap  # 重叠一部分
    return chunks

# 在main中继续调用切分
chunks = chunk_text(document)
print(f"分了 {len(chunks)} 块，每块约 500 字")

生产环境用更智能的分块------按句子、段落边界切，不把一句话拦腰截断。

3. 向量化（Embedding）

计算机不理解"语义"，只理解数字。 Embedding 把文字变成一串数字（向量），含义相近的文字向量也相近。

模型	维度	中文	特点
bge-small-zh-v1.5	512	好	开源本地，轻量，推荐上手
bge-large-zh-v1.5	1024	最好	开源本地，效果最优
m3e-base	768	好	开源本地，社区活跃
text-embedding-3-small	1536	中	OpenAI API，便宜但要联网

用本地 Embedding 模型，免费、离线、数据不出域：

pip install sentence-transformers

from sentence_transformers import SentenceTransformer

# 加载中文 Embedding 模型
model = SentenceTransformer("BAAI/bge-small-zh-v1.5",cache_folder="./models")

def get_embedding(text):
    return model.encode(text, normalize_embeddings=True).tolist()

# 为每个块生成嵌入
embeddings = [get_embedding(chunk) for chunk in chunks]
print(embeddings[0][:10])

知道两段文字有多相似------算两个向量的余弦相似度：

import numpy as np

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# 计算两个块的相似度
similarity = cosine_similarity(embeddings[0], embeddings[1])
print(f"相似度：{similarity:.4f}")

4. 向量数据库

把所有文档块的向量存起来，查询时找最相似的。

数据库	适用场景	特点
Chroma	学习、原型	pip 装直接用，轻量
FAISS	百万级数据	Meta 出品，纯向量检索，快
Milvus	生产、海量数据	分布式，支持混合检索
Pinecone	不想运维	云服务，按量付费

用 Chroma（轻量，pip 装完直接用）：

pip install chromadb

import chromadb

# 创建/打开数据库
chroma_client = chromadb.PersistentClient(path="./chroma_db")
# 清理
chroma_client.delete_collection("my_docs")
# 创建/打开集合
collection = chroma_client.get_or_create_collection(name="my_docs")

# 存储块
collection.add(
  ids=[str(i) for i in range(len(chunks))],
  documents=chunks,
  metadatas=[{"source": "README.md"}] * len(chunks),
  embeddings=embeddings,
)
print(collection.count())

查询时，把用户问题向量化，找最相似的文档,top_k 找最相似的几条结果：

def search(query, top_k=3):
    query_embedding = get_embedding(query)
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
    )
    return results["documents"][0]  # 返回文档列表

# 测试,最相似的三条
results = search("环境配置", top_k=3)
for doc in results:
  print(doc)

5. 检索 + 生成

把搜到的文档和用户问题拼在一起，喂给 LLM：

def rag_answer(question):
    # 1. 检索相关文档
    relevant_docs = search(question)

    # 2. 拼接 Prompt
    context = "\n".join(relevant_docs)
    prompt = f"""
    根据以下参考资料回答问题。如果资料中没有相关信息，如实说不知道。
    参考资料：
    {context}

    问题：{question}

    回答：
    """

    # 3. 调 LLM
    response = call_model([{"role": "user", "content": prompt}])
    return response

print(rag_answer("环境配置"))
print(rag_answer("Python基础知识"))

测试结果基本都是从给的文档里引用的，还是比较靠谱的。

用 LangChain / LlamaIndex 做 RAG

上面是裸写------自己分块、调 Embedding API、存向量、拼接 Prompt、调 LLM。理解了原理，但代码量不小。LangChain 和 LlamaIndex 把这些步骤封装好了。

LangChain

安装：

pip install langchain langchain-openai langchain-huggingface langchain_community langchain_chroma

1. 加载文档

通过DirectoryLoader批量加载指定目录下的某些格式文件，指定加载器TextLoader 解析文件内容。

from langchain_community.document_loaders import DirectoryLoader, TextLoader

loader = DirectoryLoader(".", glob="*.md", loader_cls=TextLoader)
docs = loader.load()

2. 分块

RecursiveCharacterTextSplitter 按语义边界切，不会拦腰截断一句话。

from langchain_text_splitters import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500, chunk_overlap=100, separators=["\n## ", "\n", "。", ".", " "]
)
chunks = text_splitter.split_documents(docs)

3. 向量化（Embedding）

使用langchain二次封装的langchain_huggingface,切换模型时更方便。

from langchain_huggingface import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(
    model_name="BAAI/bge-small-zh-v1.5", cache_folder="./models"
)

4. 存 Chroma

使用langchain二次封装的langchain_chroma，切换数据库更方便。

from langchain_chroma import Chroma

vectorstore = Chroma.from_documents(
    chunks, embeddings, persist_directory="./chroma_langchain"
)

5. 创建模型调用链

使用了langchain_openai对OpenAI模型封装，使用更便捷。构建模板，然后通过langchain特有的写法|管道操作，使得链式调用更简单。

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

llm = ChatOpenAI(model="deepseek-v4-pro", api_key=api_key, base_url=base_url, temperature=0, max_tokens=500)
template = """
根据以下参考资料回答问题。如果资料中没有相关信息，如实说不知道。
参考资料：
{context}

问题：{question}

回答：
"""
prompt = ChatPromptTemplate.from_template(template)
chain = (
    {
        "context": vectorstore.as_retriever(search_kwargs={"k": 3}),
        "question": RunnablePassthrough(),
    }
    | prompt
    | llm
    | StrOutputParser()
)

6. 问答

print(chain.invoke("环境配置"))
print(chain.invoke("Python基础知识"))

LlamaIndex 版

LlamaIndex 比 LangChain 更专注 RAG，API 更简洁。

安装所需的包：

pip install llama-index llama-index-embeddings-huggingface llama-index-vector-stores-chroma llama-index-llms-openai-like

1. 加载文档

from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader(".", required_exts=[".md"]).load_data()

2. 配置向量化模型

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings

Settings.embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5", cache_folder="./models"
)
Settings.chunk_size = 500
Settings.chunk_overlap = 100

3. 建索引（自动分块、向量化和存 Chroma）

存储Chroma 并没有和之前一样存储到本地，直接调用内存存储，程序结束时会自动销毁。如果有需要可以加持久化缓存。

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

4. 配置模型

也是直接封装的OpenAI，注意参数有点变化api_base ≠ base_url，因为用的是deepseek要定义is_chat_model=True.

from llama_index.llms.openai_like import OpenAILike

Settings.llm = OpenAILike(
    api_key=api_key,
    api_base=base_url,
    model="deepseek-v4-pro",
    temperature=0,
    max_tokens=500,
    is_chat_model=True,
)

5. 问答

不再需要构建模板了，LlamaIndex 内部直接封装了。也可以通过参数覆盖默认模板。

query_engine = index.as_query_engine(similarity_top_k=3)

print(query_engine.query("环境配置"))
print(query_engine.query("Python基础知识"))

高级 RAG

基础 RAG 能用，但有些场景不够好。几个常见优化：

LangChain 和LlamaIndex 默认都不做这些高级优化，需要的话要自己手动创建模型，进行配置。

重排序（Rerank）

第一次检索返回粗筛结果，再用更强的模型精排。

基础：用户问题 → Embedding 检索 → Top 3 → 喂 LLM

Rerank：
  用户问题 → Embedding 检索 → Top 20（粗筛）
  → Rerank 模型精排 → Top 3（精选）→ 喂 LLM

效果：检索精度显著提升，成本增加不多

在第一章节手写过自定义实现Rag过程代码，修改方法rag_answer,增加Rerank 过程。

from sentence_transformers import CrossEncoder

# 加载rerank 模型
reranker = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L-6-v2",
    cache_folder="./models",
)


def rag_answer(question):
    # 1. 检索相关文档
    # relevant_docs = search(question)
    relevant_docs = search(question, 10)

    pairs = [(question, doc) for doc in relevant_docs]

    # 精排打分
    scores = reranker.predict(pairs)
    ranked = sorted(zip(scores, relevant_docs), key=lambda x: x[0], reverse=True)

    top_docs = [doc for _, doc in ranked[:3]]
    # 2. 拼接 Prompt
    # context = "\n".join(relevant_docs)
    context = "\n".join(top_docs)

    # 省略剩余代码

本地测试为了方便，使用的小模型。

混合检索

纯向量检索有时会漏掉精确的关键词匹配。

问题："2024 年 Q3 财报"

向量检索        → 可能返回语义相近但无关的"财务"相关内容
关键词检索（BM25）→ 精确匹配"2024 Q3 财报"出现的文档

混合检索 = 向量检索 + 关键词检索，取并集再排序

继续增加混合检索的实现代码,首先要找到关键词，使用jieba对语句分词，然后使用rank_bm25算法，进行匹配得分计算：

from rank_bm25 import BM25Okapi
import jieba

def main():
    # ...省略其他代码

    # 混合检索
    tokenized_chunks = [jieba.lcut(chunk) for chunk in chunks]
    bm25 = BM25Okapi(tokenized_chunks)
    # 检索
    results = hybrid_search(bm25, chunks, "环境配置")
    for doc in results:
        print(doc)

def hybrid_search(bm25, chunks, query, top_k=5):
    # 1. 向量检索
    results = search(query, top_k=top_k)

    # 2. 关键词搜索
    kw_scores = bm25.get_scores(jieba.lcut(query))
    kw_results = sorted(zip(chunks, kw_scores), key=lambda x: x[1], reverse=True)
    kw_results = [c for c, _ in kw_results[:10]]

    # 3. 合并结果
    seen = set(results)
    result = list(results)
    for doc in kw_results:
        if doc not in seen:
            result.append(doc)
    return result[:top_k]

多轮检索

用户问题太模糊，先让 LLM 改写问题再检索。

用户："上次说的那个方案"
  → LLM 改写："关于 XX 项目的技术方案"
  → 检索"XX 项目 技术方案"
  → 喂 LLM 回答

一个主要点就是要有历史记录，通过历史记录来推断用户的意图，然后进行改写。