gradio_test.py

# -*- coding: utf-8 -*-
import sys

sys.path.append('../')
import gradio as gr
from langchain.prompts import PromptTemplate

from src.llm.chatglm import ChatGLMSerLLM
from src.llm.ernie_with_sdk import ChatERNIESerLLM
from qianfan import ChatCompletion
import os
import asyncio
from src.pgdb.chat.c_db import UPostgresDB
from src.pgdb.knowledge.similarity import VectorStore_FAISS
from src.config.consts import (
    CHAT_DB_USER,
    CHAT_DB_HOST,
    CHAT_DB_PORT,
    CHAT_DB_DBNAME,
    CHAT_DB_PASSWORD,
    EMBEEDING_MODEL_PATH,
    FAISS_STORE_PATH,
    INDEX_NAME,
    VEC_DB_HOST,
    VEC_DB_PASSWORD,
    VEC_DB_PORT,
    VEC_DB_USER,
    VEC_DB_DBNAME,
    SIMILARITY_SHOW_NUMBER,
    GR_PORT,
    GR_SERVER_NAME,
    ICON_PATH
)
from src.server.qa import QA

prompt1 = """'''
{context}
'''
请你根据上述已知资料回答下面的问题，问题如下：
{question}"""

PROMPT1 = PromptTemplate(input_variables=["context", "question"], template=prompt1)


def main():
    c_db = UPostgresDB(host=CHAT_DB_HOST, database=CHAT_DB_DBNAME, user=CHAT_DB_USER, password=CHAT_DB_PASSWORD,
                       port=CHAT_DB_PORT, )
    vecstore_faiss = VectorStore_FAISS(
        embedding_model_name=EMBEEDING_MODEL_PATH,
        store_path=FAISS_STORE_PATH,
        index_name=INDEX_NAME,
        info={"port": VEC_DB_PORT, "host": VEC_DB_HOST, "dbname": VEC_DB_DBNAME, "username": VEC_DB_USER,
              "password": VEC_DB_PASSWORD},
        show_number=SIMILARITY_SHOW_NUMBER,
        reset=False)
    base_llm = ChatERNIESerLLM(
        chat_completion=ChatCompletion(ak="pT7sV1smp4AeDl0LjyZuHBV9", sk="b3N0ibo1IKTLZlSs7weZc8jdR0oHjyMu"))
    # base_llm = ChatGLMSerLLM(url='http://192.168.22.106:8088')

    my_chat = QA(PROMPT1, base_llm, {"temperature": 0.9}, ['context', 'question'], _db=c_db,
                 _faiss_db=vecstore_faiss)

    def clear():  # 清空输入框
        return ''

    def show_history():  # 显示对话历史
        return my_chat.get_history()

    def stop_btn():
        return gr.Button(interactive=False)

    def restart_btn():
        return gr.Button(interactive=True)

    def get_users():
        o_users = my_chat.get_users()
        users_l = [item[0] for item in o_users]
        return gr.components.Radio(choices=users_l, label="选择一个用户", value=users_l[0], interactive=True), users_l[0]

    def create_chat(user_account):
        my_chat.create_chat(user_account)

    def get_chats(user_account):
        o_chats = my_chat.get_chats(user_account)
        chats_l = [item[0] + ':' + item[1] for item in o_chats]
        return gr.components.Radio(choices=chats_l, label="选择一个对话", value=chats_l[0], interactive=True)

    def set_info(question):
        my_chat.set_info(question)

    def set_chat_id(chat_id_info):
        chat_id = chat_id_info.split(':')[0]
        my_chat.set_chat_id(chat_id)

    def load():
        l_users, t_user = get_users()
        l_chats = get_chats(t_user)
        return l_users, l_chats

    with gr.Blocks(css='index(1).css') as demo:
        gr.HTML("""<h1 align="center">低空经济知识问答</h1>""")
        with gr.Row():
            with gr.Column(scale=2):
                users = gr.components.Radio(choices=[], label="选择一个用户", interactive=True,
                                            visible=False, show_label=False)
                chats = gr.components.Radio(choices=[], label="选择一个对话", interactive=True,
                                            show_label=False)
                new_chat_btn = gr.Button("新建对话")
            with gr.Column(scale=8):
                chatbot = gr.Chatbot(bubble_full_width=False,
                                     avatar_images=(ICON_PATH + '\\user2.png', ICON_PATH + "\\bot2.png"),
                                     value=show_history(), height=400, show_copy_button=True,
                                     show_label=False, line_breaks=True)
                with gr.Row():
                    input_text = gr.Textbox(show_label=False, lines=1, label="文本输入", scale=9, container=False)
                    sub_btn = gr.Button("提交", scale=1)

        demo.load(load, [], [users, chats])

        new_chat_btn.click(create_chat, [users], []).then(
            get_chats, [users], [chats]
        )

        users.change(get_chats, [users], [chats]).then(
            set_chat_id, [chats], None
        ).then(
            show_history, None, chatbot
        )

        chats.change(set_chat_id, [chats], None).then(
            show_history, None, chatbot
        )

        sub_btn.click(my_chat.async_chat, [input_text], [chatbot]
                      ).then(
            stop_btn, None, sub_btn
        ).then(
            set_info, [input_text], []
        ).then(
            get_chats, [users], [chats]
        ).then(
            my_chat.update_history, None, None
        ).then(
            show_history, None, chatbot
        ).then(
            clear, None, [input_text]
        ).then(
            restart_btn, None, sub_btn
        )

    demo.queue().launch(share=False, inbrowser=True, server_name='192.168.22.80', server_port=GR_PORT)


if __name__ == "__main__":
    main()