import gradio as gr
import os
from smolagents import HfApiModel


model = HfApiModel(model_id="mistralai/Mixtral-8x7B-Instruct-v0.1", token=os.environ.get("HF_TOKEN"))
system_data = [
    {
            "role":"system",
            "content":[
                {
                    "type":"text",
                    "text": "You are a doctor who specializes on helping patients with addiction issues"
                }
            ]
        }
]

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)

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="Hello this is chatbot is created for only educational purpose and is powered by mistral 8x 7b model").queue()

demo.launch()