90行Python代码实现的迷你RAG流程：检索、增强、生成

检索增强生成并非任何模型的固有功能，而是一个打包问题。上下文窗口很小，文档却极为庞大，因此必须有人决定哪些内容可以被纳入其中。这个功能仅需约90行纯Python代码即可实现，无需使用LangChain或向量数据库服务。其实现流程分为三步：

1. 检索。将问题进行嵌入，根据相似度对文本片段进行排序，并保留排名靠前的k个片段。
2. 增强。将选中的片段整合到系统消息中。
3. 生成。使用与普通LLM脚本相同的聊天接口来生成内容。

像LangChain这样的框架所提供的所有抽象功能，最终都可以归结为上述三步中的某一步，因此先了解最基础的实现方式有助于理解整个框架的运作逻辑。

嵌入技术将含义转化为几何形态

生成嵌入的模型会将一段文本转换为向量，即一串按特定方式排列的浮点数，使得含义相近的文本彼此靠近。距离是通过余弦相似度来判定的：方向完全相同则为1.0，毫无关联则为0.0。整个原理就在于此。一旦将含义转化为几何形态，查找相关文档就变成了寻找附近的向量。

# pip install openai
# export OPENAI_API_KEY=sk-...
import math

from openai import OpenAI

client = OpenAI()


def embed(texts):
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts,
    )
    return [d.embedding for d in resp.data]


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb)


q, hit, miss = embed([
    "How do I get my money back?",
    "Refunds are possible within 14 days "
    "of purchase.",
    "Data is stored in Frankfurt and "
    "encrypted at rest.",
])

print(f"refund doc:  {cosine(q, hit):.3f}")
print(f"storage doc: {cosine(q, miss):.3f}")

该脚本的典型输出是：退款文档的得分为0.531，存储文档的得分为0.126。看看匹配是如何发生的：问题中提到“money back”，而文档中则写的是“refunds”。两者没有共享任何关键词，但相似度却很高。这就是为什么要使用嵌入方式而非匹配搜索的原因。

向量存储是一种可排序的列表

“向量存储”听起来像是某种基础设施。在较小规模下，它其实就是由（向量，文本）对组成的列表以及排序功能。Pinecone、Chroma和pgvector能够在处理数百万个向量时保持高效，但真正重要的是其接口设计，而这其实可以归为同一类概念。

class VectorStore:
    def __init__(self):
        self.texts = []
        self.vectors = []

    def add(self, texts):
        self.texts.extend(texts)
        self.vectors.extend(embed(texts))

    def top_k(self, query, k=3):
        q = embed([query])[0]
        scored = [
            (cosine(q, v), t)
            for v, t in zip(
                self.vectors, self.texts
            )
        ]
        scored.sort(
            key=lambda s: s[0], reverse=True
        )
        return scored[:k]


store = VectorStore()
store.add([
    "Refunds are possible within 14 days.",
    "The Pro plan costs 12 EUR per month.",
    "Data is stored in Frankfurt.",
])

for score, text in store.top_k(
    "can I cancel and get a refund?", k=2
):
    print(f"{score:.3f}  {text}")

top_k决定了哪些内容会被保留——具体操作为：对所有内容进行排序，保留前k个，舍弃其余的。k的值相当于一个调节参数，用于在召回率、上下文窗口大小以及所需计算量之间取得平衡。

整合检索、增强与生成功能

这里只有一个新函数，即build_messages。它会给选中的内容编号，将其放入系统消息中，并指示模型仅基于这些内容进行回答。用户消息则只包含问题本身。消息列表与普通聊天会话使用的相同，因此RAG仅改变其内容。该文件可直接按原样运行：

# examples/part07_rag.py
# Mini RAG in ~90 lines: embed -> store -> top-k
# -> stitch into messages -> generate.
#
# Setup:
#   pip install openai
#   export OPENAI_API_KEY=sk-...
#
# Run:
#   python part07_rag.py

import math

from openai import OpenAI

client = OpenAI()

EMBED_MODEL = "text-embedding-3-small"
CHAT_MODEL = "gpt-4o-mini"

# A tiny corpus. In real life these are chunks
# of your docs, not whole documents.
DOCS = [
    "Support hours are 9am to 6pm CET, "
    "Monday to Friday.",
    "Refunds are possible within 14 days "
    "of purchase, no questions asked.",
    "The Pro plan costs 12 EUR per month "
    "and includes API access.",
    "The Free plan allows 3 projects and "
    "community support only.",
    "Data is stored in Frankfurt and "
    "encrypted at rest with AES-256.",
    "You can export all your data as "
    "JSON from the settings page.",
]


def embed(texts):
    resp = client.embeddings.create(
        model=EMBED_MODEL,
        input=texts,
    )
    return [d.embedding for d in resp.data]


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb)


class VectorStore:
    """A list of (vector, text) pairs. That is
    all a vector store is, before you need
    scale."""

    def __init__(self):
        self.texts = []
        self.vectors = []

    def add(self, texts):
        self.texts.extend(texts)
        self.vectors.extend(embed(texts))

    def top_k(self, query, k=3):
        q = embed([query])[0]
        scored = [
            (cosine(q, v), t)
            for v, t in zip(
                self.vectors, self.texts
            )
        ]
        scored.sort(key=lambda s: s[0], reverse=True)
        return scored[:k]



def build_messages(question, hits):
    context = "\n\n".join(
        f"[{i + 1}] {text}"
        for i, (_, text) in enumerate(hits)
    )
    system = (
        "Answer using only the context below. "
        "If the answer is not in the context, "
        "say you do not know.\n\n"
        "Context:\n" + context
    )
    return [
        {"role": "system", "content": system},
        {"role": "user", "content": question},
    ]


def ask(store, question):
    hits = store.top_k(question, k=3)
    print(f"Q: {question}")
    for score, text in hits:
        print(f"  {score:.3f}  {text[:44]}...")
    messages = build_messages(question, hits)
    resp = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=messages,
    )
    return resp.choices[0].message.content


if __name__ == "__main__":
    store = VectorStore()
    store.add(DOCS)

    answer = ask(
        store,
        "Where is my data stored, and can "
        "I get it out?",
    )
    print(f"A: {answer}")

示例测试会询问用户数据存储的位置以及是否可以导出。得分最高的三个条目分别是：关于在法兰克福以加密方式存储数据的条目，得分为0.588；关于将所有数据以JSON格式导出的条目，得分为0.472；关于退款政策的条目，得分为0.244。模型回答称数据存储在法兰克福，采用AES-256进行静态加密，并且设置页面支持完整JSON格式的导出。

与存储和导出相关的两个必要条目都进入了候选范围。关于退款的条目虽排名第三，但模型并未考虑它。只要相关的信息能够出现在候选列表中，检索结果即便不完美也是可以接受的。

测试版存在的局限

这90行代码确实存在明显的局限性，每一处都反映出实际系统还需要添加的功能：

• 分块处理。演示版本存储的是完整句子。实际文档需要以有重叠的方式分割，这是预处理阶段的决策，而非模型层面的决定。
• 重新嵌入。每次调用时都对查询进行嵌入是可行的，但每次重启时都重新对整个语料库进行嵌入则不可行，因此需要保存向量。
• 线性扫描。对每个向量进行余弦相似度计算的时间复杂度为O(n)。当分块数量超过约10万时，使用ANN索引才具有实际意义，这也是向量数据库发挥作用的地方。
• 排序质量。单纯的余弦相似度Top-k算法较为粗糙，而使用重排器或混合关键词搜索方法则能提升排序精度。

LangChain将上述各项都封装为可替换的组件。一旦手动编写出基础循环逻辑，这些组件就只是实现特定功能的工具，并无任何神秘之处。