import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# 1. Carrega Tokenizer e Modelo separadamente para garantir o controle do Contexto
model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Força o tokenizer a entender que aceitamos até 16k tokens (DeepSeek suporta 128k, mas RAM limita)
tokenizer.model_max_length = 16384 

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cpu",
    torch_dtype="auto",
    max_length=16384
)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

def predict(message, history):
    """
    Com type="messages", o 'history' já vem no formato correto:
    [{'role': 'user', 'content': 'oi'}, {'role': 'assistant', 'content': 'olá'}]
    """
    
    # Criamos a lista de mensagens adicionando a mensagem atual ao histórico
    # O history vem do Gradio, o message é o input atual
    if message is None:
        message = ""
    
    print(f"DEBUG INPUT: message={repr(message)}")
    print(f"DEBUG INPUT: history={repr(history)}")

    new_messages = history + [{"role": "user", "content": message}]
    
    # Sanitize messages to ensure no None content
    messages = []
    for msg in new_messages:
        content = msg.get("content")
        if content is None:
            content = ""
        messages.append({"role": msg.get("role"), "content": content})

    print(f"DEBUG FINAL MESSAGES: {messages}")

    # Gera a resposta (Aumentando output para raciocínio longo)
    results = pipe(messages, max_new_tokens=2048, truncation=True)
    
    # Retorna apenas o texto novo gerado
    return results[0]['generated_text'][-1]['content']

# 2. Configuração da Interface
# Adicionamos type="messages" para corrigir o aviso e padronizar com o OpenClaw/MCP
demo = gr.ChatInterface(
    fn=predict,
    type="messages",  # <--- ESTA LINHA CORRIGE O AVISO E MUDA O FORMATO DO HISTÓRICO
    title="Qwen Agent Node",
    description="Bot rodando Qwen 2.5 Coder 1.5B via CPU no Hugging Face."
)

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