import gradio as gr
from huggingface_hub import InferenceClient
import os

HF_API_TOKEN = os.getenv("HF_TOKEN")
MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2"

client = InferenceClient(model=MODEL_ID, token=HF_API_TOKEN)

def chat_with_bot(message, history):
    # Format history for the chat API
    messages = []
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    # Use chat completion
    response = client.chat_completion(
        messages=messages,
        max_tokens=500,
        temperature=0.7,
    )
    
    return response.choices[0].message.content

iface = gr.ChatInterface(
    fn=chat_with_bot,
    title="🤖 Mistral Chatbot",
    description="Chat with the Mistral-7B model hosted on Hugging Face!",
    theme="soft"
)

if __name__ == "__main__":
    iface.launch()