import openai
import os
import argparse
import gradio as gr
from chat_completion import ChatCompletion

api_key = os.environ.get("OPENAI_API_KEY")
openai.api_base="https://apai.zyai.online/v1"

parser = argparse.ArgumentParser()
parser.add_argument('--share', action='store_true', default=False)
parser.add_argument('--welcome', type=str, default='对GPT说点啥...')
parser.add_argument('--title', type=str, default='ChatGPT')
parser.add_argument('--setting', type=str, default=None)
args = parser.parse_args()

bot = ChatCompletion(api_key=api_key)

with gr.Blocks(title=args.title) as demo:
    chatbot = gr.Chatbot(show_label=False)
    msg = gr.TextArea(show_label=False, placeholder=args.welcome)
    send_btn = gr.Button('发送')
    retry_btn = gr.Button('重试')
    reset_btn = gr.Button('清空')

    def send(user_message, history):
        if not user_message:
            return '', history

        response = bot(user_message, setting=args.setting) if user_message != 'retry' else bot.retry()
        return '', history + [[user_message, response]]

    def reset():
        bot.reset()
        return None, [[None, None]]

    def retry(history):
        return send('retry', history)

    send_btn.click(send, inputs=[msg, chatbot], outputs=[msg, chatbot], show_progress=True)
    reset_btn.click(reset, inputs=None, outputs=[msg, chatbot])
    retry_btn.click(retry, inputs=chatbot, outputs=[msg, chatbot])


demo.launch(share=args.share)