import gradio as gr
import requests
import os

# Настройки
API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-3.2-3B-Instruct"
token = os.getenv("HF_TOKEN")
headers = {"Authorization": f"Bearer {token}"}

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

def respond(message, history):
    # Flare - имя ассистента
    prompt = f"System: Ты — Flare, профессиональный ИИ-ассистент. Пиши грамотно на русском.\n"
    
    # Берем только последние 2 сообщения для максимальной стабильности
    for user_msg, bot_msg in history[-2:]:
        prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n"
    
    prompt += f"User: {message}\nAssistant:"

    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": 500,
            "temperature": 0.7,
            "return_full_text": False
        }
    }

    try:
        output = query(payload)
        # Если API вернул ошибку (например, модель грузится)
        if "error" in output:
            return f"Flare: Сейчас модель подгружается, подожди 10 секунд. (Ошибка: {output['error']})"
        
        return output[0]['generated_text'].strip()
    except Exception as e:
        return f"Flare: Ошибка связи. Проверь интернет или токен. ({str(e)})"

# Самый простой интерфейс без стриминга (так надежнее)
demo = gr.ChatInterface(fn=respond)

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