用纯 Python 实现 RAG：分块、TF-IDF、余弦相似度与提示工程

典型的RAG教程会让你安装五个库，然后调用.from_documents()方法。十分钟后演示就能运行，但你却根本不明白其原理。当检索结果毫无用处或模型产生错误预测时，你就束手无策了，因为所有思考工作都由这些库代劳完成了。

这里的流程则是以相反的方式构建的，完全使用纯Python编写，仅依赖math模块，最后再调用一次普通文本的LLM。没有嵌入模型，没有FAISS，也没有Pinecone。了解这些原理后，当你日后使用LangChain或真正的向量数据库时，就能明白其背后的运作机制。最终不仅能得到一个可运行的流程，还能理解它产生的每一个数值的含义。

RAG解决的问题

语言模型的知识在训练结束后就会固定下来，仅存在于其权重中。若要让它回答有关公司内部文件、代码库或今日新闻的问题，两种显而易见的方案都会令人大失所望：

1. 重新训练模型。 这既耗时又昂贵，而且明天的变更还需要再次进行训练。
2. 将所有内容都放入提示词中。 这种方法无法扩展。即便有百万个标记的上下文窗口，也装不下5万份文档，而且关于长上下文模型的研究不断指出“中间信息丢失”的现象——那些深藏在庞大提示词中的信息虽然存在，但利用率却很低。额外的标记还会增加每次调用的延迟和成本。

RAG的真正优势其实比听起来要有限。文档仍存储在外部存储中，只有在需要回答问题时，才会检索出与该问题相关的少量片段并传递给模型。大语言模型永远不会存储你的数据，它每次接收的都是全新的、相关的内容片段。

因此，任何RAG系统都是由两个部分组成的：**检索模块**，负责在庞大的文档集合中为特定查询找到合适的文本；以及**生成模块**，负责将该文本与查询一起传递给大语言模型以获得答案。接下来的内容将介绍如何用基础组件构建这两个模块。

第一步：分块

将整份50页的PDF文件转换为一个向量会几乎丢失所有有用的信息，因为那些能够回答问题的细节会被平均化而变得毫无意义。首要任务就是将文档切割成小型、目标明确的片段。

最简单的方法是使用按词计数的固定大小滑动窗口，相邻块之间会有重叠：

def chunk_text(text: str, chunk_size: int = 200, overlap: int = 50) -> list[str]:
    if overlap >= chunk_size:
        raise ValueError("overlap must be smaller than chunk_size")

    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunks.append(" ".join(words[start:end]))
        if end >= len(words):
            break
        start = end - overlap  # step back so chunks share context
    return chunks

那个保护性条款并非装饰性的存在。如果overlap的值达到或超过chunk_size，起始位置将停止前进，循环也将永远无法结束。为了“获取更多上下文”而提高overlap的值很容易导致这个问题，因此函数会直接拒绝运行，而不会默默挂起。

重叠的作用比看上去更为重要。没有重叠的话，一个用于回答问题的句子可能会跨越分块边界而被拆分，这样两个分块中都不会包含可以检索的完整信息。而有了重叠之后，边界区域至少会在一个分块中保持完整。

需要了解的一个特殊情况是：固定词长的分词方式会忽略句子和段落的结构，因此可能会将一个句子或表格的一行一分为二。而考虑句子或段落结构的分词方式可以避免这种情况，但会导致分块大小不均——因为一个10词的句子和一个150词的句子都会被归为同一个分块。这会扰乱词元分配情况，以及单个分块在后续检索计算中所占的权重。没有任何一种分词方法是没有代价的，选择哪种方法实际上取决于你更愿意去调试哪种缺陷。

第2步：使用TF-IDF将文本转换为数值

计算机无法直接比较两个字符串的相似度。后续的所有步骤都需要向量，也就是可以进行数学运算的数字列表。构建向量的经典且无需依赖上下文的方法就是TF-IDF，即词频与文档频率倒数的乘积。

• TF用于统计某个词在当前文本片段中出现的频率。如果“kidney”这个词的TF值很高，说明该文本片段很可能与肾脏相关。
• IDF则用于衡量该词在整个语料库中的稀有程度。像“the”和“is”这样的词随处可见，无法说明当前阅读的是哪个文本片段，因此它们的IDF值接近零；而在某个特定文本片段中出现的稀有词则能提供更强的指示意义。

该代码会从所有文本片段中构建词汇表，计算每个词的IDF值，然后将每个文本片段转换为TF-IDF向量：

import math
from collections import Counter

def build_vocabulary(chunks: list[str]) -> list[str]:
    vocab = set()
    for chunk in chunks:
        vocab.update(chunk.lower().split())
    return sorted(vocab)

def compute_idf(chunks: list[str], vocab: list[str]) -> dict[str, float]:
    num_chunks = len(chunks)
    idf = {}
    for word in vocab:
        containing = sum(1 for c in chunks if word in c.lower().split())
        idf[word] = math.log(num_chunks / (1 + containing))
    return idf

def compute_tf(chunk: str, vocab: list[str]) -> dict[str, float]:
    words = chunk.lower().split()
    counts = Counter(words)
    total = len(words)
    return {word: counts.get(word, 0) / total for word in vocab}

def vectorize(chunk: str, vocab: list[str], idf: dict[str, float]) -> list[float]:
    tf = compute_tf(chunk, vocab)
    return [tf[word] * idf[word] for word in vocab]

在三个文本片段上尝试一下，情况就会变得清晰。那些几乎出现在每个片段中的词，比如“the”，其IDF值接近零，因此重复出现也不会对向量产生任何影响。而那些罕见、特定于某个主题的词，则会获得明显的正IDF值。TF-IDF实际上就是一个用来判断某个词是否能为文本片段提供独特信息的公式。

它的缺点在于数据规模。当有10万份文档时，词汇表中可能包含超过50万个唯一单词，这样一来，即便是三个词的文本片段，其向量长度也会达到50万，而且大部分都是零。这类稀疏向量在存储和比较时成本很高。这也是生产系统转而使用维度为384到1536、几乎不含零的密集型神经嵌入的主要原因。

步骤3：用余弦相似度衡量相关性

将文本块视为向量后，接下来的问题是某个文本块与查询的关联程度如何。单纯的距离度量并不可靠，因为对于文本而言，向量之间的角度比长度更为重要。一个只提到一次“猫”的文本块与一个重复提及二十次的文本块其实属于同一主题，而余弦相似度会将它们一视同仁，不会考虑哪个提及的频率更高。

实现这一功能需要点积、向量的模长以及二者之间的比值：

def dot_product(a, b):
    return sum(x * y for x, y in zip(a, b))

def magnitude(v):
    return math.sqrt(sum(x * x for x in v))

def cosine_similarity(a, b):
    denom = magnitude(a) * magnitude(b)
    return 0.0 if denom == 0 else dot_product(a, b) / denom

余弦相似度的取值范围原则上在-1到1之间，采用这种TF-IDF设置后，其数值几乎总是处于0到1的区间内：

• 1.0表示两个向量指向完全相同的方向，即相似度达到最大值。
• 0.0表示两者相互正交，没有共同的显著特征维度，换句话说，查询词与文本片段之间没有任何非平凡的共现词汇。

相似度为0并非异常现象，而是系统正确地表明该文本片段与查询词毫无关联。

需要说明的一点是：经过平滑处理的IDF公式可能会给出现在所有文本片段中的词汇赋予略负的数值，因此余弦相似度偶尔会略低于0。对于实际文本而言这几乎无关紧要，但声称它永远不会为负则有些言过其实了。

TF-IDF结合余弦相似度的缺陷在于无法处理同义词。当查询词为“automobile”而文档内容为“car”时，其相似度得分接近零，因为该算法仅考虑词汇的直接重叠，而不理解语义含义。密集嵌入模型通过将同义词置于相近位置，正是为弥补这一缺陷而存在的。

第4步：检索器

在完成向量化与相似度计算后，检索过程包括：将每个文本块进行一次向量化处理，对输入的查询词也进行向量化，计算每个文本块与查询词的相似度得分，最后对结果进行排序并保留得分最高的k个：

class VectorStore:
    def __init__(self):
        self.chunks, self.vocab, self.idf, self.vectors = [], [], {}, []

    def index(self, documents, chunk_size=50, overlap=10):
        all_chunks = []
        for doc in documents:
            all_chunks.extend(chunk_text(doc, chunk_size, overlap))
        self.chunks = all_chunks
        self.vocab = build_vocabulary(all_chunks)
        self.idf = compute_idf(all_chunks, self.vocab)
        self.vectors = [vectorize(c, self.vocab, self.idf) for c in all_chunks]

    def search(self, query, top_k=3):
        query_vector = vectorize(query, self.vocab, self.idf)
        scored = [(c, cosine_similarity(query_vector, v))
                  for c, v in zip(self.chunks, self.vectors)]
        scored.sort(key=lambda pair: pair[1], reverse=True)
        return scored[:top_k]

一个细节就足以悄悄毁掉一切。查询必须使用从文档中计算出的词汇表和IDF值来进行转换，而仅依据查询本身构建词汇表是错误的。余弦相似度假设向量长度相等，且每个向量中位置对应的词是相同的。如果为查询单独设置词汇表，那么在某个向量中位置5可能是“heart”，而在另一个向量中则是“and”，这样一来相似度得分在结构上就失去了意义，而不仅仅是略有偏差。

第5步：完善提示词

这一步不需要任何数学计算，只需要遵循一定的规则。你需要将检索到的内容片段与用户的问题结合起来，形成一个能够明确告知模型应如何操作的提示词：

def build_augmented_prompt(query: str, retrieved_chunks: list[str]) -> str:
    context = "\n\n".join(
        f"[Source {i+1}]: {chunk}" for i, chunk in enumerate(retrieved_chunks)
    )
    return f"""Answer the question using ONLY the context below.
If the answer is not contained in the context, say "I don't know based on the provided context."
Cite which source number(s) you used.

Context:
{context}

Question: {query}

Answer:"""

这里的两个细节并非仅出于美观考虑。

“仅使用下方的上下文。” 若没有这一明确规则，模型会将检索到的内容片段与训练时记忆的信息混合在一起，此时便无法分辨答案中的哪部分来自可信文档，哪部分是猜测。这种看似有依据的混合结果实际上比完全不使用RAG更具危害性，因为它给人一种有来源的假象，而实际上并不存在。

带编号的来源标签及引用说明。 这些元素能让答案具有可追溯性。当答案出错时——这种情况难免会发生——标签化的来源信息可以显示是哪个内容片段误导了模型，是该片段本身过时或错误，还是模型误解了正确的内容。如果没有这些标签，要找出错误答案的成因几乎只能靠猜测。

第6步：整合所有组件

整个流程首先进行检索，接着进行数据增强，最后调用一次大型语言模型。关键的设计决策是将索引构建与查询操作分为两个独立的流程，而非由同一个函数同时完成这两项任务：

class RagPipeline:
    def __init__(self, client, model="openai/gpt-oss-20b"):
        self.store = VectorStore()
        self.client = client
        self.model = model

    def build_index(self, documents: list[str]):
        """Run once, whenever documents are loaded or updated."""
        self.store.index(documents, chunk_size=50, overlap=10)

    def answer(self, query: str, top_k: int = 3) -> str:
        """Run on every user query — never re-indexes."""
        results = self.store.search(query, top_k=top_k)
        chunks = [c for c, _ in results]
        prompt = build_augmented_prompt(query, chunks)

        response = self.client.chat.completions.create(
            model=self.model,
            max_tokens=500,
            messages=[{"role": "user", "content": prompt}],
        )
        return response.choices[0].message.content

如果将索引构建放在每次查询的函数内部，那么每个查询都会再次对文本进行分割、重新构建词汇表，并为整个集合重新计算所有的TF-IDF向量，尽管这些数据在多次查询之间其实并未发生变化。在10,000份文档和1,000次查询的情况下，这将导致999次重复的构建工作。正确的做法是一次性构建索引后再多次查询，这也是FAISS、Chroma和Pinecone所推荐的方案。这是从一开始就确定的正确架构，而非过早的优化。

该系统的缺陷是什么

这些限制应当被坦诚面对，因为它们并非错误，而是直接源于上述设计选择。

• 没有语义理解能力。“汽车”和“automobile”永远无法匹配，因为TF-IDF仅能识别词汇表层面的字面重叠。

向量数量少但维度极高。 词汇表规模随语料库增长而扩大，因此每个向量也会随之变宽。当文档数量超过几千份时，这种情况就会变得难以处理。

代码中没有任何机制来验证检索结果是否有用。 即使所有排名前k的片段得分都接近零，流程仍会将其传递给模型。提示词要求模型承认自身无知，这完全依赖于大语言模型的配合，因此即便检索效果不佳，内容仍可能未经检查就被用于生成。