app.py

import os
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import List, Dict, Union, Optional

# --- NOVAS IMPORTAÇÕES PARA TRANSFORMERS E PyTorch ---
try:
    from transformers import AutoTokenizer, AutoModelForCausalLM
    import torch
except ImportError:
    print("Erro: As bibliotecas 'transformers' ou 'torch' não foram encontradas.")
    print("Certifique-se de instalá-las: pip install transformers torch")
    exit(1)

# Carregar variáveis de ambiente (MODEL_PATH não será mais usado diretamente aqui)
load_dotenv()

# --- Configurações do Modelo ---
# O nome do modelo no Hugging Face Hub
HF_MODEL_ID = "pierreguillou/gpt2-small-portuguese"

# Instâncias globais para o modelo e tokenizer
tokenizer: Optional[AutoTokenizer] = None
model: Optional[AutoModelForCausalLM] = None

# --- Inicialização do FastAPI ---
app = FastAPI(
    title="AutisMind AI Server",
    description="API para interação com o modelo GPT2-small-portuguese e análise de conversas.",
    version="1.0.0",
)

# --- Configuração CORS ---
origins = [
    "http://localhost:3000",
    "http://localhost:5000",
    "http://127.0.0.1:3000",
    "http://127.0.0.1:5000",
    # Adicione aqui os domínios do seu frontend e do Express API quando estiverem em produção
    # Ex: "https://seunome-seu-space-name.hf.space"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- Modelos Pydantic para Requisições e Respostas (Mantidos) ---
class ChatMessage(BaseModel):
    role: str # "user" ou "assistant"
    content: str

class ChatRequest(BaseModel):
    message: str = Field(..., description="A mensagem do usuário.")
    history: List[ChatMessage] = Field(default_factory=list, description="Histórico da conversa (opcional).")
    persona: str = Field(
        "Você é um assistente útil e amigável. Responda de forma clara e empática.",
        description="A personalidade do personagem para a IA."
    )
    chatId: Optional[str] = Field(None, description="ID do chat para contexto (não usado pela IA, mas pode ser útil para logs).")

class AnalysisResult(BaseModel):
    sentiment: str = Field("neutro", description="Sentimento da mensagem (positivo, negativo, neutro).")
    emotion: str = Field("desconhecida", description="Emoção detectada (ex: alegria, tristeza, raiva).")

class ChatResponse(BaseModel):
    ai_response: str = Field(..., description="A resposta gerada pela IA.")
    analysis: AnalysisResult = Field(..., description="Resultados da análise da mensagem do usuário ou da resposta da IA.")

# --- Evento de inicialização da aplicação (carrega o modelo) ---
@app.on_event("startup")
async def startup_event():
    global model, tokenizer
    print(f"Carregando modelo Transformers: {HF_MODEL_ID}")
    try:
        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID)
        # O GPT-2 não tem um token PAD padrão, define um.
        # 50256 é o token EOS padrão do GPT-2, pode ser usado como PAD também.
        tokenizer.pad_token = tokenizer.eos_token 
        
        # Ajusta o tamanho máximo de sequência do tokenizer para algo comum como 1024
        tokenizer.model_max_length = 1024

        model = AutoModelForCausalLM.from_pretrained(HF_MODEL_ID)
        model.eval() # Coloca o modelo em modo de avaliação (desativa dropout, etc.)
        
        # Tenta mover o modelo para a GPU se disponível e compatível com PyTorch/Transformers
        if torch.cuda.is_available():
            model.to("cuda")
            print("Modelo movido para GPU!")
        else:
            print("Modelo carregado na CPU.")

        print(f"Modelo {HF_MODEL_ID} carregado com sucesso!")
    except Exception as e:
        print(f"Erro ao carregar o modelo {HF_MODEL_ID}: {e}")
        raise RuntimeError(f"Falha ao carregar o modelo LLM: {e}")

# --- Endpoint de Saúde ---
@app.get("/health")
async def health_check():
    if model and tokenizer: # Verifica se ambos foram carregados
        return {"status": "ok", "model_loaded": True}
    return {"status": "loading", "model_loaded": False}

# --- Endpoint Principal para Geração de Resposta e Análise ---
@app.post("/generate-response", response_model=ChatResponse)
async def generate_ai_response(request: ChatRequest):
    global model, tokenizer
    if not model or not tokenizer:
        raise HTTPException(status_code=503, detail="Modelo AI ainda não carregado. Tente novamente em breve.")

    user_message = request.message
    chat_history = request.history
    character_persona = request.persona

    # --- Construir o Prompt para GPT-2 ---
    # GPT-2 é um modelo generativo, não segue o formato <|system|> como Phi.
    # A persona e o histórico são concatenados.
    prompt_parts = []
    # Inclui a persona no início do prompt, como uma instrução inicial para o modelo.
    prompt_parts.append(f"{character_persona}\n\n") 
    
    # Adicionar histórico da conversa
    # É importante limitar o histórico para não exceder o model_max_length (1024 tokens para GPT-2 small).
    # Uma boa estratégia é adicionar o histórico mais recente.
    # O GPT-2 não tem roles, então formatamos como um diálogo.
    for turn in chat_history[-5:]: # Inclui os últimos 5 turnos (adaptar conforme necessidade)
        if turn.role == 'user':
            prompt_parts.append(f"Usuário: {turn.content}\n")
        elif turn.role == 'assistant':
            prompt_parts.append(f"Assistente: {turn.content}\n")
    
    # Adiciona a mensagem atual do usuário e pede para a IA responder
    prompt_parts.append(f"Usuário: {user_message}\nAssistente:") 
    
    full_prompt = "".join(prompt_parts)

    try:
        # Codifica o prompt para tokens
        # return_tensors="pt" para PyTorch
        # max_length para garantir que o prompt não exceda o limite do modelo
        inputs = tokenizer(
            full_prompt, 
            return_tensors="pt", 
            max_length=tokenizer.model_max_length, 
            truncation=True
        )
        
        # Mover os inputs para a GPU se o modelo estiver na GPU
        if torch.cuda.is_available():
            inputs = {k: v.to("cuda") for k, v in inputs.items()}

        # Gerar a resposta
        # pad_token_id: O token usado para preencher sequências, aqui usamos o EOS token do GPT-2
        # do_sample=True: Habilita amostragem (para criatividade)
        # top_k, top_p: Métodos de amostragem para controlar a diversidade da resposta
        # max_new_tokens: O número máximo de novos tokens que a IA pode gerar APÓS o prompt
        sample_outputs = model.generate(
            inputs.input_ids,
            pad_token_id=tokenizer.pad_token_id, # Usando o token de pad que definimos
            do_sample=True,
            max_new_tokens=150, # Gera até 150 tokens novos de resposta
            temperature=0.7,
            top_k=50,
            top_p=0.9,
            num_return_sequences=1
        )
        
        # Decodificar a resposta gerada
        # skip_special_tokens=True: Remove tokens especiais como <bos>, <eos>, <pad>
        generated_text = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
        
        # --- Pós-processamento da Resposta do GPT-2 ---
        # GPT-2 vai gerar o prompt INTEIRO MAIS a resposta.
        # Precisamos remover o prompt original para obter apenas a resposta do "Assistente:".
        # Isso pode ser um pouco delicado e pode precisar de ajustes finos dependendo da saída real do modelo.
        ai_response_raw = generated_text.replace(full_prompt, "").strip()

        # O GPT-2 pode continuar gerando texto ou até começar um "Usuário:" de novo.
        # Tentamos cortar a resposta na primeira ocorrência de "Usuário:" ou "Assistente:" para evitar isso.
        if "\nUsuário:" in ai_response_raw:
            ai_response = ai_response_raw.split("\nUsuário:")[0].strip()
        elif "\nAssistente:" in ai_response_raw:
            ai_response = ai_response_raw.split("\nAssistente:")[0].strip()
        else:
            ai_response = ai_response_raw


        # --- Lógica de Análise (Mantida) ---
        analysis_result = AnalysisResult() 
        user_message_lower = user_message.lower()
        if "triste" in user_message_lower or "chateado" in user_message_lower or "mal" in user_message_lower:
            analysis_result.sentiment = "negativo"
            analysis_result.emotion = "tristeza"
        elif "feliz" in user_message_lower or "alegre" in user_message_lower or "bom" in user_message_lower:
            analysis_result.sentiment = "positivo"
            analysis_result.emotion = "alegria"
        else:
            analysis_result.sentiment = "neutro"
            analysis_result.emotion = "desconhecida"
        
        return ChatResponse(ai_response=ai_response, analysis=analysis_result)

    except Exception as e:
        print(f"Erro na geração da IA: {e}")
        raise HTTPException(status_code=500, detail=f"Erro interno do servidor AI: {str(e)}")

# --- Executar o Servidor (apenas para teste direto via Python) ---
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5001)