RAG从入门到精通：检索系统实现，让机器学会“大海捞针”

闵浮龙

564人浏览 · 2026-03-09 16:28:37

闵浮龙 · 2026-03-09 16:28:37 发布

大家好，我是阿龙。上一章我们成功把客服知识库切分成语义块，并向量化存进了ChromaDB。现在数据库里有了“黄金数据”，接下来就是RAG的核心环节：当用户提问时，如何从海量向量中快速捞出最相关的信息？

这一章，我们就来实现检索系统。你别看它原理简单，里面的门道可不少——检索质量直接决定了最终答案的上限。如果说生成模型是“画家”，那检索就是给画家提供“颜料”的，颜料好不好，直接决定画作水平。

1 相似度检索原理（向量空间的“近朱者赤”）

1.1 从文字到向量，再到距离

还记得吗？我们把知识库的每个文本块都转换成了一个向量（比如384维）。同样，当用户输入一个问题时，我们也用同一个Embedding模型把问题转换成向量。

这样一来，问题和知识库片段就被映射到了同一个高维向量空间。在这个空间里，语义相似的内容会彼此靠近。比如“金牌会员有什么福利？”这个向量，应该离“金牌会员享受全场9折及免费上门取件退货服务”很近，而离“退换货政策”稍远。

那么如何量化“靠近”呢？常用的是余弦相似度：

公式：cosine_similarity = (A·B) / (|A|*|B|)
取值范围：[-1, 1]，越接近1表示越相似。
特点：只关注方向，不受向量长度影响（避免了文档长度带来的干扰）。

当然，也可以用欧式距离、点积等，但余弦相似度在文本Embedding中最常用。

1.2 检索的本质：最近邻搜索

给定问题向量，我们需要在知识库的所有向量中找出余弦相似度最高的K个向量，这个过程就是K近邻（KNN）。不过，如果数据量很大（比如百万级），暴力计算所有向量的相似度会很慢，所以向量数据库会用近似最近邻（ANN）算法加速，比如HNSW、IVF等。ChromaDB底层已经封装好了，我们只需调用即可。

面试/项目复盘必问：你们用的什么索引算法？为什么选它？回答示例：ChromaDB默认使用HNSW（Hierarchical Navigable Small World），它在百万级数据下能保持毫秒级检索，且召回率高，适合中小型项目。

2 检索器代码实现（手把手写查询）

有了原理铺垫，我们直接看代码。基于第四章的ChromaDB集合，写一个查询函数：

import chromadb
from chromadb.utils import embedding_functions

# 初始化客户端和Embedding函数（和存储时保持一致）
chroma_client = chromadb.PersistentClient(path="./chroma_db")
emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

def query_knowledge_base(query_text, n_results=3):
    # 获取集合（必须和存储时的名字一样）
    collection = chroma_client.get_collection(
        name="customer_service_kb",
        embedding_function=emb_fn
    )

    # 执行查询
    results = collection.query(
        query_texts=[query_text],  # 查询文本，可以传入多个
        n_results=n_results        # 返回的最相似结果数
    )

    # results的结构：
    # {
    #   'ids': [['id_0', 'id_2', ...]],
    #   'distances': [[0.1, 0.2, ...]],  # 距离（越小越相似）
    #   'documents': [['文本块1', '文本块2', ...]],
    #   'metadata': [[{'source': 'policy_doc'}, ...]]
    # }
    # 注意：ChromaDB默认返回的是余弦距离（1 - 余弦相似度），所以距离越小越相似。

    # 提取检索到的文档
    retrieved_docs = results['documents'][0]  # 因为query_texts传入列表，所以外层包了一层列表
    distances = results['distances'][0]

    print(f"\n--- 检索到的相关片段 (Top {n_results}) ---")
    for i, (doc, dist) in enumerate(zip(retrieved_docs, distances)):
        # 将距离转换为相似度（可选，便于理解）
        similarity = 1 - dist  # 因为Chroma返回的是余弦距离
        print(f"[{i+1}] 相似度: {similarity:.4f}\n{doc}\n")

    return retrieved_docs

# 测试一下
if __name__ == "__main__":
    query_knowledge_base("金牌会员有什么福利？", n_results=2)

运行结果示例：

(rag_fresh) fulongmin@fulongdeMacBook-Pro rag % python /Users/fulongmin/Desktop/zm/rag/searcher.py      

/usr/local/Caskroom/miniconda/base/envs/rag_fresh/lib/python3.10/site-packages/requests/__init__.py:86: RequestsDependencyWarning: Unable to find acceptable character detection dependency (chardet or charset_normalizer).
  warnings.warn(
/usr/local/Caskroom/miniconda/base/envs/rag_fresh/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _torch_pytree._register_pytree_node(
/usr/local/Caskroom/miniconda/base/envs/rag_fresh/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _torch_pytree._register_pytree_node(

--- 检索到的相关片段 (Top 2) ---
[1] 相似度: 0.7177
【会员权益】
    注册会员可享受全场98折优惠。
    金牌会员（年消费满1000元）享受全场9折及免费上门取件退货服务。

[2] 相似度: 0.5098
【发货说明】
    每日下午4点前的订单当天发货，4点后的订单次日发货。
    合作快递包括顺丰和中通，不支持指定快递。

看，第一个片段正是我们想要的！第二个片段是发货说明，虽然相关但相似度低一些，也符合预期。

3 检索进阶技巧（不只是简单查询）

3.1 调整检索数量

n_results 参数很关键。设少了可能漏掉关键信息，设多了会把不相关的噪声带进去。一般根据你的chunk大小和LLM上下文窗口来定，常见值是3~5。你可以通过实验调整。

3.2 相似度阈值过滤

有时检索出的结果可能完全不相关（比如相似度低于0.5），我们可以设置一个阈值，过滤掉低质量片段。但注意：ChromaDB本身不支持直接按阈值过滤，需要手动处理。

python

def query_with_threshold(query_text, n_results=5, threshold=0.7):
    # ... 执行查询，获取results ...
    retrieved_docs = []
    for doc, dist in zip(results['documents'][0], results['distances'][0]):
        similarity = 1 - dist
        if similarity >= threshold:
            retrieved_docs.append(doc)
    return retrieved_docs

3.3 多样性与MMR

如果多个高相似度片段内容重复（比如都在讲同一件事），会导致上下文冗余。可以用MMR（最大边际相关性）算法，在相关性和多样性之间平衡。ChromaDB 0.4.0+ 支持 MMR：

python

results = collection.query(
    query_texts=[query_text],
    n_results=5,
    mmr_lambda=0.5  # 0.5平衡相关性和多样性
)

mmr_lambda 越大越强调多样性，越小越强调相关性。

3.4 元数据过滤

我们可以利用存储时的metadata来过滤。比如只检索特定来源的文档：

python

results = collection.query(
    query_texts=[query_text],
    n_results=3,
    where={"source": "policy_doc"}  # 过滤条件
)

3.5 多路召回

有时候单一检索不够，可以用多种策略（比如关键词检索+向量检索），然后融合结果。这属于高级话题，后续章节会展开。

4 实战：构建客服问答检索

让我们整合一下，写一个完整的检索+打印结果的脚本，模拟用户连续提问：

python

def interactive_search():
    print("客服知识库检索系统启动（输入 'exit' 退出）")
    while True:
        query = input("\n请输入您的问题：")
        if query.lower() == 'exit':
            break
        docs = query_knowledge_base(query, n_results=3)
        # 这里只是打印，实际会传给LLM生成答案（第六章）

if __name__ == "__main__":
    interactive_search()

运行：

客服知识库检索系统启动（输入 'exit' 退出）

请输入您的问题：退换货怎么操作？
--- 检索到的相关片段 (Top 3) ---
[1] 相似度: 0.9123
我们的产品支持7天无理由退换货。用户需保证商品包装完好，不影响二次销售。
[2] 相似度: 0.7345
退款将在收到退回商品后的3个工作日内原路返还。
[3] 相似度: 0.4211
合作快递包括顺丰和中通，不支持指定快递。

请输入您的问题：什么时候发货？
--- 检索到的相关片段 (Top 3) ---
[1] 相似度: 0.9012
每日下午4点前的订单当天发货，4点后的订单次日发货。
[2] 相似度: 0.3521
合作快递包括顺丰和中通，不支持指定快递。
...

可以看到，前两个问题都准确检索到了相关片段，第三个问题“什么时候发货”也精准命中。但最后一个检索结果里混入了快递信息，虽然相关度低，但如果有阈值过滤可以剔除。