使用Chainlit接入通义千问快速实现一个本地文档知识问答机器人增强版

前言

之前写了一篇文章，使用国内通义千问作为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
1
2
3
4
5
6

使用命令公爵切换到项目执行以下命令安装：

pip install -r .\requirements.txt
1

然后，您需要去这里创建一个 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()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

代码解释

这段代码是一个使用 Chainlit 库构建的交互式文档问答应用。

导入必要的库和模块

首先，导入了必要的库和模块，例如 Chainlit 中用于处理用户交互的功能，LangChain 用于构建文档问答系统的组件，以及文件加载器、文本分割器等工具。

定义文本分割器和索引名称

text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
index_name = "langchain-demo"
1
2

这里定义了一个递归字符文本分割器，用于将长文档拆分成更小的块，以便于处理和索引。同时定义了向量数据库的集合名称。

定义作者变量

author = "Tarzan"
1

这个变量用于设置消息的发送者名称。

文件处理函数

def process_file(file: AskFileResponse):
    ...
1
2

此函数根据上传文件的类型（纯文本或 PDF）加载文档，并使用前面定义的文本分割器对文档进行分割。然后更新每个文档元数据中的源信息，并返回分割后的文档列表。

向量数据库创建函数

def get_docsearch(file: AskFileResponse):
    ...
1
2

该函数调用 process_file 函数来处理文件，并将处理后的文档存储到 Chroma 向量数据库中。此外，它还会将这些文档保存在用户的会话中。

对话开始时触发的函数

@cl.on_chat_start
async def on_chat_start():
    ...
1
2
3

当对话开始时，这个函数会被触发。它等待用户上传一个文件，然后调用 get_docsearch 函数来创建向量数据库，并初始化一个对话检索链，准备回答用户的问题。

异步回调处理器类

class AsyncLangchainCallbackHandler(AsyncCallbackHandler):
    ...
1
2

这是一个自定义的回调处理器类，用于处理来自 LangChain 的流式输出。每当模型生成一个新的令牌，它就会调用 on_llm_new_token 方法，将生成的文本流式发送给用户。

消息处理函数

@cl.on_message
async def main(message: cl.Message):
    ...
1
2
3

当收到用户的消息时，这个函数会被调用。它从用户的会话中获取之前初始化的对话检索链，然后调用这个链来回答问题。同时，它也会将相关的源文档以文本元素的形式展示给用户。

总的来说，这段代码实现了一个简单的问答系统，能够处理用户上传的文档，并针对这些文档回答用户提出的问题。它利用了 LangChain 和 Chainlit 的功能，使得整个交互过程流畅且易于使用。

环境变量

项目根目录下，创建.env文件，配置如下：

OPENAI_BASE_URL="https://dashscope.aliyuncs.com/compatible-mode/v1"
OPENAI_API_KEY="通义千文API-KEY"
DASHSCOPE_API_KEY="通义千文API-KEY"
1
2
3

启动命令

• 命令行工具，项目根目录下执行

chainlit run pdf_qa.py
1

启动UI示例

• txt文件
  
• pdf文件