###########################################################################
## @author: Avinash Sampath
## @date: 20/08/2023
## @Description: This file has the API call for the chatbot
###########################################################################

import gradio as gr
import openai
import random
import time

openai.api_key = "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"

system_message = {"role": "system", "content": "You are a helpful assistant."}

with gr.Blocks(title="Ainstein",theme=gr.themes.Default(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.orange), css="footer{display:none !important}") as demo:
    chatbot = gr.Chatbot(label="Ainstein")
    msg = gr.Textbox(placeholder="Type here...", label="User")

    state = gr.State([])

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history, messages_history):
        user_message = history[-1][0]
        bot_message, messages_history = ask_gpt(user_message, messages_history)
        messages_history += [{"role": "assistant", "content": bot_message}]
        history[-1][1] = bot_message
        time.sleep(1)
        return history, messages_history

    def ask_gpt(message, messages_history):
        messages_history += [{"role": "user", "content": message}]
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=messages_history
        )
        return response['choices'][0]['message']['content'], messages_history

    def init_history(messages_history):
        messages_history = []
        messages_history += [system_message]
        return messages_history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, state], [chatbot, state]
    )

demo.title = "Ainstein"
demo.launch(share=True)