web_demo.py

#!/usr/bin/env python
# coding=utf-8
import gradio as gr
from openai_utils import get_completion
from prompt_utils import build_prompt
from vectordb_utils import InMemoryVecDB
from pdf_utils import extract_text_from_pdf
from text_utils import split_text

vec_db = InMemoryVecDB()  # 创建一个chromadb数据库对象

# init_db函数用于初始化数据库。它首先从PDF文件中提取文本，然后将文本分割成一系列的段落，最后将这些段落添加到vec_db中。
def init_db(file):
    paragraphs = extract_text_from_pdf(file.name)
    documents = split_text(paragraphs, 500, 100)
    vec_db.add_documents(documents)

# 调用chat函数时，先对用户输入的文本内容在vec_db中进行检索，然后构建一个提示Prompt，最后调用大模型，获得检索回复。
def chat(user_input, chatbot, context, search_field):
    search_results = vec_db.search(user_input, 3)   # 在向量数据库中搜索用户输入的文本
    search_field = "\n\n".join(search_results)  # 将search_results列表中的所有元素连接成一个字符串，每个元素之间用两个换行符("\n\n")分隔.如果search_results是['result1', 'result2', 'result3']，那么"\n\n".join(search_results)的结果就是'result1\n\nresult2\n\nresult3'。这样，每个搜索结果都会在新的一行开始，使得输出更加清晰易读。
    prompt = build_prompt(info = search_results, query = user_input) # 构建一个提示 Prompt
    response = get_completion(prompt, context)  # 调用大模型，获得检索回复
    chatbot.append((user_input, response))      # 将用户输入和检索回复添加到聊天机器人的对话中
    context.append({'role': 'user', 'content': user_input})     # 将用户的问题添加到user角色下 
    context.append({'role': 'assistant', 'content': response})  # 将检索回复添加到assistant角色下
    return "", chatbot, context, search_field   # 返回空字符串，聊天机器人，context和检索结果

# 重置聊天机器人的状态
def reset_state():
    return [], [], "", ""


def main():
    with gr.Blocks() as demo:
        gr.HTML("""<h1 align="center">PDF内容检索器</h1>""")
        gr.Markdown("本demo为您提供了一个简单的界面，用于从PDF文件中提取文本，并使用OpenAI的大模型来检索相关的信息。")

        with gr.Row():
            with gr.Column():
                fileCtrl = gr.File(label="上传文件", file_types=[',pdf'])

        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot()  # 创建一个Gradio聊天机器人对象
            with gr.Column(scale=2):
                # gr.HTML("""<h4>检索结果</h4>""")
                search_field = gr.Textbox(show_label=True, label="检索结果", info="本检索采用的是单一RAG方式", placeholder="空...", lines=10)
                user_input = gr.Textbox(show_label=True, label="用户输入", placeholder="请点击此处输入...", lines=3)
                with gr.Row():
                    submitBtn = gr.Button("提 交", variant="primary", size="lg")
                    emptyBtn = gr.Button("清 空", size="sm", variant="secondary")

        context = gr.State([])

        # 当点击提交按钮时，调用chat函数
        submitBtn.click(chat, [user_input, chatbot, context, search_field],
                        [user_input, chatbot, context, search_field])
        # 当点击清空按钮时，调用reset_state函数
        emptyBtn.click(reset_state, outputs=[chatbot, context, user_input, search_field])

        fileCtrl.upload(init_db, inputs = [fileCtrl])

    demo.queue().launch(share=True, server_name='0.0.0.0', server_port=8080, inbrowser=True)


if __name__ == "__main__":
    main()