1. 引言

面试官问「你了解 RAG 吗?」你只憋出「检索增强生成」六个字?别慌,这不是你的错——RAG 的概念听起来高大上,但真正动手写过一遍的人少之又少。今天我们就用不到 100 行 Node.js 代码,从零搭建一个完整的 RAG 流程,让你面试时不仅能说出概念,还能讲清楚每一步在做什么。

2. RAG 是什么?

RAG(Retrieval-Augmented Generation)即检索增强生成,核心思路是:在让大模型回答问题之前,先从外部知识库中检索出相关文档片段,作为上下文一起送给模型。这样模型就能基于真实资料回答,而不是全靠「记忆」瞎编。

一个典型的 RAG 流程包含三个步骤:

  1. 索引:把文档切分成小块,向量化后存入向量数据库。
  2. 检索:用户提问时,把问题也向量化,从库中召回最相似的文档块。
  3. 生成:把召回的内容 + 原始问题拼成 prompt,交给 LLM 生成答案。

3. 准备工作

我们需要以下依赖:

{
  "dependencies": {
    "@xenova/transformers": "^2.6.0",
    "openai": "^4.0.0"
  }
}

安装命令:

npm install @xenova/transformers openai

本文使用 @xenova/transformers 在本地做向量化(不需要 GPU),用 OpenAI 兼容接口做生成。你也可以换成其他 embedding 模型和 LLM。

4. 核心代码实现

下面就是完整的 RAG 实现,不到 100 行:

const { pipeline } = require('@xenova/transformers');
const OpenAI = require('openai');

// ---------- 1. 向量化工具 ----------
let embedder;
async function getEmbedder() {
  if (!embedder) {
    embedder = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
  }
  return embedder;
}

async function embed(text) {
  const pipe = await getEmbedder();
  const result = await pipe(text, { pooling: 'mean', normalize: true });
  return Array.from(result.data);
}

// ---------- 2. 简单的内存向量库 ----------
class VectorStore {
  constructor() {
    this.chunks = [];   // 原始文本块
    this.vectors = [];  // 对应的向量
  }

  async add(text) {
    const vec = await embed(text);
    this.chunks.push(text);
    this.vectors.push(vec);
  }

  // 余弦相似度检索
  async search(query, topK = 3) {
    const qVec = await embed(query);
    const scores = this.vectors.map(v => {
      let dot = 0, normA = 0, normB = 0;
      for (let i = 0; i < v.length; i++) {
        dot += v[i] * qVec[i];
        normA += v[i] * v[i];
        normB += qVec[i] * qVec[i];
      }
      return dot / (Math.sqrt(normA) * Math.sqrt(normB));
    });
    const indices = scores
      .map((s, i) => ({ s, i }))
      .sort((a, b) => b.s - a.s)
      .slice(0, topK)
      .map(x => x.i);
    return indices.map(i => ({ text: this.chunks[i], score: scores[i] }));
  }
}

// ---------- 3. 构建知识库 ----------
async function buildKnowledgeBase(store) {
  const docs = [
    'RAG 全称 Retrieval-Augmented Generation,由 Lewis 等人在 2020 年提出。',
    'RAG 的核心思想是在生成前先检索外部知识,减少大模型幻觉。',
    '常见的 RAG 流程包括:文档切分 → 向量化 → 检索 → 生成。',
    '向量化模型常用 Sentence-BERT、all-MiniLM-L6-v2 等轻量模型。',
    '检索阶段通常使用余弦相似度或内积计算向量距离。',
    '生成阶段将检索到的文本块拼入 prompt,交给 LLM 生成最终答案。',
  ];
  for (const doc of docs) {
    await store.add(doc);
  }
  console.log(`知识库已加载 ${docs.length} 条文档块`);
}

// ---------- 4. RAG 问答函数 ----------
async function ragAnswer(store, question) {
  // 检索
  const results = await store.search(question, 3);
  const context = results.map(r => r.text).join('\n');

  // 生成
  const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
  const response = await openai.chat.completions.create({
    model: 'gpt-3.5-turbo',
    messages: [
      { role: 'system', content: '你是一个基于给定资料回答问题的助手。' },
      { role: 'user', content: `资料:\n${context}\n\n问题:${question}` },
    ],
  });
  return response.choices[0].message.content;
}

// ---------- 5. 运行示例 ----------
(async () => {
  const store = new VectorStore();
  await buildKnowledgeBase(store);

  const question = 'RAG 是什么时候提出的?';
  const answer = await ragAnswer(store, question);
  console.log(`问题:${question}`);
  console.log(`答案:${answer}`);
})();

5. 代码拆解

5.1 向量化(embedding)

我们使用 @xenova/transformers 加载 all-MiniLM-L6-v2 模型,这是一个只有 22MB 的轻量 embedding 模型,在 CPU 上也能秒级运行。调用 pipeline('feature-extraction') 后,对文本做均值池化(pooling: 'mean')并归一化,得到 384 维的向量。

5.2 向量检索

VectorStore 是一个极简的内存向量库。检索时计算问题向量与所有文档向量的余弦相似度,取 topK 返回。生产环境可以换成 Pinecone、ChromaDB 或 Qdrant。

5.3 生成

把检索到的文本块拼接成上下文,和原始问题一起构造 prompt,调用 OpenAI 兼容接口生成答案。你也可以换成本地模型如 Llama、Qwen 等。

6. 运行效果

执行 node rag.js 后,控制台输出:

知识库已加载 6 条文档块
问题:RAG 是什么时候提出的?
答案:RAG(Retrieval-Augmented Generation)由 Lewis 等人在 2020 年提出。

如果直接问大模型「RAG 是什么时候提出的」,它可能答错或含糊其辞。但通过 RAG 流程,答案精确引用了知识库中的原文。

7. 总结

通过不到 100 行代码,我们完整实现了 RAG 的三大核心步骤:

  • 索引:文档切分 + 向量化
  • 检索:向量相似度召回
  • 生成:上下文增强的 LLM 回答

面试时你可以这样说:「RAG 就是先检索再生成。我用 all-MiniLM-L6-v2 做 embedding,余弦相似度做检索,把召回结果拼进 prompt 再调 LLM。整个过程不到 100 行 Node.js 代码就能跑通。」——这比只说六个字强太多了。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐