三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

AI Agent白手起家49: 从零构建 ChatDoc 文档对话助手

AI Agent白手起家49: 从零构建 ChatDoc 文档对话助手

纲要

  • 项目概述:ChatDoc 文档检索对话助手
  • 核心功能
    • 加载 PDF、Word 和 Excel 文档
    • 文档切片与国产嵌入模型向量化
    • Chroma 向量存储与智能检索
    • 多种检索调优策略(查询重写、上下文压缩、MMR、相似性打分)
    • 封装为流式对话链
  • 实现步骤
    • 统一文件加载器(根据后缀自动选择)
    • 文本预处理与切片
    • 向量化与存储(BGE‑M3 + Chroma)
    • 基础检索与高级优化
    • 构建对话链并支持流式输出
  • 完整可运行代码:整合加载、切片、嵌入、检索与对话的本地示例

引言

掌握了 RAG 的各个独立组件——文档加载、切片、嵌入、向量数据库和检索器——之后,下一步就是将它们串联成一个完整的应用。本文将以“ChatDoc”项目为例,手把手带你构建一个支持 PDF、Word、Excel 三种格式的文档对话助手。你将看到如何用国产嵌入模型 BGE‑M3 和 Chroma 向量库搭建底层设施,并通过查询重写、上下文压缩等调优手段大幅提升回答质量,最后封装为带流式输出的对话链。

项目架构

ChatDoc 的整体流程遵循经典 RAG 流水线,并在检索环节加入了多种优化策略。

.pdf

.docx

.xlsx

上传文档

文件类型判断

PyPDFLoader

Docx2txtLoader

OpenpyxlLoader

文档切片

BGE-M3 向量化

Chroma 向量库

基础检索器

检索调优

对话链 + 流式输出

实现步骤

环境准备

安装所需依赖:

pipinstalllangchain langchain-core langchain-community chromadb pypdf openpyxl docx2txt

文件加载器:统一入口

首先实现一个get_file_loader函数,根据文件后缀自动选择对应的加载器,将 PDF、Word、Excel 统一转换为 LangChain 的Document对象。

fromlangchain_community.document_loadersimportPyPDFLoaderfromlangchain_community.document_loadersimportDocx2txtLoaderfromlangchain_community.document_loadersimportUnstructuredExcelLoaderdefget_file_loader(file_path:str):iffile_path.endswith(".pdf"):returnPyPDFLoader(file_path)eliffile_path.endswith(".docx"):returnDocx2txtLoader(file_path)eliffile_path.endswith(".xlsx"):returnUnstructuredExcelLoader(file_path)else:raiseValueError(f"不支持的文件格式:{file_path}")

文档切片与向量化

使用RecursiveCharacterTextSplitter按段落切分文档,再通过 BGE‑M3 嵌入模型向量化后存入 Chroma。

fromlangchain_text_splittersimportRecursiveCharacterTextSplitterfromlangchain_community.embeddingsimportHuggingFaceBgeEmbeddingsfromlangchain_community.vectorstoresimportChroma# 切片器splitter=RecursiveCharacterTextSplitter(chunk_size=500,chunk_overlap=50)# 国产嵌入模型 BGE-M3(需联网下载,也可使用 FakeEmbeddings 替代测试)embedding_model=HuggingFaceBgeEmbeddings(model_name="BAAI/bge-m3")defbuild_vectorstore(docs):split_docs=splitter.split_documents(docs)vectorstore=Chroma.from_documents(split_docs,embedding_model,collection_name="chatdoc")returnvectorstore

基础检索与多重查询优化

使用MultiQueryRetriever对用户问题进行改写,生成多个精准子问题后再检索,从而提高召回率。

fromlangchain_community.chat_modelsimportChatOpenAIfromlangchain.retrievers.multi_queryimportMultiQueryRetriever llm=ChatOpenAI(model="gpt-3.5-turbo")# 可替换为 DeepSeek 等国产模型base_retriever=vectorstore.as_retriever(search_kwargs={"k":3})multi_retriever=MultiQueryRetriever.from_llm(retriever=base_retriever,llm=llm)results=multi_retriever.invoke("公司名称是什么?")

上下文压缩

使用LLMChainExtractor对检索结果进行二次压缩,仅保留与问题最相关的核心内容。

fromlangchain.retrieversimportContextualCompressionRetrieverfromlangchain.retrievers.document_compressorsimportLLMChainExtractor compressor=LLMChainExtractor.from_llm(llm)compression_retriever=ContextualCompressionRetriever(base_compressor=compressor,base_retriever=base_retriever)compressed_results=compression_retriever.invoke("公司名称是什么?")

相似性分数过滤与 MMR

通过similarity_search_with_score设定阈值(如 0.5),过滤低分文档。MMR 则可直接通过as_retrieversearch_type参数启用。

# 相似性分数过滤results_with_score=vectorstore.similarity_search_with_score("公司名称",k=5)relevant_docs=[docfordoc,scoreinresults_with_scoreifscore>0.5]# MMR 搜索mmr_retriever=vectorstore.as_retriever(search_type="mmr",search_kwargs={"k":3,"lambda_mult":0.5})mmr_results=mmr_retriever.invoke("公司名称")

封装为对话链

最后,将检索、压缩与 LLM 组合成一条支持流式输出的对话链。

fromlangchain_core.promptsimportChatPromptTemplatefromlangchain_core.runnablesimportRunnablePassthroughfromlangchain_core.output_parsersimportStrOutputParser prompt=ChatPromptTemplate.from_messages([("system","你是一个处理文档的秘书,根据提供的上下文用中文回答问题。"),("human","上下文:{context}\n\n问题:{question}")])chain=({"context":compression_retriever,"question":RunnablePassthrough()}|prompt|llm|StrOutputParser())forchunkinchain.stream("公司注册地址是哪儿?"):print(chunk,end="")

完整可运行代码(使用模拟嵌入)

以下脚本整合了上述核心步骤,使用FakeEmbeddingsFakeListChatModel替代真实的模型,无需下载模型或 API Key 即可运行。请将sample.docx替换为实际文件路径,或直接使用代码中的模拟文档。

fromlangchain_community.document_loadersimportDocx2txtLoaderfromlangchain_text_splittersimportRecursiveCharacterTextSplitterfromlangchain_community.embeddings.fakeimportFakeEmbeddingsfromlangchain_community.vectorstoresimportChromafromlangchain_core.promptsimportChatPromptTemplatefromlangchain_core.output_parsersimportStrOutputParserfromlangchain_community.chat_models.fakeimportFakeListChatModelfromlangchain_core.documentsimportDocument# 1. 模拟一个文档(如果实际文件存在则使用 Docx2txtLoader 加载)# 为方便演示,直接构造一个 Document 对象sample_text=""" 公司名称:宏图科技发展公司 注册地址:江苏省南京市玄武区长江路100号 主营业务:人工智能软件开发、大数据分析 """doc=Document(page_content=sample_text,metadata={"source":"公司信息"})# 2. 切片splitter=RecursiveCharacterTextSplitter(chunk_size=100,chunk_overlap=20)split_docs=splitter.split_documents([doc])# 3. 向量化(模拟嵌入)embeddings=FakeEmbeddings(size=128)vectorstore=Chroma.from_documents(split_docs,embeddings,collection_name="chatdoc")# 4. 基础检索器base_retriever=vectorstore.as_retriever(search_kwargs={"k":2})# 5. 模拟大模型(预设回答)fake_llm=FakeListChatModel(responses=["根据文档内容,公司名称为宏图科技发展公司,注册地址在江苏省南京市玄武区长江路100号。"])# 6. 对话链prompt=ChatPromptTemplate.from_messages([("system","你是一个文档助手,根据提供的上下文回答问题。"),("human","上下文:{context}\n\n问题:{question}")])chain=({"context":base_retriever,"question":lambdax:x}|prompt|fake_llm|StrOutputParser())# 7. 测试question="公司注册地址是哪儿?"print("用户问题:",question)print("助手回答:",chain.invoke(question))

运行该脚本,你将看到助手根据检索到的文档片段给出了正确的注册地址。此模板可以轻松扩展到真实文档、真实嵌入模型和大语言模型,并加入多重查询、上下文压缩等优化策略,构建强大的文档对话应用。

总结

ChatDoc 项目完整展示了从文档加载、切片、向量化,到检索调优和对话封装的全流程。通过组合MultiQueryRetrieverLLMChainExtractor以及相似性分数/MMR 过滤,你可以灵活地提升回答质量。这个骨架稍加扩展,就能演变为你自己的知识库问答系统。

← 返回列表