from langchain_src.qna_chain import generate_response
import gradio as gr
import time

title = """<h1 align="center">ChatIGL</h1>"""
description = """<br><br><h3 align="center">This is a chat model, which can currently answer questions to IGL Docs provided.</h3>"""


with gr.Blocks(theme="glass") as demo:
    
    gr.HTML(title)

    chatbot = gr.Chatbot(type="messages", height=750)
    msg = gr.Textbox()
    send_button = gr.Button("Send")
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history):
        bot_message = generate_response(message, chat_history)
        bot_message = bot_message.replace("\\times", "*").replace("\\text{", "").replace("}", "")
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": bot_message})
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    send_button.click(respond, [msg, chatbot], [msg, chatbot])

    gr.Examples([
        ["How is penalty to be deducted from O&M bill of CNG 1200 SCMH compressor package against gas loss of 60,000 SCM in a month. Explain with detailed calculations."],
        ["How should a MDPE pipe cross a Nallah?"],
        ["At what depth should a MDPE pipe be laid?"],
        ["Explain with example RFC shortfall calculation for vendor of Pool 1."],
    ], inputs=msg, label= "Examples"
    )

    clear.click(lambda: None, None, chatbot, queue=False)
    gr.HTML(description)


demo.launch()