LangChain.js做RAG检索(chroma)

6. RAG 检索

RAG（检索增强生成）前面我们已经了解过，经典四个阶段：

加载文档：DocumentLoader 读文件
切分：Splitter 把长文本切成块
向量化 + 存储：Embedding 模型转向量，存入向量库
检索 + 问答：用问题检索相关块，喂给 LLM 回答

阶段 1：文档加载（Load）    阶段 2：切分（Split）
   原始文档                   长文档切成小块
   ├ PDF/Word/txt    →        ├ chunk1
   ├ 网页                     ├ chunk2
   └ 数据库                   └ chunk3...
        │                          │
        ▼                          ▼
阶段 3：向量化+存储（Embed & Store）
   每个 chunk 用 embedding 模型转成向量
   存入向量库（Chroma 等）
        │
        ▼
阶段 4：检索+问答（Retrieve & Generate）—— 运行时
   用户问题 → 向量化 → 在向量库找最相似的 k 个 chunk
   → 把 chunk 拼成 context → 喂给 LLM 回答

阶段 1-3 是离线一次性做的（建库），阶段 4 是每次提问时做的（查询）。所以 RAG 分"建库"和"查询"两个独立流程。

封装 SimpleChroma 包

Chroma 集成放在了 @langchain/community/vectorstores/chroma，但community 已被官方标记 deprecated（停止维护，等拆成独立包），目前独立包并没有，我们按照官方的集成说明(https://docs.langchain.com/oss/javascript/integrations/vectorstores) 进行封装自己的独立包。

SimpleChroma 的核心结构

// 构造：配置外置，清晰
const store = new SimpleChroma({
    host, port, username, password,    // 连接 + Basic Auth
    collectionName,
    embedder,                          // 写入和查询共用同一个实例
})

// ✅ 官方契约方法
await store.addDocuments(docs, { ids })                 // 写入（传 ids 支持增量去重）
const hits = await store.similaritySearchWithScore(q, k) // 检索（带距离）
const docs = await store.similaritySearch(q, k)          // 检索（纯 Document[]）
await store.delete({ ids? / filter? })                   // 删除
const retriever = store.asRetriever(k)                   // 转 Runnable，能接 LCEL 链

// ✅ 生产实用方法
const n = await store.count()             // 查总数
await store.dropCollection()              // 删整个集合（重建用）

源文件：[SimpleChroma.ts]、[bgeEmbedder.ts]。
其它使用时候可能用到的源文件：[splitText.ts]、[loadNovelsFromDir.ts]。会在后面发出来。

关键细节（实现见 library/simpleChroma.ts，有详细注释）：

写入端用 embedder.generate()（不加 BGE 前缀）
检索端用 embedder.generateForQueries()（自动加 BGE 前缀，保证检索质量）
输出是 LangChain 的 Document 类型，能直接接 LCEL 链
metadata 写入前过滤：RecursiveCharacterTextSplitter 会加 loc（对象），但 Chroma 只接受基本类型，所以剔除

完整的

/**
 * RAG检索.ts（使用 simpleChroma.ts 生产级封装）
 * ============================================================
 *
 * 【学习目标】完整 RAG 流程：加载 → 切分 → 向量化存储 → 检索 → 问答
 *
 * 【本示例的技术栈】
 *   - 向量库：library/simpleChroma.ts 的 SimpleChroma（自封装，对齐官方 VectorStore 契约）
 *   - Embedding：library/bgeEmbedder.ts 的 BgeEmbeddingFunction（bge-small-zh-v1.5）
 *   - 切分器：@langchain/textsplitters
 *   - LLM：DeepSeek（生成回答）
 *
 * 【两个流程】
 *   1. buildKnowledgeBase：入库流程（读目录所有 txt → 切分 → 写入向量库）
 *   2. query：检索问答流程（问题 → 检索 top-k → 拼 context → LLM 回答）
 *
 * 删除示例见 06.1.删除示例.ts
 *
 * 【可配置项】
 *   所有可变参数都集中在文件顶部的 CONFIG 对象里，方便调整：
 *   - 集合名、Chroma 连接、embedding 模型、切分参数、检索 top-k
 */

import 'dotenv/config'
import { Document } from '@langchain/core/documents'
import { RecursiveCharacterTextSplitter } from '@langchain/textsplitters'
import { ChatPromptTemplate } from '@langchain/core/prompts'
import { StringOutputParser } from '@langchain/core/output_parsers'
import { ChatDeepSeek } from '@langchain/deepseek'

import { SimpleChroma } from '../library/simpleChroma'
import { BgeEmbeddingFunction } from '../library/bgeEmbedder'
import { loadNovelsFromDir } from '../library/loadNovelsFromDir'

// ============================================================================
// 可配置项（集中在这里，方便调整）
// ============================================================================

const CONFIG = {
    // Chroma 连接配置
    chroma: {
        host: '192.168.31.249',
        port: 7777,
        username: 'admin',
        password: '123456',
    },
    // 集合名（改这里就能换库）
    // ⚠️ 注意：一个集合的向量维度由首次写入时的 embedding 模型决定（bge=512，MiniLM=384）。
    //    换 embedding 模型必须换新集合名（或先 dropCollection），否则维度不匹配会报错。
    collectionName: 'langchain_rag_demo_bge',
    // 要入库的文档目录（目录下所有 .txt 都会被读入）
    documentDir: 'txt',
    // embedding 模型（写入和查询必须一致，换模型要重建库）
    embedModel: 'Xenova/bge-small-zh-v1.5',
    embedDtype: 'fp16' as const,
    // 切分参数
    split: {
        chunkSize: 800,
        chunkOverlap: 100,
    },
    // 检索 top-k
    topK: 3,
} as const

// 共享的 embedding 实例（写入和查询必须用同一个）
const embedder = new BgeEmbeddingFunction({
    modelName: CONFIG.embedModel,
    dtype: CONFIG.embedDtype,
})

// 共享的向量库连接（入库和查询复用）
function createStore() {
    return new SimpleChroma({
        ...CONFIG.chroma,
        collectionName: CONFIG.collectionName,
        embedder,
    })
}

// 读目录下所有 .txt → LangChain Document[]
// 每个文件的 metadata 带 source（文件名）和 book（书名，去扩展名），便于按书过滤检索
function loadDocumentsFromDir(dir: string): Document[] {
    return loadNovelsFromDir(dir).map(n => new Document({
        pageContent: n.text,
        metadata: { source: n.source, book: n.id },
    }))
}

// ============================================================================
// 流程 1：入库
// ============================================================================

/**
 * 入库流程：读目录下所有 .txt → 切分 → 写入向量库（支持增量去重）
 */
async function buildKnowledgeBase() {
    console.log('=== 入库流程 ===\n')

    const store = createStore()

    // 1. 读目录下所有 .txt
    const rawDocs = loadDocumentsFromDir(CONFIG.documentDir)
    console.log(`1. 读取目录 ${CONFIG.documentDir}：${rawDocs.length} 本书`)
    rawDocs.forEach(d => console.log(`     - ${d.metadata.source}（${d.metadata.book}）：${d.pageContent.length} 字符`))

    // 2. 逐本切分，合并所有 chunk；切分后 metadata 会带上 loc（行号），书信息透传
    const splitter = new RecursiveCharacterTextSplitter(CONFIG.split)
    const allChunks: Document[] = []
    for (const doc of rawDocs) {
        const chunks = await splitter.splitDocuments([doc])
        allChunks.push(...chunks)
    }
    console.log(`2. 切分：共 ${allChunks.length} 块`)

    // 3. 写入（id 用 `${book}_chunk_${i}` 格式，多本书不会冲突；增量去重跳过已存在的）
    console.log(`3. 向量化并写入 Chroma（首次加载 embedding 模型稍慢）...`)
    const ids = allChunks.map((c, i) => `${c.metadata.book}_chunk_${i}`)
    const written = await store.addDocuments(allChunks, { ids })
    console.log(`   实际写入 ${written.length} 块（共 ${allChunks.length} 块，增量去重跳过 ${allChunks.length - written.length}）`)
    console.log(`   集合当前共 ${await store.count()} 条记录\n`)
}

// ============================================================================
// 流程 2：检索问答
// ============================================================================

/**
 * 检索问答流程：问题 → 检索 top-k → 拼 context → LLM 回答
 */
async function query(question: string) {
    console.log(`=== 检索问答 ===`)
    console.log(`问题: ${question}\n`)

    const store = createStore()

    // 1. 检索相关片段（带距离）
    const hits = await store.similaritySearchWithScore(question, CONFIG.topK)
    console.log(`检索到 ${hits.length} 条相关片段（距离越小越相关）:\n`)
    hits.forEach(([doc, distance], i) => {
        const preview = doc.pageContent.slice(0, 60).replace(/\n/g, ' ')
        console.log(`  [${i + 1}] 距离 ${distance.toFixed(4)}  "${preview}..."`)
    })
    console.log()

    // 2. 拼 context，用 LCEL 链做问答
    const context = hits.map(([doc]) => doc.pageContent).join('\n\n---\n\n')
    const model = new ChatDeepSeek({ model: 'deepseek-v4-flash', apiKey: process.env.DEEPSEEK_API_KEY })
    const ragPrompt = ChatPromptTemplate.fromTemplate(`
你是一个知识库助手。请根据以下参考资料回答问题。
如果资料不足，请明确说明。

参考资料：
{context}

问题：{question}
`.trim())

    const ragChain = ragPrompt.pipe(model).pipe(new StringOutputParser())
    const answer = await ragChain.invoke({ context, question })
    console.log('回答:')
    console.log(answer)
}

// ============================================================================
// 入口：入库 → 查询
// ============================================================================

async function main() {
    // 入库
    // await buildKnowledgeBase()
    // 查询
    await query('政府是怎么从伊文思手中拿到三体人相关的资料的？')
}

main().catch(err => {
    console.error('运行失败:', err)
    process.exit(1)
})

6. RAG 检索

封装 SimpleChroma 包

SimpleChroma 的核心结构

完整的

💬 评论 0