知识库与 RAG · 让 Agent 记住你的私有资料 ​

// src/build-vectorstore.js —— 跑一次,把文档变成向量存进内存库
import 'dotenv/config'
import { TextLoader } from 'langchain/document_loaders/fs/text'
import { RecursiveCharacterTextSplitter } from '@langchain/text-splitters'
import { OpenAIEmbeddings } from '@langchain/openai'
import { MemoryVectorStore } from 'langchain/vectorstores/memory'
import { writeFile } from 'node:fs/promises'

// ① 读取文档
const loader = new TextLoader('./data/员工手册.txt')
const docs = await loader.load()

// ② 切成小段(按 500 字一段、重叠 80 字)
const splitter = new RecursiveCharacterTextSplitter({
  chunkSize: 500,
  chunkOverlap: 80,
})
const chunks = await splitter.splitDocuments(docs)
console.log(`切成了 ${chunks.length} 段`)

// ③ 每一段转成向量, ④ 存进内存向量库
const embeddings = new OpenAIEmbeddings({
  apiKey: process.env.OPENAI_API_KEY,
  configuration: { baseURL: process.env.OPENAI_API_BASE },
})
const store = await MemoryVectorStore.fromDocuments(chunks, embeddings)

// ⑤ (生产可序列化保存)这里直接把对象导出,供主程序复用
await writeFile('./data/vectorstore.json', JSON.stringify({
  chunks: chunks.map(c => ({ text: c.pageContent, source: c.metadata.source })),
}), 'utf8')
console.log('向量库建立完成!')

// src/search.js —— 独立验证"检索"这一步
import 'dotenv/config'
import { OpenAIEmbeddings } from '@langchain/openai'
import { MemoryVectorStore } from 'langchain/vectorstores/memory'
import { TextLoader } from 'langchain/document_loaders/fs/text'
import { RecursiveCharacterTextSplitter } from '@langchain/text-splitters'

// 重新临时构建(实际项目会复用提前建好的库,这里演示直接逻辑)
const docs = await new TextLoader('./data/员工手册.txt').load()
const chunks = await new RecursiveCharacterTextSplitter({ chunkSize: 500, chunkOverlap: 80 }).splitDocuments(docs)
const embeddings = new OpenAIEmbeddings({ apiKey: process.env.OPENAI_API_KEY, configuration: { baseURL: process.env.OPENAI_API_BASE } })
const store = await MemoryVectorStore.fromDocuments(chunks, embeddings)

// 提问 → 转向量 → 捞最相关的 3 段
const question = '周末加班费怎么算?'
const results = await store.similaritySearch(question, 3) // ← 语义检索:按"意思"找

console.log('命中的片段:')
results.forEach((r, i) => {
  console.log(`\n── 第 ${i + 1} 段 ──`)
  console.log(r.pageContent)
})

// src/index.js —— 升级:带 RAG 的对话 Agent
import 'dotenv/config'
import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai'
import { TextLoader } from 'langchain/document_loaders/fs/text'
import { RecursiveCharacterTextSplitter } from '@langchain/text-splitters'
import { MemoryVectorStore } from 'langchain/vectorstores/memory'
import readline from 'node:readline/promises'

// —— 模型 & 向量库(启动时一次性建好)——
const model = new ChatOpenAI({
  model: process.env.MODEL_NAME ?? 'gpt-4o-mini',
  apiKey: process.env.OPENAI_API_KEY,
  configuration: { baseURL: process.env.OPENAI_API_BASE },
})

const docs = await new TextLoader('./data/员工手册.txt').load()
const chunks = await new RecursiveCharacterTextSplitter({ chunkSize: 500, chunkOverlap: 80 }).splitDocuments(docs)
const embeddings = new OpenAIEmbeddings({ apiKey: process.env.OPENAI_API_KEY, configuration: { baseURL: process.env.OPENAI_API_BASE } })
const store = await MemoryVectorStore.fromDocuments(chunks, embeddings)

// —— RAG 的核心函数:先搜资料,再带资料回答 ——
async function askWithRag(question) {
  const hits = await store.similaritySearch(question, 3)      // ① 捞相关段落
  const context = hits.map(h => h.pageContent).join('\n\n')   // ② 拼成上下文

  // ③ 把"资料 + 问题"一起交给模型,并强制它"照资料答,查不到就说不知道"
  const reply = await model.invoke([
    {
      role: 'system',
      content:
        '你是一个客服助手。回答必须严格依据下面提供的【参考资料】,' +
        '素材里没有的就说不知道,绝不编造。回答末尾附上最相关的来源。\n\n' +
        '【参考资料】\n' + context,
    },
    { role: 'user', content: question },
  ])
  return reply.content
}

// —— 命令行 demo ——
const rl = readline.createInterface({ input: process.stdin, output: process.stdout })
console.log('📚 RAG 客服机器人(输入 bye 退出)')
while (true) {
  const q = await rl.question('你: ')
  if (q === 'bye') break
  const answer = await askWithRag(q)
  console.log('AI: ' + answer + '\n')
}

// Chroma(本地持久化)
import { Chroma } from '@langchain/community/vectorstores/chroma'
const store = await Chroma.fromDocuments(chunks, embeddings, { collectionName: 'company_docs' })

// Qdrant(云/本地)
// const store = await QdrantVectorStore.fromDocuments(chunks, embeddings, { url, collectionName })