import gradio as gr
import time
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.schema import HumanMessage, SystemMessage

# Initialize the chat model
chat = ChatOpenAI(model="gpt-3.5-turbo-16k")

# Define the system message
system_message_content = "You are Michael of Nebadon, also known as Jesus Christ strictly as described in the Urantia Book. You speak in the tone, spirits, and manner of the Christ as described in the Urantia Book."

# Create templates
system_template = system_message_content
system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)
human_template = "{text}"
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

# Build a chat prompt template
chat_prompt = ChatPromptTemplate.from_messages(
    [system_message_prompt, human_message_prompt]
)

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

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

    def bot(history):
        # Get a chat completion from the formatted messages
        messages = chat_prompt.format_prompt(text=history[-1][0]).to_messages()
        response = chat(messages)
        bot_message = response.content  # assuming response is an AIMessage

        history[-1][1] = ""
        for character in bot_message:
            history[-1][1] += character
            time.sleep(0.05)
            yield history

    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()