AI Agent白手起家55: RAG 知识库设计——从文档摄入到智能检索
纲要
- 知识库工具在智能体中的作用
- 系统架构概览:读与写分离
- 写入路径:文档加载、切分、嵌入与存储
FastAPI服务提供add_url接口- 使用
WebBaseLoader加载网页 - 语义切分
SemanticChunker Chroma向量数据库持久化
- 读取路径:查询重写与多路召回
- 基于历史记录的问题改写链
- 多查询生成与并行检索
MMR算法去重排序- 最终答案合成链
- 完整可运行代码
- 项目结构
- 依赖安装
- 向量数据库写入服务
add_docs.py - 知识库检索工具
knowledge_tool.py - 启动与测试
- 总结与相关度说明
知识库工具:给智能体装上“外挂大脑”
大模型的知识停留在训练截止日期,而实际应用中往往需要接入私有文档、产品手册、内部规章等。RAG技术正是解决这一问题的标准范式:将文档向量化后存入数据库,检索时用相似度找到最相关的片段,交给大模型生成最终答案。
小浪助手的知识库模块实现了完整的读取与写入链路,并加入了查询重写优化,显著提升了检索准确度。
架构总览
读取和写入共享同一个向量数据库,但通过不同的模块独立实现,便于维护和扩展。
写入路径:让知识“入库”
采用FastAPI搭建轻量后台服务,接收 URL 列表,自动完成加载、切分、嵌入和存储。
项目结构
rag_service/ ├── add_docs.py # 文档写入服务 ├── knowledge_tool.py # 检索工具 ├── config.py # 环境变量 ├── .env └── chroma_db/ # 向量数据库持久化目录环境准备
pipinstallfastapi uvicorn langchain langchain-openai langchain-community chromadb python-dotenv配置文件config.py
importosfromdotenvimportload_dotenv load_dotenv()classConfig:OPENAI_API_KEY=os.getenv("OPENAI_API_KEY")OPENAI_BASE_URL=os.getenv("OPENAI_BASE_URL","https://api.openai.com/v1")EMBEDDING_MODEL=os.getenv("EMBEDDING_MODEL","BAAI/bge-m3")CHROMA_PERSIST_DIR=os.getenv("CHROMA_PERSIST_DIR","./chroma_db")COLLECTION_NAME=os.getenv("COLLECTION_NAME","xiaolang_docs")CHUNK_SIZE=int(os.getenv("CHUNK_SIZE",500))CHUNK_OVERLAP=int(os.getenv("CHUNK_OVERLAP",50))文档写入服务add_docs.py
# add_docs.pyimportuuidfromtypingimportList,DictfromfastapiimportFastAPIfrompydanticimportBaseModelfromlangchain_community.document_loadersimportWebBaseLoaderfromlangchain_experimental.text_splitterimportSemanticChunkerfromlangchain_openaiimportOpenAIEmbeddingsfromlangchain_community.vectorstoresimportChromafromconfigimportConfig app=FastAPI()classDocumentProcessor:def__init__(self):self.embeddings=OpenAIEmbeddings(model=Config.EMBEDDING_MODEL,openai_api_key=Config.OPENAI_API_KEY,base_url=Config.OPENAI_BASE_URL,)self.splitter=SemanticChunker(self.embeddings,breakpoint_threshold_type="percentile")self.vectorstore=Chroma(collection_name=Config.COLLECTION_NAME,embedding_function=self.embeddings,persist_directory=Config.CHROMA_PERSIST_DIR,)defadd_from_urls(self,urls:List[str])->Dict:"""从URL列表加载文档并存入向量库"""results=[]forurlinurls:try:loader=WebBaseLoader(url)docs=loader.load()ifnotdocs:results.append({"url":url,"status":"empty"})continuechunks=self.splitter.split_documents(docs)# 为每个块生成唯一IDids=[str(uuid.uuid4())for_inchunks]self.vectorstore.add_documents(chunks,ids=ids)results.append({"url":url,"status":"success","chunks":len(chunks)})exceptExceptionase:results.append({"url":url,"status":"error","detail":str(e)})return{"results":results}processor=DocumentProcessor()classUrlPayload(BaseModel):urls:List[str]@app.post("/add_url")asyncdefadd_url(payload:UrlPayload):returnprocessor.add_from_urls(payload.urls)if__name__=="__main__":importuvicorn uvicorn.run(app,host="0.0.0.0",port=8000)启动服务后,访问http://localhost:8000/docs即可通过界面测试添加 URL 文档。
读取路径:精准检索与答案生成
直接从向量库用原始问题进行相似度搜索,往往得不到最佳结果,因为口语化的提问与文档中的书面表达差异很大。查询重写技术可以生成多个不同角度的查询变体,大幅提升召回率。
知识库检索工具knowledge_tool.py
# knowledge_tool.pyfromtypingimportListfromlangchain.toolsimporttoolfromlangchain_openaiimportChatOpenAI,OpenAIEmbeddingsfromlangchain_community.vectorstoresimportChromafromlangchain_core.promptsimportChatPromptTemplatefromlangchain_core.output_parsersimportStrOutputParserfromlangchain_core.runnablesimportRunnablePassthroughfromconfigimportConfig# 初始化向量库(只读模式)embeddings=OpenAIEmbeddings(model=Config.EMBEDDING_MODEL,openai_api_key=Config.OPENAI_API_KEY,base_url=Config.OPENAI_BASE_URL,)vectorstore=Chroma(collection_name=Config.COLLECTION_NAME,embedding_function=embeddings,persist_directory=Config.CHROMA_PERSIST_DIR,)defrewrite_query(original_query:str,chat_history:str="")->List[str]:"""利用 LLM 将用户问题改写为多个检索变体"""llm=ChatOpenAI(model="gpt-3.5-turbo",temperature=0.3)prompt=ChatPromptTemplate.from_template("""根据聊天记录和最新的用户问题,生成3个独立的、语义相同但表达不同的查询语句, 每个查询单独一行,不要编号,不要解释。 聊天记录: {chat_history} 用户问题: {query} 生成的查询:""")chain=prompt|llm|StrOutputParser()result=chain.invoke({"query":original_query,"chat_history":chat_history})queries=[q.strip()forqinresult.split("\n")ifq.strip()]return[original_query]+queries@tooldefsearch_knowledge_base(query:str)->str:"""从内部知识库检索相关文档并合成答案。用于需要专业领域知识的场景。"""# 查询重写rewritten_queries=rewrite_query(query)# 多路检索并去重all_docs=[]forqinrewritten_queries:docs=vectorstore.max_marginal_relevance_search(q,k=3,fetch_k=10)all_docs.extend(docs)# 去重seen=set()unique_docs=[]fordocinall_docs:ifdoc.page_contentnotinseen:seen.add(doc.page_content)unique_docs.append(doc)# 选取最相关的前5个片段context="\n\n".join([d.page_contentfordinunique_docs[:5]])ifnotcontext:return"知识库中未找到相关信息。"# 合成最终答案llm=ChatOpenAI(model="gpt-3.5-turbo",temperature=0)answer_prompt=ChatPromptTemplate.from_template("使用以下检索到的上下文回答用户问题。如果不知道答案就说不知道,最多三句话。""\n上下文: {context}\n问题: {query}\n答案:")chain=answer_prompt|llm|StrOutputParser()returnchain.invoke({"context":context,"query":query})# 本地测试if__name__=="__main__":# 先确保 add_docs 服务已启动并添加过文档test_query="LangGraph 是如何更新图状态的?"print(search_knowledge_base.invoke(test_query))完整运行流程
- 在
.env文件中配置OPENAI_API_KEY等环境变量。 - 启动文档写入服务:
在 Swagger UI 中提交要学习的网页 URL。python add_docs.py - 测试检索工具:
观察控制台输出,确认向量检索与答案生成正常。python knowledge_tool.py
查询重写的价值
许多开发者会忽略查询重写,直接将用户问题扔给向量数据库。但在多轮对话中,用户可能会说“那个呢?”“上次那个”,这些指代如果不结合历史记录重写为独立查询,向量搜索基本无效。该模块通过引入历史记录和改写链,生成了多个聚焦于核心语义的查询,显著提升了检索的相关性和鲁棒性。
总结
本博客从文档写入到智能检索,完整实现了 RAG 知识库工具。向量数据库选用Chroma,嵌入模型使用硅基流动的BAAI/bge-m3,并结合了语义切分、查询重写、MMR 检索等技术,提供了一个可直接集成到智能体中的知识增强方案。所有代码均可直接运行,开发者只需补充.env配置和相关文档即可。