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))

完整运行流程

1. 在.env文件中配置OPENAI_API_KEY等环境变量。
2. 启动文档写入服务：

   python add_docs.py

   在 Swagger UI 中提交要学习的网页 URL。
3. 测试检索工具：

   python knowledge_tool.py

   观察控制台输出，确认向量检索与答案生成正常。

查询重写的价值

许多开发者会忽略查询重写，直接将用户问题扔给向量数据库。但在多轮对话中，用户可能会说“那个呢？”“上次那个”，这些指代如果不结合历史记录重写为独立查询，向量搜索基本无效。该模块通过引入历史记录和改写链，生成了多个聚焦于核心语义的查询，显著提升了检索的相关性和鲁棒性。

总结

本博客从文档写入到智能检索，完整实现了 RAG 知识库工具。向量数据库选用Chroma，嵌入模型使用硅基流动的BAAI/bge-m3，并结合了语义切分、查询重写、MMR 检索等技术，提供了一个可直接集成到智能体中的知识增强方案。所有代码均可直接运行，开发者只需补充.env配置和相关文档即可。