使用 langChain.js 实现 RAG 知识库语义搜索

5,048 阅读5分钟

大家好,我是双越。前百度 滴滴 资深前端工程师,慕课网金牌讲师,PMP。我的代表作有:

  • wangEditor 开源 web 富文本编辑器,GitHub 18k star,npm 周下载量 20k
  • 划水AI Node 全栈 AIGC 知识库,包括 AI 写作、多人协同编辑。复杂业务,真实上线。
  • 前端面试派 系统专业的面试导航,刷题,写简历,看面试技巧,内推工作。开源免费。

我最近整理了一些 AI Agent 开发相关的资料,有兴趣的同学可以加入分享和讨论~

开始

LangChain 是一个开发 LLM 应用的框架,是目前 LLM 开发最流行的解决方案之一。最早是 Python 开发的,后来也出来 JS 语言的,现在已经更新到 v0.3 版本。

在上篇博客中我使用 langChain.js 开发了一个基础的 AI Agent ,使用自然语言查询天气。

本文将继续使用 langChain.js 实现一个基础的 RAG 语义搜索,这也是开发 AI Agent 时常用的解决方案。

RAG 是什么

RAG - Retrieval Augmented Generation 检索增强生成,一般用于 LLM 整合知识库,模糊搜索非结构化数据,找到相似的结构以后,再交给 LLM 去处理,这样会大大增加搜索结果的准确性。

image.png

创建代码环境

新建一个 nodejs 代码库,安装 langchain 和 dotenv ,后面我们需要使用环境变量。

npm i langchain @langchain/community @langchain/core dotenv

新建一个 rag.js 下文的代码都会在这个文件中写。

加载文档到向量数据库

先加载文档到内存,然后拆分文档内容为小 chunk ,再生成 embed 格式,最后存储到向量数据库。

image.png

加载文档

先准备一个 PDF 文档,不易太短(如几页的简历)也不易太长(如几百页的需求文档),没有的可以下载这个。

把 pdf 放在代码库,然后使用 new PDFLoader 加载到内存中。

import path from 'path'
import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf'

const pdfPath = path.resolve('data/nke-10k-2023.pdf')
const loader = new PDFLoader(pdfPath)

const docs = await loader.load()
// console.log(docs.length)
// console.log(docs[0].metadata)

可以执行 node rag.js 打印一下 docs.length

拆分文档

把当前的文档拆分为更小的 chunk ,size 是 1000 字符,overlap 是相邻 chunk 可以重叠 200 字符。

import { RecursiveCharacterTextSplitter } from '@langchain/textsplitters'

const textSplitter = new RecursiveCharacterTextSplitter({
  chunkSize: 1000,
  chunkOverlap: 200,
})
const allSplits = await textSplitter.splitDocuments(docs)
// console.log('allSplits.length:', allSplits.length)

可以执行 node rag.js 打印 allSplits.length ,它会比 docs.length 更大,因为 chunk 拆分的更小了。

Embeddings

接下来要把刚拆分出来的 allSplit 转换为 embed 向量格式,然后才能存储到向量数据库。

langChain 内置了很多 text_embedding js.langchain.com/docs/integr…

它默认的是 OpenAIEmbeddings 但是我们国内网络不能用,我测试了其他几个推荐的,也都有各种问题,在这卡住了。

最后我找到了 Alibaba Tongyi 这个 embeddings 模型可以使用 js.langchain.com/docs/integr…

首先登录阿里云百炼平台去申请一个 API key ,然后写入 .env 文件中

ALIBABA_API_KEY=xxxx

然后修改 rag.js 代码,把 allSplits 写入到 vectorStore 中

import { AlibabaTongyiEmbeddings } from '@langchain/community/embeddings/alibaba_tongyi'
import { MemoryVectorStore } from 'langchain/vectorstores/memory'
import 'dotenv/config'

const embeddings = new AlibabaTongyiEmbeddings({})
const vectorStore = new MemoryVectorStore(embeddings)
await vectorStore.addDocuments(allSplits)

然后可以写代码测试一下,从向量数据库中查询一个信息

const results1 = await vectorStore.similaritySearch(
  'When was Nike incorporated?'
)
console.log('results1:', results1.length, results1[0])

执行 node rag.js 可以看到打印的结果,这是根据 PDF 内容搜索出来的答案。

image.png

检索和生成

以上是把文档内容转换并存储到 vectorStore 向量数据库了,接下来再加入 llm 进行检索并生成答案。

image.png

加载网页内容到 vectorStore

新建一个文件 rag2.js,这次我们不用本地 PDF 文档了,而用 CheerioWebBaseLoader 去加载一个网页内容。

import 'cheerio'
import { CheerioWebBaseLoader } from '@langchain/community/document_loaders/web/cheerio'

// Load and chunk contents of blog
const pTagSelector = 'p'
const cheerioLoader = new CheerioWebBaseLoader(
  'https://www.wangeditor.com/v5/development.html',
  {
    selector: pTagSelector,
  }
)
const docs = await cheerioLoader.load()
// console.assert(docs.length === 1)
// console.log(`Total characters: ${docs[0].pageContent.length}`)

然后把内容拆分为小的 chunk 和之前一样

import { RecursiveCharacterTextSplitter } from '@langchain/textsplitters'

const splitter = new RecursiveCharacterTextSplitter({
  chunkSize: 1000,
  chunkOverlap: 200,
})
const allSplits = await splitter.splitDocuments(docs)
// console.log('allSplits.length:', allSplits.length)

然后转换为 Embeddings 并存储到向量数据库

import 'dotenv/config'
import { AlibabaTongyiEmbeddings } from '@langchain/community/embeddings/alibaba_tongyi'
import { MemoryVectorStore } from 'langchain/vectorstores/memory'

const embeddings = new AlibabaTongyiEmbeddings({})
const vectorStore = new MemoryVectorStore(embeddings)
await vectorStore.addDocuments(allSplits)

选择一个 LLM 大模型

langChain 集成了有很多 LLM 可供选择 js.langchain.com/docs/integr…

它默认推荐的是 OpenAI 但是在国内我们没法直接调用它的 API ,所以我当前选择的是 DeepSeek 。

注册登录 DeepSeek 创建一个 API key 并把它放在 .env 文件中

DEEPSEEK_API_KEY=xxx

安装 langChain deepseek 插件

npm i @langchain/deepseek

写代码,定义 llm

import { ChatDeepSeek } from '@langchain/deepseek'

const llm = new ChatDeepSeek({
  model: 'deepseek-chat',
  temperature: 0,
  // other params...
})

定义 Agent 工作流

定义数据结构。在 Agent 工作流执行过程中,多个流程之间的数据传输,使用如下结构。

import { Annotation } from '@langchain/langgraph'

const StateAnnotation = Annotation.Root({
  question: Annotation, // 用户输入的问题
  context: Annotation,  // 从 vectorStore 中检索出来的结果
  answer: Annotation,   // 最终输入给用户的结果
})

定义检索方法。通过用户输入的问题 state.question ,在 vectorStore 中检索相似数据 vectorStore.similaritySearch ,最终返回检索结果。

const retrieve = async (state) => {
  console.log('retrieve... question: ', state.question)
  const retrievedDocs = await vectorStore.similaritySearch(state.question)
  return { context: retrievedDocs }
}

定义生成答案的方法。先远程获取一个 promptTemplate (也可以手写),然后结合 retrievedDocs 一起生成 messages ,最后调用 llm 生成自然语言的结果。

import { pull } from 'langchain/hub'

// 从 langChain hub 中获取 promptTemplate
const promptTemplate = await pull('rlm/rag-prompt')
// console.log('promptTemplate ', promptTemplate)

const generate = async (state) => {
  console.log('generate... context: ', state.context.length)
  const docsContent = state.context.map((doc) => doc.pageContent).join('\n')
  const messages = await promptTemplate.invoke({
    question: state.question,
    context: docsContent,
  })
  const response = await llm.invoke(messages)
  return { answer: response.content }
}

定义 workflow 工作流。定义两个节点 retrieve 和 generate ,再定义三个边用于连接这两个节点。

import { StateGraph } from '@langchain/langgraph'

const graph = new StateGraph(StateAnnotation)
  .addNode('retrieve', retrieve)
  .addNode('generate', generate)
  .addEdge('__start__', 'retrieve')
  .addEdge('retrieve', 'generate')
  .addEdge('generate', '__end__')
  .compile()

整体的流程图如下:

image.png

调用 Agent

定义 question 然后使用 invoke 方法调用 Agent

let inputs = { question: '什么是 ModalMenu?' }
const result = await graph.invoke(inputs)
console.log('res ', result.answer)

执行 node rag2.js 可以看到如下结果,先执行 retrieve 再执行 generate 最后返回结果

image.png

还可以使用 stream 流式输出,配合前端能力可以实现打字效果。

const stream = await graph.stream(inputs, { streamMode: 'messages' })
for await (const [message, _metadata] of stream) {
  process.stdout.write(message.content + '|')
}

最后

在实际项目中,向量数据不能存储在内存中,还使用商用向量数据库,例如 Pinecone ,它可以免费试用。关注我,后续将继续分享 Agent MCP 等 AI 开发相关话题。

最后,前端想学全栈 + AI 开发,可以看看我做的 划水AI 项目,AI 写作、多人协同编辑。复杂业务,真实上线,可注册试用。