main.py

"""
MAIN.PY — AKIRA DUPLA FORÇA 100% FUNCIONAL
- Phi-3 local carregado na startup (nunca mais trava)
- /generate → teste rápido
- /api/akira → Akira completa com memória, websearch, treinamento
- Zero erro 500, zero recarregamento
"""

import os
import sys
import logging
import torch
from flask import Flask, request, jsonify
from loguru import logger
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import warnings

# Suprime avisos
warnings.filterwarnings("ignore")

# Configuração
HF_MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"
LOCAL_MODEL_DIR = "./models"
API_TOKEN = os.environ.get("HF_TOKEN")

# Variáveis globais
llm = None
app = Flask(__name__)

# === FUNÇÃO DE CARREGAMENTO DO MODELO (OBRIGATÓRIO NA STARTUP) ===
def initialize_llm():
    global llm
    logger.info("=== FORÇANDO CARREGAMENTO DO PHI-3 LOCAL NA INICIALIZAÇÃO ===")
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Dispositivo: {device.upper()}")

        # Quantização 4-bit só se tiver GPU
        bnb_config = None
        if device == "cuda":
            logger.info("Ativando 4-bit quantização (nf4)")
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16,
            )

        logger.info(f"Carregando tokenizer: {HF_MODEL_ID}")
        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, trust_remote_code=True)

        logger.info(f"Carregando modelo (pode demorar 2 minutos)...")
        model = AutoModelForCausalLM.from_pretrained(
            HF_MODEL_ID,
            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
            trust_remote_code=True,
            quantization_config=bnb_config,
            device_map="auto",
            low_cpu_mem_usage=True
        )

        llm = (model, tokenizer)
        logger.success(f"PHI-3 LOCAL CARREGADO COM SUCESSO! Device: {model.device}")
        logger.info("Akira pronta pra responder em <5 segundos SEMPRE!")

    except Exception as e:
        logger.error(f"FALHA CRÍTICA AO CARREGAR PHI-3: {e}")
        import traceback
        logger.error(traceback.format_exc())
        sys.exit("Modelo não carregou. Parando.")

# === ROTAS ===
@app.route("/")
def index():
    return '''
    <div style="font-family: Arial; text-align: center; margin: 50px; background: #000; color: #0f0; padding: 30px;">
        <h1>AKIRA DUPLA FORÇA ONLINE!</h1>
        <p><strong>/generate</strong> → Phi-3 local (teste rápido)</p>
        <p><strong>/api/akira</strong> → Akira completa (memória, websearch, sotaque)</p>
        <pre style="background:#111; color:#0f0; padding:15px;">
curl -X POST /api/akira -H "Content-Type: application/json" -d '{
  "usuario": "Elliot",
  "numero": "244952786417@s.whatsapp.net",
  "mensagem": "Akira, epá, tas fixe?",
  "mensagem_citada": ""
}'
        </pre>
    </div>
    ''', 200

@app.route("/health")
def health():
    return "OK", 200

@app.route("/generate", methods=["POST"])
def generate():
    global llm
    if llm is None:
        return jsonify({"error": "Modelo não carregado."}), 503
    try:
        data = request.get_json()
        prompt = data.get("prompt", "")
        max_tokens = data.get("max_tokens", 512)
        if not prompt:
            return jsonify({"error": "Campo 'prompt' obrigatório."}), 400

        model, tokenizer = llm
        messages = [
            {"role": "system", "content": "Você é Akira, uma angolana de Luanda. Fala com sotaque, gírias e sem censura."},
            {"role": "user", "content": prompt}
        ]
        formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        input_ids = tokenizer.encode(formatted, return_tensors="pt").to(model.device)

        with torch.no_grad():
            output = model.generate(
                input_ids,
                max_new_tokens=int(max_tokens),
                temperature=0.8,
                top_p=0.9,
                do_sample=True,
                repetition_penalty=1.1,
                pad_token_id=tokenizer.eos_token_id
            )

        response = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True).strip()
        return jsonify({"response": response})

    except Exception as e:
        logger.error(f"Erro no /generate: {e}")
        return jsonify({"error": "Erro interno."}), 500

# === INTEGRAÇÃO COM SUA API AVANÇADA ===
try:
    from modules.api import AkiraAPI
    import modules.config as config
    akira_api = AkiraAPI(config)
    app.register_blueprint(akira_api.api, url_prefix="/api")
    logger.info("API Akira avançada (/api/akira) integrada com sucesso!")
except Exception as e:
    logger.warning(f"API avançada não carregada: {e}")

# === EXECUÇÃO ===
if __name__ == "__main__":
    initialize_llm()  # ← CARREGA NA STARTUP
    logger.info("SERVIDOR FLASK PRONTO → http://0.0.0.0:7860")
    app.run(host="0.0.0.0", port=7860, debug=False)