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 = "Your IT assistant"

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

history = ['']

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    btn = gr.Button("Submit", variant="primary")
    clear = gr.Button("Clear")
    
    def user(user_message, history):
        history.append([user_message, None])
        return gr.update(value=""), history
    
    def bot(history):
        try:
            bot_message = palm.chat(
                context=context,
                examples=examples,
                messages=[h[0] for h in history]
            )
            
            history[-1][1] = ""
            for character in bot_message.last:
                history[-1][1] += character
                time.sleep(0.005)
        except Exception as e:
            # Handle the exception here
            print("Error occurred:", str(e))
            # You can customize the error handling as per your requirements
            # For example, return an error message to the user
            
            history[-1][1] = "Incorrect input please retry with a longer sentence in english"
        
        return history

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

demo.queue()
demo.launch(debug=True)