import gradio as gr
import replicate
import os

key = ""
prompt = ""
personality = ""
def keys(newkey):
  global key
  key = str(newkey)
  print(key)
  return key

def AIPerson(personal):
  global personality
  personality = personal
  return personality
  
def user(user_message, history):
  global prompt
  prompt = user_message
  return "", history + [[user_message, None]]

def bot(history):
  global personality
  global prompt
  global key
  os.environ["REPLICATE_API_TOKEN"] = key
  print(prompt)
  print(personality)
  output = replicate.run(
      "nateraw/goliath-120b:39c0be5d2e96df975da94e8923535733d81179abed5f37c999afc5f7c33bf04e",
      input={
          "top_k": 50,
          "top_p": 1,
          "prompt": prompt,
          "temperature": 0.8,
          "max_new_tokens": 128,
          "prompt_template": f"{personality}\n\nUSER: {prompt}\nASSISTANT:",
          "presence_penalty": 1,
          "frequency_penalty": 0
      },
  )
  history[-1][1] = ''.join([w for w in output]).strip()
  return history

with gr.Blocks(css="style.css") as run:
  apikey = gr.Textbox(label="Enter Replicate API key", elem_classes="apikey")
  chatbot = gr.Chatbot(elem_classes="chatbox")
  msg = gr.Textbox(elem_classes="txtmsg")
  with gr.Row():
    submit = gr.Button("Submit", elem_classes="submit")
    clear = gr.Button("Clear", elem_classes="clear")
  pertxt = gr.Textbox(label="Enter AI Personality", elem_classes="aiperson")

  apikey.input(keys, apikey)
  os.environ["REPLICATE_API_TOKEN"] = key
  pertxt.input(AIPerson, pertxt)
  submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
      bot, chatbot, chatbot
  )
  clear.click(lambda: None, None, chatbot, queue=False)
run.launch(share=True)