import openai
import gradio as gr
import os

openai.api_key = os.environ['API_KEY']

def query_gpt(message, chat_history):
  chat_history.append({"role": "user", "content": message})
  out = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=chat_history)
  response = out["choices"][0]["message"]["content"]
  chat_history.append({"role": "assistant", "content": response})
  return response, chat_history

def chatbot_to_chat_history(chatbot):
  chat_history = []
  chat_history.append({"role": "system", "content": "You are a helpful assistant."})
  for message_pair in chatbot:
    chat_history.append({"role": "user", "content": message_pair[0]})
    chat_history.append({"role": "assistant", "content": message_pair[1]})
  return chat_history

with gr.Blocks() as chat:
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Message")
    pwd = gr.Textbox(label="Password")
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, pwd, chatbot):
        if pwd == os.environ['PASSWORD']:
          chat_history = chatbot_to_chat_history(chatbot)
          bot_message,_ = query_gpt(message, chat_history)
          chatbot.append((message, bot_message))
        return "", pwd, chatbot

    msg.submit(respond, [msg, pwd, chatbot], [msg, pwd, chatbot])

chat.launch()