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

palm.configure(api_key=os.environ.get("palm_key"))

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

context = "You're a computer failure assistant"

examples = [
    [
        "Hey my computer is broken",
        "Hey, what is the issue with your computer?"
    ]
]

user_message = []
history = []


with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])

    def user(user_message, history):
        user_message.append(msg.value)
        return gr.update(value="", interactive=False), history + [[user_message[-1], None]]

    def bot(history):
        # Initialize the GenerateMessageRequest object
        request = generate_message_pb2.GenerateMessageRequest()

        # Set the prompt field
        request.prompt.context = context
        request.prompt.examples = examples

        # Set the messages field
        request.prompt.messages = history

        # Generate the message
        response = client.generate_message(request)

        return response.message

    response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)

    demo.queue()
    demo.launch()