import gradio as gr
from transformers import pipeline

pipe = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct", device_map="cpu")

# Dê nomes descritivos e tipos claros. O MCP usa isso para o agente "entender" a ferramenta.
def generate_story_or_chat(prompt: str, history: list):
    """
    Use esta função para conversar com o assistente Qwen ou gerar textos criativos.
    
    Parameters:
    - prompt: A mensagem atual ou pedido do usuário.
    - history: O contexto da conversa anterior.
    """
    # (Mesma lógica de memória que corrigimos antes)
    messages_formatted = []
    for user_msg, bot_msg in history:
        messages_formatted.append({"role": "user", "content": user_msg})
        messages_formatted.append({"role": "assistant", "content": bot_msg})
    messages_formatted.append({"role": "user", "content": prompt})
    
    results = pipe(messages_formatted, max_new_tokens=512)
    return results[0]['generated_text'][-1]['content']

# A interface define como o MCP vai enxergar isso
demo = gr.ChatInterface(
    fn=generate_story_or_chat,
    title="Qwen Agent Node",
    description="Um endpoint MCP para geração de texto e chat usando Qwen 2.5."
)

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