import gradio as gr 
from transformers import Conversation, pipeline
from typing import Union

chatbot = pipeline(model="facebook/blenderbot-400M-distill")

def process(prompt:Union[str,None]=None,history:Union[str,None]=None):
    """" process prompt
    history has the following format 
    ```
    ****user \t hello
    ****assistant \t hi
    ****user \t how are you
    ****assistant \t good
    """
    conv = Conversation()
    if history != None : 
        for hist in history.split("****") : 
            if len(hist.strip()) == 0 : continue
            entry = hist.split("\t")
            conv.add_message({"role": entry[0].strip(), "content": entry[1].strip()})
    if prompt is not None or prompt != "":
        conv.add_message({"role": "user", "content": prompt})
    return chatbot(conv).messages[-1]["content"]


prompt = gr.Textbox(label="Prompt")
history = gr.Textbox(label="History")
out = gr.Textbox(label="Response")
examples = gr.Examples([
    ["",""" ****user \t hello \n ***assistant \t hi \n ****user \t how are you"""],
    ["what's 2+10",""" ****user \t hi \n ***assistant \t hello \n"""]
    ],
    inputs=[prompt,history],
    outputs=[out],
    cache_examples=True)

iface = gr.Interface(fn=process,inputs=[prompt,history],outputs=[out])
iface.launch()