import gradio as gr
import requests

from dotenv import load_dotenv
import os


load_dotenv()
API_URL = "https://router.huggingface.co/v1/chat/completions"
headers = {
    "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()
history  = []


def chat(message, history):
    messages = []

    # Add previous conversation
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})

    # Add current user message
    messages.append({"role": "user", "content": message})

    response = query({
        "model": "meta-llama/Llama-3.1-8B-Instruct:novita",
        "messages": messages
    })

    return response["choices"][0]["message"]["content"]

gr.ChatInterface(
    fn=chat,
    save_history=True
).launch(theme="ocean")