import gradio as gr
import os
from smolagents import HfApiModel


model = HfApiModel(model_id="mistralai/Mistral-7B-Instruct-v0.3", token=os.environ.get("HF_TOKEN"))
system_data = [
    {
            "role":"system",
            "content":[
                {
                    "type":"text",
                    "text": """
                    You are a helpful healthcare assistant who specializes on helping patients with addiction issues.
                    Make sure to be as emphaethetic as possible and be more personal than generic in 
                    your responses as possible. Try your best to help people fight their addiction. 
                    You should never give dangerous advice at any point in time. Also include this term at
                    then end of your response 'Any Other Question Do You Have?'. 
                    Always keep your answers short and crisp.
                    """
                }
            ]
        }
]

def get_user_data(prompt: str):
    return [
        {
            "role":"user",
            "content":[
                {
                    "type":"text",
                    "text": prompt
                }
            ]
        }
    ]

def get_history(history):
    mod_history = []
    for user_message, bot_message in history:
        user_dict = {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": user_message
                }
            ]
        }
        bot_dict = {
            "role": "assistant",
            "content": [
                {
                    "type": "text",
                    "text": bot_message
                }
            ]
        }
        mod_history.append(user_dict)
        mod_history.append(bot_dict)
    return mod_history

def chat(prompt, history):
    
    return model(system_data + get_history(history)+ get_user_data(prompt)).content

demo = gr.ChatInterface(chat, 
                        chatbot=gr.Chatbot(),
                        title="ArunGPT",
                        theme = gr.themes.Soft(), 
                        description="""
                        **DISCLAIMER:** THIS CHATBOT IS POWERED BY GENAI MODELS, AUTHOR OR CREATOR DOESN'T TAKE
                        ANY RESPONSIBLITY FOR AN ADVERSE EVENT HAPPENS AS RESULT OF MODEL'S SUGGESTIONS.
                        REMEMBER IT IS THE USERS RESPONSIBLITY TO CROSSCHECK THE ADVICE WITH PROFESSIONALS
                        AND ACT ON IT. 
                        """).queue()

demo.launch()