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

1. 引言为什么需要手写 RAG大语言模型虽然强大但存在知识截止、幻觉和无法引用来源等问题。检索增强生成Retrieval-Augmented GenerationRAG通过先检索再生成的方式把外部知识注入模型从而提升回答的准确性和可溯源性。本文不依赖 LangChain 等框架从零手写一个可运行的 RAG 系统覆盖文档加载、切分、向量化、检索、生成和溯源全流程。手写 RAG 的好处在于你能清楚看到每个环节的输入输出便于调试、定制和教学。下面我们逐步搭建。2. 系统架构与整体流程一个最小可用的 RAG 系统包含以下模块文档加载器读取本地文本或 Markdown 文件。文本切分器把长文档切成适合检索的块Chunk。向量化器把文本块编码为向量。向量存储与检索保存向量并支持相似度检索。提示词组装把检索结果和用户问题拼成 Prompt。生成器调用大模型生成回答。溯源模块把回答关联到原始文档块。整体流程如下flowchart TD A[加载文档] -- B[文本切分] B -- C[向量化] C -- D[向量存储] E[用户提问] -- F[问题向量化] F -- G[相似度检索] D -- G G -- H[组装 Prompt] H -- I[大模型生成] I -- J[返回回答与溯源]3. 环境准备与依赖安装本文使用 Python 3.10核心依赖如下pip install openai numpy说明openai用于调用 Embedding 模型和对话模型兼容 OpenAI 接口的服务均可。numpy用于向量相似度计算。如果你希望完全本地运行可以把 Embedding 和生成模型替换为本地模型如 BGE、Qwen本文以 OpenAI 兼容接口为例便于快速跑通。4. 文档加载与文本切分首先实现文档加载器支持读取纯文本和 Markdown 文件from pathlib import Path def load_document(file_path: str) - str: 读取本地文本或 Markdown 文件内容。 path Path(file_path) if not path.exists(): raise FileNotFoundError(f文件不存在: {file_path}) return path.read_text(encodingutf-8)接下来实现文本切分器。切分策略直接影响检索质量这里采用「按段落优先、按长度兜底」的策略import re def split_text(text: str, chunk_size: int 500, overlap: int 50) - list[str]: 把长文本切成多个块。 - chunk_size: 每个块的最大字符数 - overlap: 相邻块之间的重叠字符数用于保持上下文连贯 # 先按空行切分为段落 paragraphs [p.strip() for p in re.split(r\n\s*\n, text) if p.strip()] chunks [] current for para in paragraphs: # 如果当前块加上新段落会超长先保存当前块 if len(current) len(para) 1 gt; chunk_size and current: chunks.append(current) # 保留尾部 overlap 字符作为下一块的开头保持上下文 current current[-overlap:] if overlap gt; 0 else current (current \n para).strip() if current: chunks.append(current) return chunks切分时保留重叠区域可以避免关键信息恰好落在块边界而被截断。5. 向量化与向量存储向量化负责把文本块转换为向量。这里封装一个 Embedding 客户端from openai import OpenAI class EmbeddingClient: def init(self, api_key: str, base_url: str, model: str text-embedding-3-small): self.client OpenAI(api_keyapi_key, base_urlbase_url) self.model model def embed(self, texts: list[str]) -gt; list[list[float]]: 批量把文本编码为向量。 resp self.client.embeddings.create(modelself.model, inputtexts) return [item.embedding for item in resp.data]向量存储使用简单的内存实现支持添加和检索import numpy as np class VectorStore: def init(self): self.vectors [] # 向量列表 self.metadata [] # 每个向量对应的文本块和来源信息 def add(self, vectors: list[list[float]], metadatas: list[dict]): 批量添加向量及其元数据。 self.vectors.extend(vectors) self.metadata.extend(metadatas) def search(self, query_vector: list[float], top_k: int 3) -gt; list[dict]: 余弦相似度检索返回最相似的 top_k 个结果。 每个结果包含 text、score 和 source 字段。 if not self.vectors: return [] q np.array(query_vector) mat np.array(self.vectors) # 余弦相似度 点积 / (模长乘积) scores (mat q) / (np.linalg.norm(mat, axis1) * np.linalg.norm(q) 1e-9) top_indices np.argsort(scores)[::-1][:top_k] results [] for idx in top_indices: results.append({ text: self.metadata[idx][text], score: float(scores[idx]), source: self.metadata[idx].get(source, ), chunk_index: self.metadata[idx].get(chunk_index, 0), }) return results这里使用余弦相似度衡量相关性分数越高表示越相关。6. 索引构建把文档写入向量库把加载、切分、向量化、存储串起来构建索引def build_index(file_path: str, embed_client: EmbeddingClient, store: VectorStore): 读取文档、切分、向量化并写入向量存储。 text load_document(file_path) chunks split_text(text) vectors embed_client.embed(chunks) metadatas [ { text: chunk, source: file_path, chunk_index: i, } for i, chunk in enumerate(chunks) ] store.add(vectors, metadatas) print(f已索引 {len(chunks)} 个文本块)这一步完成后向量库就具备了检索能力。7. 检索与提示词组装用户提问时先把问题向量化再从向量库检索相关块def retrieve(query: str, embed_client: EmbeddingClient, store: VectorStore, top_k: int 3) - list[dict]: 检索与问题最相关的文本块。 query_vector embed_client.embed([query])[0] return store.search(query_vector, top_ktop_k)把检索结果组装成带上下文的 Promptdef build_prompt(query: str, retrieved: list[dict]) - str: 把检索结果和用户问题组装成 Prompt。 context \n\n.join( f[来源 {i1}] {item[text]} for i, item in enumerate(retrieved) ) prompt f请根据以下参考资料回答用户问题。如果资料中没有相关信息请明确说明“资料中未找到相关内容”不要编造。 参考资料 {context} 用户问题{query} 请给出准确、简洁的回答并在回答末尾列出引用的来源编号。 return prompt提示词中明确要求模型引用来源编号这是实现可溯源回答的关键。8. 生成回答与溯源输出调用对话模型生成回答并把来源信息一并返回class ChatClient: def __init__(self, api_key: str, base_url: str, model: str gpt-4o-mini): self.client OpenAI(api_keyapi_key, base_urlbase_url) self.model model def generate(self, prompt: str) -gt; str: resp self.client.chat.completions.create( modelself.model, messages[{role: user, content: prompt}], temperature0.2, ) return resp.choices[0].message.content def ask(query: str, embed_client: EmbeddingClient, store: VectorStore, chat_client: ChatClient, top_k: int 3) - dict: 完整的 RAG 问答入口返回回答和溯源信息。 retrieved retrieve(query, embed_client, store, top_k) prompt build_prompt(query, retrieved) answer chat_client.generate(prompt) return { answer: answer, sources: [ { text: item[text], score: item[score], source: item[source], chunk_index: item[chunk_index], } for item in retrieved ], }返回结果中同时包含回答和来源块前端可以据此展示引用来源实现可溯源。9. 完整实战跑通一个问答示例下面把以上模块串起来用一个示例文档跑通全流程。假设有一份产品说明文档product.md# 智能音箱 X1 使用说明 基本功能 智能音箱 X1 支持语音播放音乐、设置闹钟、查询天气和控制智能家居设备。 连接方式 首次使用请下载 App并按照提示连接 Wi-Fi。设备支持 2.4G 和 5G 双频段。 常见问题 如果音箱无法唤醒请检查电源和网络连接并确认麦克风未被静音。主程序如下def main(): api_key your-api-key base_url https://api.openai.com/v1 embed_client EmbeddingClient(api_key, base_url) chat_client ChatClient(api_key, base_url) store VectorStore() 1. 构建索引 build_index(product.md, embed_client, store) 2. 提问 query 音箱无法唤醒怎么办 result ask(query, embed_client, store, chat_client, top_k2) 3. 输出回答 print(回答, result[answer]) print(\n溯源信息) for i, src in enumerate(result[sources], 1): print(f [{i}] 来源: {src[source]} 第{src[chunk_index]}块 相似度: {src[score]:.4f}) print(f 内容: {src[text][:80]}...) if name main: main()运行后系统会先索引文档再根据问题检索相关块最终生成带来源编号的回答。10. 可溯源机制的设计要点可溯源是本文 RAG 的核心目标设计上包含三个层面块级溯源每个文本块都记录来源文件和块序号检索结果天然携带来源。提示词约束在 Prompt 中要求模型引用来源编号从生成侧约束回答可追溯。结果透出问答接口同时返回回答和来源块前端可展示引用卡片。此外相似度分数可以作为溯源可信度的参考分数越高说明该块与问题越相关回答依据越可靠。11. 优化方向与进阶建议手写版本跑通后可以从以下方向继续优化切分策略引入语义切分或按标题结构切分提升块质量。混合检索结合关键词检索BM25和向量检索兼顾精确匹配和语义匹配。重排序在检索后增加重排序模型进一步提升 top_k 结果的相关性。持久化存储把向量库落盘如使用 SQLite 或 FAISS避免每次重启重建索引。多文档支持扩展加载器支持 PDF、Word、网页等格式。这些优化都可以在本文的模块化结构上逐步叠加而无需推翻重写。12. 总结本文从零手写了一个可运行的 RAG 系统覆盖文档加载、文本切分、向量化、向量存储、检索、提示词组装、生成和溯源全流程。核心收获有三点理解 RAG 的完整数据流文档到块、块到向量、向量到检索、检索到生成。掌握可溯源的实现方法块级元数据 提示词约束 结果透出。获得一个可扩展的模块化代码骨架便于后续接入更复杂的优化策略。你可以把这份代码作为基础结合自己的业务文档和模型服务快速搭建一套属于自己的检索增强问答系统。