bin_public/app/Chatbot.py

# -*- coding:utf-8 -*-

from overwrites import *
from chat_func import *
from openai_func import *

from bin_public.utils.utils import *
from bin_public.utils.utils_db import *
from bin_public.config.presets import *
from bin_public.utils.Pinecone import *

my_api_key = ""

# if we are running in Docker
if os.environ.get('dockerrun') == 'yes':
    dockerflag = True
else:
    dockerflag = False

authflag = False

if dockerflag:
    my_api_key = os.environ.get('my_api_key')
    if my_api_key == "empty":
        print("Please give a api key!")
        sys.exit(1)
    # auth
    username = os.environ.get('USERNAME')
    password = os.environ.get('PASSWORD')
    if not (isinstance(username, type(None)) or isinstance(password, type(None))):
        authflag = True
else:
    '''if not my_api_key and os.path.exists("api_key.txt") and os.path.getsize("api_key.txt"):  # API key 所在的文件
        with open("api_key.txt", "r") as f:
            my_api_key = f.read().strip()'''

    if os.path.exists("auth.json"):
        with open("auth.json", "r") as f:
            auth = json.load(f)
            username = auth["username"]
            password = auth["password"]
            if username != "" and password != "":
                authflag = True

gr.Chatbot.postprocess = postprocess
PromptHelper.compact_text_chunks = compact_text_chunks

with open(CSS_path, "r", encoding="utf-8") as f:
    customCSS = f.read()

with gr.Blocks(css=customCSS) as demo:
    history = gr.State([])
    token_count = gr.State([])
    invite_code = gr.State()
    promptTemplates = gr.State(load_template(get_template_names(plain=True)[0], mode=2))
    function_template = gr.State(function)
    interviewer_prompt = gr.Textbox(INTERVIEWER_PROMPT, visible=False)
    user_api_key = gr.State(my_api_key)
    user_question = gr.State("")
    outputing = gr.State(False)
    topic = gr.State("未命名对话历史记录")
    pinecone_api_key = gr.Textbox(PINECONE_API_KEY, visible=False)
    pinecone_api_env = gr.Textbox(PINECONE_API_ENV, visible=False)
    pinecone_index_name = gr.Textbox(PINECONE_INDEX_NAME, visible=False)

    # gr.HTML("""
    # <div style="text-align: center; margin-top: 20px;">
    # """)
    gr.HTML(title)

    with gr.Row(scale=1).style(equal_height=True):
        with gr.Column(scale=5):
            with gr.Row(scale=1):
                chatbot = gr.Chatbot().style(height=600)  # .style(color_map=("#1D51EE", "#585A5B"))
            with gr.Row(scale=1):
                with gr.Column(scale=12):
                    user_input = gr.Textbox(show_label=False, placeholder="在这里输入").style(
                        container=False)
                with gr.Column(min_width=50, scale=1):
                    submitBtn = gr.Button("🚀", variant="primary")
                    cancelBtn = gr.Button("取消", variant="secondary", visible=False)
            with gr.Row(scale=1):
                emptyBtn = gr.Button("🧹 新的对话", )
                retryBtn = gr.Button("🔄 重新生成")
                delFirstBtn = gr.Button("🗑️ 删除最旧对话", visible=False)
                delLastBtn = gr.Button("🗑️ 删除最新对话")
                reduceTokenBtn = gr.Button("♻️ 总结对话")

        with gr.Column():
            with gr.Column(min_width=50, scale=1):
                status_display = gr.Markdown("status: ready")
                with gr.Tab(label="ChatGPT"):
                    keyTXT = gr.Textbox(show_label=True, placeholder=f"OpenAI API-key...",
                                        type="password", visible=not HIDE_MY_KEY, label="API-Key/Invite-Code",)
                    usageTxt = gr.Markdown("**提交key** 以显示额度", elem_id="usage_display")

                    keyTxt = gr.Textbox(visible=False)

                    key_button = gr.Button("Enter")
                    #usage_button = gr.Button("显示用量")

                    model_select_dropdown = gr.Dropdown(label="选择模型", choices=MODELS, multiselect=False,
                                                        value=MODELS[0])
                    with gr.Accordion("参数", open=False):
                        temperature = gr.Slider(minimum=-0, maximum=2.0, value=1.0,
                                                step=0.1, interactive=True, label="Temperature", )
                        top_p = gr.Slider(minimum=-0, maximum=1.0, value=1.0, step=0.05,
                                          interactive=True, label="Top-p (nucleus sampling)", visible=False)
                    use_streaming_checkbox = gr.Checkbox(label="实时传输回答", value=True, visible=enable_streaming_option)
                    use_websearch_checkbox = gr.Checkbox(label="使用在线搜索", value=False)
                    index_files = gr.Files(label="上传索引文件", type="file", multiple=True)

                with gr.Tab(label="Prompt"):
                    systemPromptTxt = gr.Textbox(show_label=True, placeholder=f"在这里输入System Prompt...",
                                                 label="System prompt", value=initial_prompt).style(container=True)
                    with gr.Accordion(label="加载Prompt模板", open=True):
                        with gr.Column():
                            with gr.Row():
                                with gr.Column(scale=6):
                                    templateFileSelectDropdown = gr.Dropdown(label="选择Prompt模板集合文件",
                                                                             choices=get_template_names(plain=True),
                                                                             multiselect=False,
                                                                             value=get_template_names(plain=True)[0])
                                with gr.Column(scale=1):
                                    templateRefreshBtn = gr.Button("🔄 刷新")
                            with gr.Row():
                                with gr.Column():
                                    templateSelectDropdown = gr.Dropdown(label="从Prompt模板中加载", choices=load_template(
                                        get_template_names(plain=True)[0], mode=1), multiselect=False, value=
                                                                         load_template(
                                                                             get_template_names(plain=True)[0], mode=1)[
                                                                             0])

                with gr.Tab(label="保存/加载"):
                    with gr.Accordion(label="保存/加载对话历史记录", open=True):
                        with gr.Column():
                            with gr.Row():
                                with gr.Column(scale=6):
                                    historyFileSelectDropdown = gr.Dropdown(
                                        label="从列表中加载对话",
                                        choices=get_history_names(plain=True),
                                        multiselect=False,
                                        value=get_history_names(plain=True)[0],
                                        visible=False
                                    )

                            with gr.Row():
                                with gr.Column(scale=6):
                                    saveFileName = gr.Textbox(
                                        show_label=True,
                                        placeholder=f"设置文件名: 默认为.json，可选为.md",
                                        label="设置保存文件名",
                                        value="对话历史记录",
                                    ).style(container=True)
                                with gr.Column(scale=1):
                                    saveHistoryBtn = gr.Button("💾 保存对话")
                                    exportMarkdownBtn = gr.Button("📝 导出为Markdown")
                            with gr.Row():
                                with gr.Column():
                                    downloadFile = gr.File(interactive=True)

                with gr.Tab(label="学术功能"):
                    with gr.Row():
                        aep = gr.Button("英语学术润色")
                        acp = gr.Button("中文学术润色")
                        sge = gr.Button("查找英文语法错误")
                        ac2e = gr.Button("学术中译英")
                        c2e = gr.Button("中译英")
                        e2c = gr.Button("英译中")

                with gr.Tab(label="角色"):
                    with gr.Row():
                        interviewer = gr.Button("面试官")
                        migraine = gr.Button("医生问诊")
                        pre_defined_q = gr.Dropdown(label="选择预设问题",
                                                    choices=qs,
                                                    multiselect=False,
                                                    value=qs[0])
                        index_pinecone = gr.Textbox(placeholder=f"Index fetched", visible=True)

                with gr.Tab(label="Davinci-003"):
                    with gr.Column():
                        with gr.Row():
                            with gr.Column():
                                davinci_user_input = gr.Textbox(show_label=False, placeholder="在这里输入").style(
                                    container=False)
                                temperature_davinci = gr.Slider(minimum=-0, maximum=1.0, value=0.7,
                                                step=0.1, interactive=True, label="Temperature", )
                                davinci_submitBtn = gr.Button("🚀", variant="primary")
                            davinci_output = gr.Textbox(show_label=False, placeholder="output").style(
                                    container=False)

    gr.HTML("""
    <div style="text-align: center; margin-top: 20px; margin-bottom: 20px;">
    """)
    gr.Markdown(description)

    get_usage_args = dict(
        fn=get_usage, inputs=[keyTxt], outputs=[usageTxt], show_progress=False
    )

    # 输入为api key则保持不变，为邀请码则调用中心的api key
    key_button.click(key_preprocessing, [keyTXT], [status_display, keyTxt, invite_code]).then(**get_usage_args)

    #usage_button.click(**get_usage_args)

    user_input.submit(predict,
                      [
                          keyTxt,
                          invite_code,
                          systemPromptTxt,
                          history,
                          user_input,
                          chatbot,
                          token_count,
                          top_p,
                          temperature,
                          use_streaming_checkbox,
                          model_select_dropdown,
                          use_websearch_checkbox,
                          index_files
                      ],
                      [chatbot, history, status_display, token_count], show_progress=True)

    user_input.submit(reset_textbox, [], [user_input])

    submitBtn.click(predict, [
        keyTxt,
        invite_code,
        systemPromptTxt,
        history,
        user_input,
        chatbot,
        token_count,
        top_p,
        temperature,
        use_streaming_checkbox,
        model_select_dropdown,
        use_websearch_checkbox,
        index_files],
                    [chatbot, history, status_display, token_count], show_progress=True)

    submitBtn.click(reset_textbox, [], [user_input])

    emptyBtn.click(reset_state, outputs=[chatbot, history, token_count, status_display], show_progress=True)

    retryBtn.click(retry,
                   [keyTxt, invite_code, systemPromptTxt, history, chatbot, token_count, top_p, temperature,
                    use_streaming_checkbox,
                    model_select_dropdown], [chatbot, history, status_display, token_count], show_progress=True)

    delFirstBtn.click(
        delete_first_conversation,
        [history, token_count],
        [history, token_count, status_display],
    )

    delLastBtn.click(
        delete_last_conversation,
        [chatbot, history, token_count],
        [chatbot, history, token_count, status_display], show_progress=True)

    reduceTokenBtn.click(reduce_token_size, [keyTxt, invite_code, systemPromptTxt, history, chatbot, token_count, top_p,
                                             temperature, use_streaming_checkbox, model_select_dropdown],
                         [chatbot, history, status_display, token_count], show_progress=True)
    # History
    saveHistoryBtn.click(
        save_chat_history,
        [saveFileName, systemPromptTxt, history, chatbot],
        downloadFile,
        show_progress=True,
    )
    saveHistoryBtn.click(get_history_names, None, [historyFileSelectDropdown])
    exportMarkdownBtn.click(
        export_markdown,
        [saveFileName, systemPromptTxt, history, chatbot],
        downloadFile,
        show_progress=True,
    )
    # historyRefreshBtn.click(get_history_names, None, [historyFileSelectDropdown])
    historyFileSelectDropdown.change(
        load_chat_history,
        [historyFileSelectDropdown, systemPromptTxt, history, chatbot],
        [saveFileName, systemPromptTxt, history, chatbot],
        show_progress=True,
    )
    downloadFile.change(
        load_chat_history,
        [downloadFile, systemPromptTxt, history, chatbot],
        [saveFileName, systemPromptTxt, history, chatbot],
    )

    # Template
    templateRefreshBtn.click(get_template_names, None, [templateFileSelectDropdown])

    templateFileSelectDropdown.change(load_template,
                                      [templateFileSelectDropdown],
                                      [promptTemplates, templateSelectDropdown],
                                      show_progress=True)

    templateSelectDropdown.change(get_template_content,
                                  [promptTemplates, templateSelectDropdown, systemPromptTxt],
                                  [systemPromptTxt],
                                  show_progress=True)

    # 功能
    function_button_list = [aep, acp, sge, ac2e, c2e, e2c]
    for i in function_button_list:
        name = gr.Dropdown(choices=list(function.keys()), value=i.value, visible=False)
        i.click(get_function_content,
                [function_template, name],
                [systemPromptTxt],
                show_progress=True)

    # 角色
    interviewer.click(predict, [
        keyTxt,
        invite_code,
        interviewer_prompt,
        history,
        user_input,
        chatbot,
        token_count,
        top_p,
        temperature,
        use_streaming_checkbox,
        model_select_dropdown,
        use_websearch_checkbox,
        index_files],
                    [chatbot, history, status_display, token_count], show_progress=True)

    interviewer.click(get_character_content,
                      [interviewer_prompt],
                      [systemPromptTxt])

    migraine.click(context_construction,[
        keyTxt,
        user_input,
        model_select_dropdown,
        pinecone_api_key,
        pinecone_api_env,
        temperature,
        pinecone_index_name],
        [systemPromptTxt, index_pinecone, status_display]
                   ).then(
        predict, [
            keyTxt,
            invite_code,
            systemPromptTxt,
            history,
            user_input,
            chatbot,
            token_count,
            top_p,
            temperature,
            use_streaming_checkbox,
            model_select_dropdown,
            use_websearch_checkbox,
            index_files],
        [chatbot, history, status_display, token_count], show_progress=True
    ).then(reset_textbox, [], [user_input])

    pre_defined_q.change(get_pre_defined_q,
                         [pre_defined_q],
                         [user_input]).then(context_construction,[
        keyTxt,
        user_input,
        model_select_dropdown,
        pinecone_api_key,
        pinecone_api_env,
        temperature,
        pinecone_index_name],
        [systemPromptTxt, index_pinecone, status_display]
                   ).then(
        predict, [
            keyTxt,
            invite_code,
            systemPromptTxt,
            history,
            user_input,
            chatbot,
            token_count,
            top_p,
            temperature,
            use_streaming_checkbox,
            model_select_dropdown,
            use_websearch_checkbox,
            index_files],
        [chatbot, history, status_display, token_count], show_progress=True
    ).then(reset_textbox, [], [user_input])

    # Davinci
    davinci_user_input.submit(predict_davinci,
                      [
                          keyTxt,
                          davinci_user_input,
                          temperature,
                      ],
                      [davinci_output], show_progress=True)

    davinci_submitBtn.click(predict_davinci,
                      [
                          keyTxt,
                          davinci_user_input,
                          temperature_davinci,
                      ],
                      [davinci_output], show_progress=True)

logging.info("\n访问 http://localhost:7860 查看界面")
# 默认开启本地服务器，默认可以直接从IP访问，默认不创建公开分享链接
demo.title = "ChatGPT-长江商学院 🚀"

if __name__ == "__main__":
    reload_javascript()

    # if running in Docker
    if dockerflag:
        if authflag:
            demo.queue().launch(server_name="0.0.0.0", server_port=7860, auth=(username, password))
        else:
            demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=False)
    # if not running in Docker
    else:
        if authflag:
            demo.queue().launch(share=False, auth=(username, password))
        else:
            demo.queue().launch(share=False)  # 改为 share=True 可以创建公开分享链接