前言
之前写了一篇文章，使用国内通义千问作为llm，结合langchain框架实现文本向量化检索和使用chainlit实现网页界面交互，实现一个本地知识问答的机器人。原文链接《使用Chainlit接入通义千问快速实现一个本地文档知识问答机器人》。本次基于上个版本做了增强优化，重要改动是：
- 处理txt文本以外支持pdf文档的知识问答
- 使用流式响应提升用户体验
下面是文章教程
教程
文档问答机器人实现示例
在此示例中，我们将构建一个聊天机器人 QA 应用。我们将学习如何：
- 上传文件
- 从文件创建向量嵌入
- 创建一个聊天机器人应用程序，能够显示用于生成答案的来源
先决条件
安装项目所需依赖。
在项目根目录下创建 requirements.txt 文件，配置需要的依赖内容如下：
chainlit~=1.1.306
openai~=1.37.0
langchain~=0.2.11
chromadb~=0.4.24
tiktoken~=0.7.0
dashscope~=1.20.3
使用命令公爵切换到项目执行以下命令安装：
pip install -r .\requirements.txt
然后，您需要去这里创建一个 OpenAI 密钥。没有可以使用国内的通义千问或者百度文心一言的。具体文章看之前的《使用Chainlit接入通义千问快速实现一个多模态的对话应用》。
使用 LangChain 进行对话式文档 QA
项目根目录下创建文件pdf_qa.py
import chainlit as cl
from chainlit.types import AskFileResponse
from langchain.callbacks.base import AsyncCallbackHandler
from langchain.chains import (
ConversationalRetrievalChain,
)
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import PyPDFLoader, TextLoader
from langchain.embeddings.dashscope import DashScopeEmbeddings
from langchain.memory import ChatMessageHistory, ConversationBufferMemory
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
index_name = "langchain-demo"
# Create a Chroma vector store
embeddings = DashScopeEmbeddings()
author = "Tarzan"
def process_file(file: AskFileResponse):
loader = None
if file.type == "text/plain":
loader = TextLoader(file.path, encoding="utf-8")
elif file.type == "application/pdf":
loader = PyPDFLoader(file.path)
documents = loader.load()
docs = text_splitter.split_documents(documents)
for i, doc in enumerate(docs):
doc.metadata["source"] = f"source_{i}"
return docs
def get_docsearch(file: AskFileResponse):
docs = process_file(file)
# Save data in the user session
cl.user_session.set("docs", docs)
docsearch = Chroma.from_documents(
docs, embeddings, collection_name=index_name
)
return docsearch
@cl.on_chat_start
async def on_chat_start():
files = None
# Wait for the user to upload a file
while files is None:
files = await cl.AskFileMessage(
content="Please upload a text file to begin!",
accept=["text/plain", "application/pdf"],
max_size_mb=20,
timeout=180,
).send()
file = files[0]
msg = cl.Message(content=f"Processing `{file.name}`...")
await msg.send()
docsearch = await cl.make_async(get_docsearch)(file)
message_history = ChatMessageHistory()
memory = ConversationBufferMemory(
memory_key="chat_history",
output_key="answer",
chat_memory=message_history,
return_messages=True,
)
# Create a chain that uses the Chroma vector store
chain = ConversationalRetrievalChain.from_llm(
ChatOpenAI(model_name="qwen-turbo", temperature=0, streaming=True),
chain_type="stuff",
retriever=docsearch.as_retriever(),
memory=memory,
return_source_documents=True,
)
# Let the user know that the system is ready
msg.content = f"Processing `{file.name}` done. You can now ask questions!"
await msg.update()
cl.user_session.set("chain", chain)
class AsyncLangchainCallbackHandler(AsyncCallbackHandler):
def __init__(self, message: cl.Message):
self.message = message
async def on_llm_new_token(self, token: str, **kwargs) -> None:
await self.message.stream_token(token)
@cl.on_message
async def main(message: cl.Message):
msg = cl.Message(content="", elements=[], author=author)
await msg.send()
chain = cl.user_session.get("chain")
# 创建回调处理器实例
cb = AsyncLangchainCallbackHandler(msg)
res = await chain.acall(message.content, callbacks=[cb])
source_documents = res["source_documents"]
text_elements = []
if source_documents:
for source_idx, source_doc in enumerate(source_documents):
source_name = f"source_{source_idx}"
# Create the text element referenced in the message
text_elements.append(
cl.Text(content=source_doc.page_content, name=source_name, display="side")
)
source_names = [text_el.name for text_el in text_elements]
if source_names:
await msg.stream_token(f"\nSources: {', '.join(source_names)}")
msg.elements = text_elements
else:
await msg.stream_token("\nNo sources found")
await msg.update()
代码解释
这段代码是一个使用 Chainlit 库构建的交互式文档问答应用。
导入必要的库和模块
首先，导入了必要的库和模块，例如 Chainlit 中用于处理用户交互的功能，LangChain 用于构建文档问答系统的组件，以及文件加载器、文本分割器等工具。
定义文本分割器和索引名称
text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
index_name = "langchain-demo"
这里定义了一个递归字符文本分割器，用于将长文档拆分成更小的块，以便于处理和索引。同时定义了向量数据库的集合名称。
定义作者变量
author = "Tarzan"
这个变量用于设置消息的发送者名称。
文件处理函数
def process_file(file: AskFileResponse):
...
此函数根据上传文件的类型（纯文本或 PDF）加载文档，并使用前面定义的文本分割器对文档进行分割。然后更新每个文档元数据中的源信息，并返回分割后的文档列表。
向量数据库创建函数
def get_docsearch(file: AskFileResponse):
...
该函数调用 process_file 函数来处理文件，并将处理后的文档存储到 Chroma 向量数据库中。此外，它还会将这些文档保存在用户的会话中。
对话开始时触发的函数
@cl.on_chat_start
async def on_chat_start():
...
当对话开始时，这个函数会被触发。它等待用户上传一个文件，然后调用 get_docsearch 函数来创建向量数据库，并初始化一个对话检索链，准备回答用户的问题。
异步回调处理器类
class AsyncLangchainCallbackHandler(AsyncCallbackHandler):
...
这是一个自定义的回调处理器类，用于处理来自 LangChain 的流式输出。每当模型生成一个新的令牌，它就会调用 on_llm_new_token 方法，将生成的文本流式发送给用户。
消息处理函数
@cl.on_message
async def main(message: cl.Message):
...
当收到用户的消息时，这个函数会被调用。它从用户的会话中获取之前初始化的对话检索链，然后调用这个链来回答问题。同时，它也会将相关的源文档以文本元素的形式展示给用户。
总的来说，这段代码实现了一个简单的问答系统，能够处理用户上传的文档，并针对这些文档回答用户提出的问题。它利用了 LangChain 和 Chainlit 的功能，使得整个交互过程流畅且易于使用。
环境变量
项目根目录下，创建.env文件，配置如下：
OPENAI_BASE_URL="https://dashscope.aliyuncs.com/compatible-mode/v1"
OPENAI_API_KEY="通义千文API-KEY"
DASHSCOPE_API_KEY="通义千文API-KEY"
启动命令
- 命令行工具，项目根目录下执行
chainlit run pdf_qa.py
启动UI示例
- txt文件
- pdf文件