☰
用Python标准库手写RAG全链路:从分块、BM25到向量混合检索
2026/10/9 3:41:49 网站建设 项目流程

这个标题我酝酿了很久。起因是我在一套用现成框架搭的知识库里连续被坑:线上问答答非所问,打开框架日志看半天,既看不到文档到底被切成了什么块,也不知道每个问题召回的到底是哪几段,更看不到分数是怎么算出来的。说白了,工具把该暴露的细节全藏起来了。所以我干脆把所有封装全部扔掉,只用Python标准库——re、collections、math、hashlib、json、urllib这些——从分词、分块、倒排索引、BM25打分,到向量相似度、混合检索、上下文拼装和LLM调用,把整条RAG链路从头写了一遍。

这篇就把这条手撕路径完整记录下来。适合两类人:一类是RAG入门者,不想一上来就被框架的黑话劝退;另一类是已经搭过知识库、却总觉得效果"说不清哪里不对"的人。手写一遍之后,RAG对你就不是什么神秘黑盒,而是一张能在脑子里展开的电路图。

1. 为什么放着框架不用,偏要自己造RAG

1.1 先分清RAG知识库和结构化知识库

很多人一上来就被几个概念绕晕,最常见的就是把RAG知识库和知识图谱(KG)混为一谈。先说清楚:RAG适合非结构化文本——人事制度、产品手册、论文、FAQ这些,数据洗干净切块,检索出来拼进提示词让大模型生成答案;而知识图谱是把实体和关系抽出来,存成三元组,适合做多跳推理、关系查询这类需要精确逻辑的问题。两者的应用场景是互补的,不是替代关系。

维度RAG知识库结构化知识图谱
数据形态原始文本切片实体、关系、属性
构建成本低,清洗+切块即可高,需要抽取建模
适合问题全文问答、资料查找多跳关系、精确约束查询
主要瓶颈检索召回质量图构建与维护成本

所以别纠结"我该选RAG还是KG",大多数场景的正确顺序是先用RAG快速落地,等真出现了多跳关系问题,再考虑补图谱。

1.2 RAG链路里的四个核心环节

手撕之前,先把链路拆清楚。一套完整的RAG有四个环节:

  1. 索引(Indexing):把原始文档读进来,清洗、分块、切词,建倒排索引或向量索引;
  2. 检索(Retrieval):用户问题进来,从索引里召回最相关的若干块;
  3. 增强(Augmentation):把召回块整理成语境,塞进提示词模板,控制预算、加引用;
  4. 生成(Generation):调用大模型,基于语境生成带依据的回答。

你会发现,框架做的一堆事情,本质上就是把上面四步封装成API。但问题也出在这里:四步里的每一步都可能出问题,而框架的日志往往只告诉你在第几步失败了,却不告诉你为什么。最典型的是第三步——检索质量差,结果大模型在"没读到正确资料"的情况下,用自己脑子里学过的知识胡编了一个很流畅的答案。这种错误最难定位,因为你从回答表面根本看不出来。

1.3 纯标准库能做到什么程度

很多人问我,不装numpy、不装transformers、不装tiktoken,这套东西能做出来吗?答案是能。我用到的库一只手数得过来:re负责切词和清洗,defaultdict和Counter负责倒排索引,math负责BM25和余弦的计算,hashlib做稳定的哈希映射,json和urllib走LLM接口。效果方面,在几百到几万段落量级的知识库上,自研的"BM25 + 哈希n-gram向量"混合检索,已经足够应付绝大多数内部问答场景。

这套代码全部可以复制运行,除了最后的LLM调用需要你准备一个API Key之外,其余部分连网络都不需要。

2. 索引前的第一道坎:文档分块和词条化

2.1 分块为什么不能一把梭

我先说一个反直觉的结论:检索质量受分块参数的影响,往往大于受检索算法的影响。块太大,一段话里有效信息被无关内容稀释,相似度被平均掉;块太小,上下文断在半句话上,人名和指代关系全丢了。我在公司制度文档里见过最经典的坑:一条"本制度自发布之日起施行"被单独切成一个块,凡是跟"发布日期"相关的问题,召回的几乎全是这块废料。

所以我的做法是三个一起上:按段落切分、按token预算合并、首尾加重叠窗口。重叠窗口的目的是让边界附近的信息至少出现在两个块里,增加被召回的概率——SQL里叫冗余,检索里这叫保险。

2.2 一个不依赖jieba的中英文切词方法

中文没有空格,正规做法是jieba分词,但那已经不是"基础库"了。纯标准库怎么办?我用正则把连续英文和数字整体保留,中文按单字输出,再用bigram把相邻信息粘起来。为什么要补bigram?因为"北京"拆成"北"和"京"两个单字之后,检索"北京"时单字全部命中,分数却会被一堆无关单字词干扰;补一个"北#京"的bigram,检索结果立刻精准一截。

import re from typing import List def tokenize(text: str) -> List[str]: # 英文/数字整体保留,中文按单字,再用 bigram 补充短语信息 latin = re.findall(r'[a-zA-Z0-9]+', text.lower()) chinese = re.findall(r'[\u4e00-\u9fff]', text) tokens = latin + chinese grams = list(tokens) for i in range(len(tokens) - 1): grams.append(tokens[i] + '#' + tokens[i+1]) return grams

这里有个容易忽略的细节:切词要和检索、embedding共用同一个函数。很多项目分词、建索引、查向量各用一套逻辑,结果索引里存的东西和查询时生成的东西对不上,线上召回率莫名其妙掉一截。统一入口,排查问题的时候少一半头疼。

2.3 Token估算与分块代码

正式的token计算需要tiktoken这类依赖模型词表的库,手撕阶段我用一个粗估函数:中文字符数加上英文单词数乘1.3,再留20%余量给各种边界情况。

def estimate_tokens(text: str) -> int: cn = len(re.findall(r'[\u4e00-\u9fff]', text)) en = len(re.findall(r'[a-zA-Z0-9]+', text)) return cn + int(en * 1.3) + 1

分块主逻辑:按空行切段落,逐段累积到预算上限,超了就落地成块,然后用上一块尾部内容充当重叠窗口。

def chunk_document(text: str, chunk_size: int = 400, overlap: int = 80) -> List[str]: paragraphs = [p.strip() for p in re.split(r'\n\s*\n', text) if p.strip()] chunks, current, current_len = [], [], 0 for para in paragraphs: para_len = estimate_tokens(para) if current_len + para_len > chunk_size and current: chunks.append('\n'.join(current)) current = [] current_len = 0 tail = chunks[-1] kept = tail[-overlap * 2:] # 粗略按字符回退,方便演示 if kept: current.append(kept) current_len = estimate_tokens(kept) current.append(para) current_len += para_len if current: chunks.append('\n'.join(current)) return chunks

这个版本胜在逻辑干净,直接能用。线上场景我建议再叠加一层"块元数据":把每个块来自哪个文件、哪个章节、第几页记下来,后面拼prompt的时候一并塞进去,模型回答时能判断资料出处,人也能溯源。

3. 检索内核:从倒排索引到BM25打分

3.1 倒排索引:检索的地基

检索器最原始的问题是"给定词,找出它出现在哪些文档里"。暴力方案是把所有文档扫一遍,文档一多就完蛋。倒排索引的思路正好反过来:预先建一个词到文档列表的映射,查询时直接查表。打个比方,全本书的目录是从前往后翻,倒排索引就是书最后的索引页,按关键词反查页码。

from collections import defaultdict, Counter def build_inverted_index(docs: List[List[str]]): postings = defaultdict(list) # term -> [(doc_id, tf), ...] doc_freq = defaultdict(int) # term -> 包含该词的文档数 for doc_id, tokens in enumerate(docs): freq = Counter(tokens) for term, tf in freq.items(): postings[term].append((doc_id, tf)) doc_freq[term] += 1 return postings, doc_freq

3.2 BM25为什么比朴素词频强

朴素词频的问题很明显:长文档天然词频高,谁长谁排前面,不公平。TF-IDF有了稀有词权重,但依然会被文档长度干扰。BM25在这个基础上做两件事:词频饱和和长度归一化。

公式长这样:

score = IDF(q) * tf * (k1 + 1) / (tf + k1 * (1 - b + b * |D| / avgdl))

其中k1控制词频的饱和程度,一般取1.2到2.0,越大越看重高频词的差异;b控制文档长度的惩罚强度,一般取0.75,越大对长文档惩罚越重。我实际调参的经验是:先固定b=0.75,只动k1,因为k1的直觉最直接——"某个词出现次数翻倍,分数要不要也翻倍"。k1越大,分数增长越不随次数线性放大。

这里有个特别容易踩的坑:BM25的IDF公式有很多变体,如果直接用经典版log((N - df + 0.5) / (df + 0.5)),当某个词出现在超过一半的文档里时,IDF可能变负,整条分数变成负数,排名逻辑直接崩掉。我一开始就中过招。稳妥写法用平滑版本:

import math def bm25_idf(doc_count: int, df: int) -> float: return math.log(1 + (doc_count - df + 0.5) / (df + 0.5))

对数里加个1,下限变成0,从根上避免负分。

3.3 纯Python实现BM25

class BM25: def __init__(self, docs: List[List[str]], k1: float = 1.5, b: float = 0.75): self.k1 = k1 self.b = b self.doc_count = len(docs) self.doc_len = [len(d) for d in docs] self.avg_len = sum(self.doc_len) / self.doc_count self.postings, self.doc_freq = build_inverted_index(docs) def idf(self, term: str) -> float: df = self.doc_freq.get(term, 0) return bm25_idf(self.doc_count, df) def score(self, query_tokens: List[str], doc_id: int) -> float: dl = self.doc_len[doc_id] total = 0.0 for term in set(query_tokens): df = self.doc_freq.get(term, 0) if df == 0: continue tf = 0 for pid, p_tf in self.postings[term]: if pid == doc_id: tf = p_tf break idf = self.idf(term) numerator = tf * (self.k1 + 1) denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_len) total += idf * numerator / denominator return total

有了这个类,检索就是一行sort的事:

def bm25_retrieve(query: str, bm25: BM25, top_k: int = 5): q_tokens = tokenize(query) scored = [(bm25.score(q_tokens, i), i) for i in range(bm25.doc_count)] scored.sort(reverse=True) return scored[:top_k]

到这一步,一个能跑的检索器已经诞生了。但别忘了,BM25只认识词面,不认识意思——下一节解决这个短板。

4. 向量检索与混合打分:不引numpy照样算余弦

4.1 自造语义向量:哈希n-gram的思路

BM25的死穴是同义改写。用户搜"怎么退款",资料里写的是"退货流程",字面不重合,BM25直接漏检。真正的解法是用预训练模型把句子编码成语义向量,但那属于第三方依赖,和本文的"基础库"原则不符。

为了把向量检索的整条管道跑通,我用了一个替代方案:哈希n-gram向量。思路是把文本切成的n-gram逐个用哈希映射到固定维度向量里,然后归一化。它没有神经网络那样的语义深度,但它的产出和真语义向量一模一样——都是一个定长稠密向量,后续的余弦计算、top-k排序、混合打分,全套流程完全一致。这也是我最想强调的一点:管道的正确性,和embedding的质量,是两个可以分开验证的问题。

import hashlib def hashed_embedding(text: str, dim: int = 512) -> List[float]: tokens = tokenize(text) grams = list(tokens) grams += [tokens[i] + '#' + tokens[i+1] for i in range(len(tokens) - 1)] vec = [0.0] * dim for gram in grams: idx = int(hashlib.md5(gram.encode('utf-8')).hexdigest(), 16) % dim vec[idx] += 1.0 norm = math.sqrt(sum(x * x for x in vec)) return [x / norm for x in vec]

一个重要细节:千万别用Python内置的hash()函数来映射。Python的hash()对字符串是进程级加盐的,同一段代码每次重启生成的向量都不一样,索引和查询稍有不一致,整个检索就废了。md5虽然慢,但胜在稳定、跨进程可复现,手撕阶段完全够用。

4.2 余弦相似度与向量归一化

余弦相似度衡量两个向量方向的接近程度,公式是a·b / (|a| * |b|)。因为上面已经把向量归一化过,分母等于1,实际计算退化成点积。为了演示完整性,还是写一个带归一化的完整版本:

def dot(a: List[float], b: List[float]) -> float: return sum(x * y for x, y in zip(a, b)) def cosine(a: List[float], b: List[float]) -> float: norm_a = math.sqrt(sum(x * x for x in a)) norm_b = math.sqrt(sum(x * x for x in b)) if norm_a == 0 or norm_b == 0: return 0.0 return dot(a, b) / (norm_a * norm_b)

性能上要说实话:几千个向量做点积,纯Python循环没问题;到十万级就要考虑numpy或者专门的向量库了。但那是工程优化问题,不影响你把手撕过程跑通。

4.3 BM25和向量分数怎么融合

两个信号单拎出来都有毛病:BM25精准但漏同义改写,哈希向量能捞到近似表达但噪声大。混合检索的思路是把两个分数归一化之后加权求和:

def normalize_scores(scores: List[float]) -> List[float]: lo, hi = min(scores), max(scores) if hi - lo < 1e-9: return [0.0] * len(scores) return [(s - lo) / (hi - lo) for s in scores] def hybrid_retrieve(query: str, chunks: List[str], bm25: BM25, embeddings: List[List[float]], top_k: int = 5, alpha: float = 0.7): q_tokens = tokenize(query) q_vec = hashed_embedding(query) bm_scores = [bm25.score(q_tokens, i) for i in range(len(chunks))] vec_scores = [cosine(q_vec, embeddings[i]) for i in range(len(chunks))] bm_norm = normalize_scores(bm_scores) vec_norm = normalize_scores(vec_scores) combo = [alpha * b + (1 - alpha) * v for b, v in zip(bm_norm, vec_norm)] ranked = sorted(range(len(combo)), key=lambda i: combo[i], reverse=True)[:top_k] return [(i, combo[i]) for i in ranked]

两个通道的分数范围和分布完全不同,不归一化直接加权,弱的通道会被强的吃掉。归一化方法我用的是最简单的min-max,效果已经不错。alpha这个权重在0.6到0.8之间常见,具体多少,留给评测数据说话——这正好是第六节的内容。

5. 把闭环跑起来:检索、上下文拼装与LLM调用

5.1 Top-K召回与上下文预算控制

检索做完,接下来是把召回结果塞进提示词。这里有一个很多人栽过的跟头:以为上下文越长越好。上下文越长,模型注意力越分散,还容易把无关块里的噪声当真。我的默认配置是top_k=3到5,上下文预算控制在2200 token左右——给系统提示词、用户问题和模型回答留出空间。

def build_context(ranked_results, chunks: List[str], budget: int = 2200) -> str: parts, used = [], 0 for rank, (doc_id, score) in enumerate(ranked_results, 1): if used >= budget: break text = chunks[doc_id] remain = budget - used # 演示用粗略截断,正式场景应按 estimate_tokens 循环切到不超预算 truncated = text[:max(int(remain * 2), 0)] parts.append(f"[{rank}] {truncated}") used += estimate_tokens(truncated) return "\n\n".join(parts)

每次截断后必须重新计算消耗的token,因为截断的是字符,消耗的预算按token算。截断丢掉的那部分内容,如果刚好是答案所在,就白白浪费了——所以overlap在这种场景下是双保险。

5.2 提示词模板与引用溯源

上下文拼好之后,模板设计直接决定答案可靠性。我的核心要求是两条:一,只准依据资料回答,资料里没有就明说;二,回答里标注引用编号,方便人工复核。

SYSTEM_PROMPT = ( "你是企业知识库助手。回答只能依据用户提供的参考资料,不得编造。" "如果资料中没有答案,请明确说'资料中未提及'。" "引用资料内容时,在句末用[编号]标注来源。" ) USER_TEMPLATE = "参考资料:\n{context}\n\n请回答:{question}"

别小看这句"资料里没有就明说",它能把模型的幻觉率压下去一大截。我们做过对比,不加这句约束时,模型遇到底层资料缺失,会本能地把训练时学过的相近回答套上来,看着像那么回事,实际全错。

5.3 只用urllib写一个最小LLM客户端

走到生成环节,仍然可以不引requests。urllib是标准库,够用。客户端本质上就是构造一个聊天补全请求,POST出去,解析JSON回来。为了通用,base_url和模型名都做成参数,任何兼容OpenAI接口的服务都能直接换。

import json import os import urllib.request def call_llm(messages: List[dict], model: str = "gpt-4o-mini", base_url: str = "https://api.openai.com/v1/chat/completions") -> str: body = json.dumps({ "model": model, "messages": messages, "temperature": 0.2, }).encode("utf-8") req = urllib.request.Request( base_url, data=body, headers={ "Content-Type": "application/json", "Authorization": "Bearer " + os.environ["LLM_API_KEY"], }, method="POST", ) with urllib.request.urlopen(req, timeout=60) as resp: data = json.loads(resp.read().decode("utf-8")) return data["choices"][0]["message"]["content"]

三个注意事项。一是超时和重试:LLM接口动辄几十秒,timeout设个60秒只是底线,正式的还得加重试;二是把API Key放到环境变量里,千万别写死在代码里,这个习惯我吃过亏之后养成了;三是temperature别调太高,知识库问答要的是稳定准确,不是发散创意,0.2已经够用。

5.4 完整流程串起来

最后把整条链路拼成一个函数:

def rag_pipeline(question: str, chunks: List[str], bm25: BM25, embeddings: List[List[float]], top_k: int = 3) -> str: ranked = hybrid_retrieve(question, chunks, bm25, embeddings, top_k=top_k) context = build_context(ranked, chunks) messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": USER_TEMPLATE.format(context=context, question=question)}, ] return call_llm(messages)

到这一步,自己手写的RAG已经能回答问题了。但先别急着庆祝——跑了几个测试样例说"效果不错",这是最危险的判断方式。

6. 检索质量调优:RAG真正的大瓶颈

6.1 先建一套评测集再谈优化

我在无数个项目里看到同一个问题:团队凭印象评估RAG效果,觉得"好像行",一上线就翻车。原因很简单,人脑对两三个样例的判断极不靠谱,而RAG的失败模式恰恰是"偶尔漏一条关键的"。正确做法是建一个20到50条的评测集,每条包含一个真实问题、一段标准答案、以及它对应的正确文档块编号。

这个问题集合怎么来?让业务人员从知识库里挑最常见的问题,一条条人工标注它应该出自哪个段落。标注过程很枯燥,但它是整条调优路的地基——没有它,你下面所有参数改动都是在猜。

6.2 用Recall@k和MRR把效果量化

评测指标别贪多,两个够用。Recall@k看"正确答案是否出现在前k个召回结果里",MRR看"正确答案排在第几位"。

def recall_at_k(gold_doc_ids: set, retrieved_ids: List[int], k: int) -> float: hit = len(gold_doc_ids & set(retrieved_ids[:k])) return hit / len(gold_doc_ids) def mrr(gold_doc_ids: set, retrieved_ids: List[int]) -> float: for rank, doc_id in enumerate(retrieved_ids, 1): if doc_id in gold_doc_ids: return 1.0 / rank return 0.0

这里我想强调一个几乎所有人都会踩的坑:看评测结果的时候,如果发现回答错误、但检索指标却很高,那么问题大概率出在生成环节;反过来,如果检索指标本身就低,那不管换多聪明的模型都没用。更阴险的是第三种情况——检索指标偏低,回答却看起来挺对。别高兴太早,那往往是模型在用自己的知识圆场,一旦换新领域的问题就会现出原形。人工复核引用编号,就是为了抓住这种假阳性。

6.3 调参方向和常见翻车点

有了评测集,参数调优就变成简单的控制变量实验。我常用的套路是按顺序调四个参数:

  1. chunk_size:试200/400/600三档,观察Recall@3变化;
  2. overlap:从20到50 token之间试,看边界信息召回是否变好;
  3. alpha:从0.5到0.9扫一遍,代码和编号多的语料,BM25权重往上提;用户提问口语化、改写多,向量权重往上提;
  4. top_k:从3到8试,观察MRR是否饱和。

下面这张表是我积累的高频翻车场景:

表现常见原因应对
召回结果和问题毫无关系分块太大,信息被稀释;分词太碎调小chunk_size,检查bigram
同一答案被多个块重复召回overlap过大,内容重复减少overlap,加去重
回答流畅但引用文不对题检索漏召回,模型自圆其说看Recall@k,不是换模型
请求报token超限token估算太乐观预算留20%余量

还有一个容易被忽视的元信息技巧:建索引时把每个块的文件名、章节标题一并记下来,拼进context。模型知道"这段出自《员工手册》第三章",回答时自然更贴合文档语境,引用也更可信。

6.4 什么时候该回到框架

手撕这套东西,目标从来不是"以后永远不用框架"。我的真实体会是:手写一遍,你获得的是调试框架的能力和判断框架好坏的基准。以后再遇到框架检索效果不好,你能直接看出是分块问题、embedding问题还是打分逻辑问题,而不是瞎换配置碰运气。

真到生产环境、数据量到几十万甚至上百万级,我肯定会用专门的向量库、正规的embedding模型,也可能重新引入框架——毕竟成熟的工程去处理并发、持久化、增量更新,比自己手写靠谱。但手撕过的那套代码,会永远留一份在项目里当"参考实现",出了问题拿它做对照,比看任何文档都快。我建议你也这样干:手撕一遍,留作自己的种子工程,后面的路会好走很多。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询