100行代码!实现一个基于私有知识库的RAG AI Agent

用户问题 → 向量化检索 → 找到相关文档 → 喂给LLM → 准确回答

项目初始化

# 创建项目mkdir rag-agent && cd rag-agent

修改package.json

{  "name": "rag-agent",  "version": "1.0.0",  "description": "",  "main": "index.js",  "type": "module",  "scripts": {    "test": "echo \"Error: no test specified\" && exit 1"  },  "keywords": [],  "author": "",  "license": "ISC",  "packageManager": "pnpm@10.24.0",  "dependencies": {    "@langchain/community": "^1.1.9",    "@langchain/core": "^1.1.17",    "@langchain/langgraph": "^1.1.2",    "@langchain/openai": "^1.2.3",    "@langchain/textsplitters": "^1.0.1",    "chromadb": "^3.2.2",    "jsonwebtoken": "^9.0.3",    "langchain": "^1.2.15",    "pdf-parse": "^1.1.4",    "zod": "^4.3.6"  },  "devDependencies": {    "tsx": "^4.7.0",    "typescript": "^5.1.6"  }}

安装依赖：

npm install

相关代码

第一步:配置向量模型和LLM

创建src/model.ts:

import { ZhipuAIEmbeddings } from"@langchain/community/embeddings/zhipuai";import { ChatOpenAI } from'@langchain/openai'import { Chroma } from"@langchain/community/vectorstores/chroma";// 知识库向量模型,使用智谱的embedding-3模型exportconst embeddings = newZhipuAIEmbeddings({modelName: "embedding-3",apiKey: 'ZHIPU_API_KEY'// 替换你的智普API Key});// LLM配置,使用智谱GLM-4.7-Flashexportconst model = newChatOpenAI({modelName: "glm-4.7-flash",maxRetries: 0,timeout: 60000,configuration: {baseURL: "https://open.bigmodel.cn/api/paas/v4/"    },apiKey: 'ZHIPU_API_KEY',// 替换你的智普API Keytemperature: 0});// ChromaDB向量存储exportconst vectorStore = newChroma(embeddings, {collectionName: "company-docs",host: "localhost",port: 8000});

关键点解析:

• embedding-3:将文本转换成向量的模型(768维数字数组)
• Chroma:专门存储和检索向量的数据库,支持相似度搜索
• temperature: 0:让AI回答更确定,减少随机性

第二步:知识库入库

创建src/rag/doc-init.ts:

import { PDFLoader } from"@langchain/community/document_loaders/fs/pdf"import { RecursiveCharacterTextSplitter } from"@langchain/textsplitters";import { vectorStore } from'../model.js'// 加载PDF知识库文档const loader = newPDFLoader("./幸福的婚姻.pdf")const docs = await loader.load()// 文档切分器const splitter = newRecursiveCharacterTextSplitter({chunkSize: 1000,      // 每块1000字符chunkOverlap: 200,    // 块之间重叠200字符(保证上下文连贯)});const allSplits = await splitter.splitDocuments(docs);console.log(`文档切分成${allSplits.length}块`)// 数据清洗(ChromaDB要求metadata必须是简单类型)allSplits.forEach(doc => {if (doc.metadata.pdf) {delete doc.metadata.pdf;    }for (const key in doc.metadata) {const value = doc.metadata[key];if (typeof value === 'object' && value !== null) {            doc.metadata[key] = JSON.stringify(value);        }    }});// 存入向量数据库await vectorStore.addDocuments(allSplits);console.log('✅ 文档保存到向量库成功')

为什么需要文档切分?

• LLM有上下文窗口限制(比如8K tokens),无法一次读完整个PDF
• 切分成小块后,检索更精准(只找相关段落,而不是整本书)
• overlap保证语义完整性(避免句子被截断)

第三步:实现RAG Agent

创建src/rag/rag-agent.ts:

import { createAgent, dynamicSystemPromptMiddleware } from"langchain";import { SystemMessage } from"langchain";import { vectorStore, model } from'../model.js'const agent = createAgent({    model,tools: [],middleware: [// 中间件:动态修改系统提示词dynamicSystemPromptMiddleware(async (state) => {// 1. 获取用户问题const lastMessage = state.messages[state.messages.length - 1];const lastQuery = typeof lastMessage.content === 'string'                ? lastMessage.content : "";// 2. 从向量库检索相似度最高的2条相关文档const retrievedDocs = await vectorStore.similaritySearch(lastQuery, 2);// 3. 提取文档内容const docsContent = retrievedDocs                .map((doc) => doc.pageContent)                .join("\n\n");// 4. 将查到的知识库内容构造系统提示词(强制AI只基于文档回答)const systemMessage = newSystemMessage(`你是一个专业的客服助手。请只基于以下上下文回答问题,不要编造信息:\n\n${docsContent}`            );console.log('📚 检索到的文档:', docsContent.slice(0, 100) + '...')return systemMessage;        })    ]});// 测试提问let response = await agent.invoke(    { messages: [{ role: "user", content: "怎样培养你的喜爱和赞美" }] });console.log('\n✅ AI 回答：', response.messages[1].content)response = await agent.invoke(    { messages: [{ role: "user", content: "如何改掉冷战的毛病" }] });console.log('\n✅ AI 回答：', response.messages[1].content)

核心原理拆解:

1. 拦截用户问题:通过middleware在LLM调用前"劫持"对话
2. 向量检索:用embedding模型把问题转向量,在ChromaDB中找最相似的文档
3. 注入上下文:把检索到的文档塞进System Prompt
4. 强制约束:System Prompt中明确要求"只基于上下文回答",避免幻觉

完整运行步骤

1. 启动ChromaDB服务

# 安装ChromaDB(如果还没装)npm install -g chromadb# 启动服务chroma run

启动成功后,默认运行在http://localhost:8000

2. 准备知识库文档

把你的PDF文档放到项目根目录(示例用的是《幸福的婚姻.pdf》)

3. 知识库入库

tsx src/rag/doc-init.ts

输出:

✅ 文档保存到向量库成功

4. 运行RAG Agent

tsx src/rag/rag-agent.ts

输出示例:

📚 检索到的文档: 培养喜爱和赞美的方法之一是...✅ AI 回答: 根据文档,培养喜爱和赞美的方法包括...📚 检索到的文档: 冷战会破坏婚姻关系...✅ AI 回答: 文档建议通过以下方式改掉冷战毛病...

进阶优化方向

1. 提升检索质量

• 混合检索:向量检索+关键词检索(BM25)
• 重排序(Reranking):用专门的模型对检索结果重新排序
• 查询扩展:用LLM把用户问题重写成多个查询,扩大召回范围

2. 多模态支持

• 支持图片、视频、音频等多媒体文档
• 用CLIP等模型处理图片向量

3. 分布式部署

• 生产环境推荐用Milvus、Qdrant等专业向量数据库
• 用Redis做缓存,加速常见问题查询

4. 监控与评估

• 记录每次检索的文档,人工评估相关性
• 用RAGAS等工具自动评测RAG系统质量

总结

通过RAG技术,我们只用不到100行代码就实现了:

✅ 让AI拥有企业私有知识库

 ✅ 大幅提升回答准确性(减少幻觉) 

✅ 支持知识库实时更新 

✅ 本地部署,数据安全可控

下一步行动:

1. 把你公司的文档整理好,替换示例PDF
2. 部署到服务器,包装成API供团队使用
3. 接入企业微信/钉钉/Slack,做成智能客服

关注我,后台回复"RAG"领取完整项目代码!

关于作者:静远AI实战,专注AI落地与效率提升。帮助企业和个人用AI提升10倍效率。

我的服务:AI自动化实战 | 企业AI咨询 | AI应用开发 | 知识付费课程