手写一个 RAG：从零搭建可溯源的检索增强问答系统

太子釢
10 阅读10分钟

不依赖 LangChain / LlamaIndex，用 numpy + OpenAI 兼容 API + 本地 BGE，把 RAG 的每个环节亲手搭一遍。

一、背景

RAG（Retrieval-Augmented Generation，检索增强生成）是让大模型回答私有/长尾知识的主流方案：先把文档切成块向量化，收到问题时检索相关片段，连同问题一起交给 LLM 生成带引用的回答。本项目的核心逻辑（切块、检索、引用溯源）全部从零手写，不依赖 LangChain / LlamaIndex，把 RAG 的每个环节亲手搭一遍。

具体的模型与工具（LLM、embedding 模型、向量检索实现）不做前置介绍，用到时再逐个讲。

语料是 10 篇原创的 iOS 开发知识笔记（Swift 基础、UI、并发、内存管理、网络等，见仓库 corpus/ 目录），RAG 的目标就是让模型能回答这些私有知识。

二、系统总览

flowchart LR
    subgraph 离线索引
        A[语料 corpus] --> B[chunk 切块]
        B --> C[embed 向量化]
        C --> D[(VectorStore numpy 矩阵)]
    end
    subgraph 在线问答
        Q[用户问题] --> E[embed 向量化]
        E --> F[retrieve 检索 top k]
        F --> G[answer 拼 prompt 调 LLM]
        G --> H[回答与来源溯源]
    end
    D -. 相似度计算 .-> F

整条系统分离线索引和在线问答两条流水线，先逐个环节看它们在做什么、原理是什么：

离线索引（提前把文档准备好，只跑一次）

  • chunk 切块：把整篇文档拆成一个个独立的小片段，每个片段单独向量化。这样检索时能精确命中某一段内容。块的大小不宜过大，否则：
    1. 不精确：一个向量是整块内容的"平均语义"，块太大时向量化后精确度下降，检索容易捞错；
    2. 超过 embedding 模型限制：embedding 模型有最大输入长度，文本超长会截断或报错。
  • embed 向量化：把文本变成数字。embedding 模型把每个 chunk 编码成一个向量（一长串浮点数），语义相近的文本向量在空间中靠得近、无关的远——这样"相似度"就有了可计算的数值。
  • 存进矩阵：所有向量拼成一个 numpy 矩阵（N 行 × 每行一个向量），N 是块数。这个矩阵就是"可检索的索引"。

在线问答（用户每次提问实时跑）

  • 问题向量化：用同一个 embedding 模型把用户的问题也编码成向量，保证和文档向量在同一个语义空间里，才能比距离。
  • retrieve 检索：算问题向量和矩阵每一行的相似度，找出最相关的 top-k 个块。这是整条链的核心——决定了 LLM 能不能"看到"正确的内容。
  • answer 生成：把检索到的块编号塞进 prompt，连同问题交给 LLM，要求用 [n] 标注依据，生成带引用的回答。

三、环境配置与技术栈

开发环境：uv（包管理 + 虚拟环境二合一）+ Python 3.12。

技术栈与参数：

环节选型
LLM 生成DeepSeek deepseek-v4-flash（OpenAI 兼容协议）
向量化本地 BGE bge-base-zh-v1.5（CPU 可跑，免外部 API）
存储与检索纯 numpy 矩阵

用到的库（环境配置）：

  • openai：调 DeepSeek API（改 base_url 即可复用）
  • numpy：向量矩阵存储 + 余弦相似度计算
  • sentence-transformers：加载 BGE 模型做向量化
  • transformers：HuggingFace 模型加载库，被 sentence-transformers 依赖
  • torch：深度学习框架，BGE 前向推理的底层计算引擎
  • python-dotenv：从 .env 读 API key，避免硬编码

Intel Mac（x86_64）下 torch 最高只能用 2.2.2。

四、chunk：切块的艺术

切块是 RAG 第一个要设计的环节。目标：每块"信息完整、语义干净、带出身"。设计决策如下：

flowchart TD
    A[整篇文本] --> B[按换行切行 去空行]
    B --> C[短行贪心合并 上限 size 减 overlap]
    C --> D{单行超过上限?}
    D -->|是| E[标点处拆分 优先级 句末 逗号 非标点]
    D -->|否| F[保留]
    E --> G[从第2块起 前置前块尾部 overlap]
    F --> G
    G --> H[组装 Chunk text source index]

几个值得说的决策：

  1. overlap 的意义：纯顺序切块可能会把一句完整的话从中间劈开，前后两半各缺半截语义，检索时两边都"残"。所以让每块结尾重复带上前一块的尾巴（本项目 40 字），保证边界处的语义连贯。实现上两个细节：上限取 size - overlap（每块还要前置前一块的尾巴，留出余量才不超长）；overlap 取自前一块的 base（不含其 overlap，避免逐块放大）。
  2. 先切分、后加 overlap：先按长度/标点切出互不重叠的 base 块，再给每块前置前一块的尾巴。这样切分逻辑只关心"边界切在哪"，重叠是后加的修饰，两件事互不纠缠，边界在哪不因 overlap 变大而漂移。
  3. 标点优先级：句末（。！？；）→ 逗号（，）。中英标点都考虑。外加一个"最大回溯"：无标点或标点离窗口末尾太远（欠填超 10%）就放弃干净边界，直接填满整块，避免切出小碎块。
@dataclass
class Chunk:
    text: str
    source: str   # 来自哪个文件（兼作标题）
    index: int    # 文件内自增编号

五、embed：把语义变成数字

embedding 模型（又称向量模型 / 文本编码器）负责把"一段文本"变成"一串数字"（向量）。它在大规模语料上训练，目标是让语义相近的句子在向量空间里靠得近、语义无关的离得远——所以用向量距离就能度量"两段话有多相关"，这是 RAG 检索的地基。

本项目用的是 BGE 的 bge-base-zh-v1.5（智源研究院开源，中文场景效果好，1.1 亿参数，模型权重约 400MB、CPU 可跑）。加载很简单，借助 sentence-transformers 库一行搞定——它会自动处理下载、分词、编码等琐碎环节，不用自己写模型前向推理代码。每个 chunk 会被编码成一个 768 维向量（base 版模型的输出维度），即 768 个浮点数。项目里把它封装成 embed 函数，一次接收一批文本、批量返回向量：

from functools import lru_cache
from sentence_transformers import SentenceTransformer

@lru_cache(maxsize=1)
def get_model() -> SentenceTransformer:
    return SentenceTransformer("BAAI/bge-base-zh-v1.5")

def embed(texts: list[str]) -> list[list[float]]:
    return get_model().encode(texts, normalize_embeddings=True).tolist()

get_model 上挂着 @lru_cache（本地缓存）：模型加载很重（读几百 MB 权重），每次调用都从磁盘加载太慢，缓存后第一次加载、之后直接复用内存里的模型。

注意 normalize_embeddings=True：向量归一化后长度统一为 1，点积就等价于余弦相似度，检索时才能直接用矩阵乘。

"怎么退钱" vs "退款流程"   → 相似度 0.732（近）
"怎么退钱" vs "今天天气真好" → 相似度 0.295（远）

为什么 embedding 不用 DeepSeek：DeepSeek 官方只提供 chat/reasoner 接口，没有 embedding API。所以 LLM 用 DeepSeek，向量化用本地 BGE——一个管生成，一个管检索，各司其职。

注意 embedding 支持批量：embed(所有块的文本) 一次算完，比逐条调用快一个量级。

六、vector_store + retrieve：检索

初版不引入任何向量数据库，一个 numpy 矩阵搞定：

class VectorStore:
    def add(self, chunk, vector):      # 矩阵加一行 (N, 768)，块存列表
        ...
    def search(self, qv, top_k=3):     # 核心：矩阵乘算所有相似度
        scores = self.vectors @ qv      # 一次算出 N 个分数
        order = np.argsort(-scores)[:top_k]   # 按相似度从大到小排，取前 top_k 个
        return [(chunks[i], scores[i]) for i in order]

vectors @ qv 一次矩阵乘 = 算完所有块和问题的余弦相似度（embed 已归一化，点积即余弦）。排序取 top-k。33 块微秒级，一万块几十毫秒——暴力全查在小数据量下完全够用。

retrieve 在其上做业务层决策：相似度低于阈值（0.4）的块扔掉，防止拿无关内容硬答；结果为空就返回明确的"无结果"，让 answer 不硬编。阈值 0.4 是初步判断，如果不准确可以根据实际情况调整。

def retrieve(store, query, top_k=3, min_score=0.4) -> list[Chunk]:
    hits = store.search(embed([query])[0], top_k=top_k)
    return [chunk for chunk, score in hits if score >= min_score]

先取最相似的 top_k 个，再筛掉低于阈值的——两道过滤，避免拿无关内容硬答。

sequenceDiagram
    participant U as 用户
    participant R as retrieve
    participant VS as VectorStore
    participant E as BGE
    U->>E: 问题文本
    E-->>R: 问题向量
    R->>VS: search
    VS-->>R: 相似度排序的前k个块
    R->>R: 过滤低于阈值的块
    R-->>U: 留下的块（可能为空）

七、answer：带引用的回答

检索到块后，把它们编号并塞进 prompt 生成最终的回答，并要求 LLM 用 [n] 标注依据。核心代码很简单，就是拼 prompt + 调 LLM：

def build_prompt(query, chunks):
    context = "\n\n".join(f"[{i}] {c.text}" for i, c in enumerate(chunks, 1))
    return f"""只依据下面编号的资料回答问题，
用 [n] 标注依据，资料没有的说"资料中没有相关信息"。

资料：
{context}

问题：{query}"""

def answer(query, chunks):
    return AnswerResult(answer=chat(build_prompt(query, chunks)), chunks=chunks)

其中 chat 负责调用 LLM：DeepSeek 兼容 OpenAI 协议，用 OpenAI SDK 改一下 base_url 就能发请求，把 prompt 传给模型、取回回答文本。

实际发出去的 prompt 长这样：

资料：
[1] 循环引用是两个对象互相持有强引用，导致引用计数永远不为 0...
[2] 解决方式是使用 weak 弱引用，不增加引用计数...

问题：循环引用是什么？怎么解决？

规则：
- 只依据资料回答，资料没有的说"资料中没有相关信息"
- 用 [n] 标注依据，紧跟句末，如"内存无法释放 [1]"；不要用括号
- 每个依据点单独一行

关键设计：[n] 表示参考资料的编号。这样不仅方便溯源，也避免了模型自己编造来源（幻觉）。

flowchart LR
    A[检索到的块] --> B[编号后拼 prompt 调 LLM] --> C[回答含编号] --> D[代码映射编号到文档]

实测效果（问"循环引用是什么？怎么解决？"）：

循环引用是两个对象互相持有强引用，导致彼此的引用计数永远不为 0，内存无法释放 [2]。
解决方法一：使用 weak 弱引用，不增加引用计数，对象释放后自动变为 nil [1]。
解决方法二：使用 unowned 无主引用，也不增加计数，但需保证对象一直存在 [1]。

--- 引用来源 ---
[1] 04_arc内存管理.md #1
[2] 04_arc内存管理.md #0

引用和内容能对上：定义在 #0 块（标 [2]），解决方案在 #1 块（标 [1]）。

八、evaluate：用数字说话

肉眼看着准不够，评测才是硬道理。做法：手工造 10 个测试用例 (问题 → 期望文档)——从每篇语料挑一个它该答得上来的问题，期望它被检索到并作为引用来源。例如：

TEST_CASES = [
    ("Swift 里的 Optional 是什么？怎么安全解包？", "01_swift基础.md"),
    ("什么是循环引用？怎么避免？", "04_arc内存管理.md"),
    ("App 的偏好设置该用什么存储？", "07_数据持久化.md"),
    ("TestFlight 怎么分发测试版给用户？", "10_app发布上架.md"),
]

每个用例测两个指标：

  • 检索 hit_rate：期望文档是否出现在 top-k 检索结果里
  • 引用有效：回答引用的来源里是否包含期望文档

评测逻辑：

def evaluate(store) -> list[tuple]:
    results = []
    for query, expected in TEST_CASES:
        chunks = retrieve(store, query)
        hit = expected in {c.source for c in chunks}                 # 指标1：检索命中
        cited = {chunks[n-1].source for n in extract_citation_nums(
            answer(query, chunks).answer)}                           # 指标2：引用里的来源
        results.append((query, expected, hit, expected in cited))
    return results

extract_citation_nums 用正则把回答里的 [n] 提取出来，再映射回块，就能核对引用是否指向期望文档：

import re

def extract_citation_nums(text: str) -> list[int]:
    return [int(n) for n in re.findall(r"\[(\d+)\]", text)]

评测的价值在于：以后调 size / overlap / min_score 参数时，不用再肉眼判断，跑一遍 evaluate 看数字。

九、踩过的坑

  1. Intel Mac 兼容性：torch 最高只出到 2.2.2（x86_64），连带 transformers 锁 4.46.3、numpy 锁 1.26.4。
  2. Markdown 代码块被切断：```swift 代码块在 chunk 边界被截断。不识别 Markdown 结构是简化决策，代价就是这种边界瑕疵。
  3. 依据引用格式不稳定：同样规则下，依据文档编号的输出格式不稳定，需要在 prompt 里加上规范。
  4. 阈值太松放无关内容：0.3 会让完全无关的查询也返回块，调整到 0.4 正合适。

如果这篇文章对你有帮助，欢迎收藏转发。RAG 不难，但值得亲手搭一遍。

avatar
文章
阅读
粉丝