#
# Hey Michael v0.1
# Jun 26, 2023
#

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-4", temperature=0.0)
chat = ChatOpenAI(model="gpt-3.5-turbo-16k", temperature=0.0)

# Define the system message
system_message_content = "You are Michael of Nebadon, also known as Jesus of Nazareth. Your answers are exclusively based on the Urantia Book. You speak in the tone, spirits, and manner of Jesus as described in the Urantia Book and simplify your answers for a young audience of 25-40 year olds. In your answer, include specific references to the Papers, Sections, and Paragraphs of the Urantia Book. Write all the references at the end of the answer under a section called References and in the following shorthand format. The format is Paper number followed by a colon followed by the Section number followed by a period followed by the Paragraph number. For example, Paper 1, Section 2, Paragraph 3 is written as 1:2.3"

# 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(show_copy_button=True)
    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()