import os
import time
import gradio as gr
import google.generativeai as genai

# Credentials
genai.configure(api_key=os.getenv('PALM_API_KEY'))

# Gradio
title = '🦒Playground w/ Google PaLM v2'
description = """Click below tab and start with your message"""

def generate_text(prompt: str):
    response = genai.generate_text(prompt=prompt)
    return response.result

chat_defaults = {
  'model': 'models/chat-bison-001',
  'temperature': 0.25,
  'candidate_count': 1,
  'top_k': 40,
  'top_p': 0,
}

chat_messages = []

def generate_chat(prompt: str):
    context = "You are an intelligent chatbot powered by biggest technology company."
    chat_messages.append(prompt)
    response = genai.chat(
      **chat_defaults,
      context=context,
      messages=chat_messages
    )
    chat_messages.append(response.last)
    return response.last

with gr.Blocks(theme='JohnSmith9982/small_and_pretty') as demo:
    
    gr.Markdown(
    f"""
    # {title}
    ## {description}
    """)

    with gr.Tab('Just Text'):
        app = gr.Interface(fn=generate_text, inputs=["text"], outputs=["text"],
                           examples=[["What is Large Language Model?"], ["Write me a story about Pokemon."]])

    with gr.Tab('Just Chat'):
        chatbot = gr.Chatbot(height=600)
        msg = gr.Textbox()
        clear = gr.Button("Clear")
    
        def user(user_message, history):
            return "", history + [[user_message, None]]
    
        def bot(history):
            bot_message = generate_chat(history[-1][0])
            history[-1][1] = ""
            for character in bot_message:
                history[-1][1] += character
                time.sleep(0.01)
                yield history
    
        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
            bot, chatbot, chatbot
        )
        clear.click(lambda: None, None, chatbot, queue=False)
    
    gr.Markdown(
    f"""
    Testing by __G__
    """)
    
demo.queue()
demo.launch()