Create api_logic.py

api_logic.py

@@ -0,0 +1,267 @@
+import os
+import json
+import faiss
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from sentence_transformers import SentenceTransformer
+from huggingface_hub import HfApi # Pode ser removido se não usar para commits externos
+import logging
+
+# Configuração de Logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# --- Configurações Globais ---
+MODEL_NAME    = "stabilityai/stablelm-3b-4e1t"
+EMBED_MODEL   = "all-MiniLM-L6-v2" # Usando nome curto que funciona
+SUMMARIZER_ID = "sshleifer/distilbart-cnn-12-6"
+MEM_FILE      = "memory.json" # Salvará no armazenamento do Space
+IDX_FILE      = "index.faiss" # Salvará no armazenamento do Space
+# HF_TOKEN      = os.getenv("HF_TOKEN") # Se precisar para algo no futuro
+# REPO_ID       = os.getenv("SPACE_ID") # Se precisar para algo no futuro
+
+# --- Variáveis Globais para Modelos e Memória (Carregados uma vez) ---
+tokenizer = None
+model = None
+chat_pipe = None
+embedder = None
+summarizer = None
+memory = []
+index = None
+faiss_dimension = 384 # Dimensão padrão para all-MiniLM-L6-v2
+
+def load_models_and_memory():
+    """Carrega modelos, pipelines e a memória/índice FAISS."""
+    global tokenizer, model, chat_pipe, embedder, summarizer, memory, index, faiss_dimension
+
+    logger.info("Iniciando carregamento de modelos e memória...")
+
+    # Carrega Tokenizer e Modelo LLM (FP16)
+    logger.info(f"Carregando tokenizer: {MODEL_NAME}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+    logger.info(f"Carregando modelo: {MODEL_NAME} (FP16, CPU)")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True, # Importante para CPU com RAM limitada
+        device_map="auto" # Deixa accelerate decidir (provavelmente 'cpu')
+    )
+    logger.info("Modelo LLM carregado.")
+
+    # Cria Pipeline de Chat
+    logger.info("Criando pipeline de text-generation...")
+    chat_pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        # max_length=512, # Definir max_new_tokens na chamada é melhor
+        do_sample=True,
+        top_p=0.9,
+    )
+    logger.info("Pipeline de chat criado.")
+
+    # Carrega Embedder
+    logger.info(f"Carregando embedder: {EMBED_MODEL}")
+    embedder = SentenceTransformer(EMBED_MODEL)
+    faiss_dimension = embedder.get_sentence_embedding_dimension() # Atualiza a dimensão
+    logger.info(f"Embedder carregado. Dimensão: {faiss_dimension}")
+
+
+    # Carrega Summarizer
+    logger.info(f"Carregando summarizer: {SUMMARIZER_ID}")
+    summarizer = pipeline("summarization", model=SUMMARIZER_ID)
+    logger.info("Summarizer carregado.")
+
+    # Carrega/Inicializa Memória e Índice FAISS
+    logger.info("Carregando/Inicializando memória e índice FAISS...")
+    if os.path.exists(MEM_FILE):
+        try:
+            with open(MEM_FILE, "r") as f:
+                memory = json.load(f)
+            logger.info(f"Arquivo de memória '{MEM_FILE}' carregado ({len(memory)} entradas).")
+        except Exception as e:
+            logger.error(f"Erro ao carregar {MEM_FILE}: {e}. Iniciando memória vazia.")
+            memory = []
+    else:
+        logger.info(f"Arquivo '{MEM_FILE}' não encontrado. Iniciando memória vazia.")
+        memory = []
+
+    if os.path.exists(IDX_FILE):
+        try:
+            index = faiss.read_index(IDX_FILE)
+            logger.info(f"Índice FAISS '{IDX_FILE}' carregado ({index.ntotal} vetores).")
+            # Validação simples da dimensão
+            if index.ntotal > 0 and index.d != faiss_dimension:
+                 logger.warning(f"Dimensão do índice FAISS ({index.d}) diferente da dimensão do embedder ({faiss_dimension})! Recriando índice.")
+                 index = faiss.IndexFlatL2(faiss_dimension)
+                 # Idealmente, aqui você re-indexaria a memória existente, mas vamos simplificar por agora
+        except Exception as e:
+            logger.error(f"Erro ao carregar {IDX_FILE}: {e}. Recriando índice FAISS.")
+            index = faiss.IndexFlatL2(faiss_dimension)
+    else:
+        logger.info(f"Arquivo '{IDX_FILE}' não encontrado. Criando novo índice FAISS.")
+        index = faiss.IndexFlatL2(faiss_dimension)
+
+    logger.info("Carregamento de modelos e memória concluído.")
+
+
+# --- Funções de Lógica (Adaptadas) ---
+
+def save_state():
+    """Salva o estado atual da memória e do índice FAISS."""
+    logger.info("Salvando estado (memória e índice)...")
+    try:
+        with open(MEM_FILE, "w") as f:
+            json.dump(memory, f, indent=2) # Adicionado indent para legibilidade
+        faiss.write_index(index, IDX_FILE)
+        logger.info("Estado salvo com sucesso.")
+    except Exception as e:
+        logger.error(f"Erro ao salvar estado: {e}")
+
+def summarize_block(txt):
+    logger.info("Chamando summarizer...")
+    instr = ("Resuma este trecho de diálogo de RPG preservando personagens, locais, itens e eventos:\n\n")
+    try:
+        summary = summarizer(instr + txt, max_length=150, min_length=50)[0]["summary_text"] # Reduzi min_length
+        logger.info("Resumo gerado.")
+        return summary
+    except Exception as e:
+        logger.error(f"Erro no summarizer: {e}")
+        return f"[[Erro ao resumir: {e}]]"
+
+def compact_memory(threshold=50): # Reduzi o threshold para testar mais rápido
+    logger.info(f"Verificando compactação (limite={threshold}), memória atual={len(memory)}...")
+    if len(memory) < threshold:
+        return False # Indica que não compactou
+
+    logger.info(f"Compactando memória ({threshold} itens)...")
+    bloco = memory[:threshold]
+    texto = "\n".join(
+        itm["text"] if itm["type"]=="summary"
+        else f"Usuário: {itm['user']}\nIA: {itm['bot']}"
+        for itm in bloco
+    )
+    resumo = summarize_block(texto)
+    memory[:threshold] = [{"type":"summary","text":resumo}] # Substitui bloco por resumo
+
+    logger.info("Recriando índice FAISS após compactação...")
+    new_idx = faiss.IndexFlatL2(faiss_dimension)
+    embeddings_to_add = []
+    for itm in memory:
+        key = itm["text"] if itm["type"]=="summary" else itm["user"]
+        try:
+            # Coleta todos os embeddings primeiro
+            embeddings_to_add.append(embedder.encode([key], convert_to_numpy=True)[0])
+        except Exception as e:
+            logger.error(f"Erro ao encodar item para reindexação: {e}")
+
+    if embeddings_to_add:
+        try:
+            new_idx.add(np.array(embeddings_to_add)) # Adiciona em lote
+            global index
+            index = new_idx
+            logger.info("Reindexação FAISS concluída.")
+            save_state() # Salva após compactar e reindexar
+            return True # Indica que compactou
+        except Exception as e:
+             logger.error(f"Erro ao adicionar embeddings ao novo índice: {e}")
+             return False
+    else:
+        logger.info("Nenhum embedding válido para adicionar ao novo índice.")
+        return False
+
+
+def add_to_memory_and_index(user_msg, bot_msg):
+    logger.info("Adicionando nova entrada à memória e índice...")
+    entry = {"type":"dialog", "user":user_msg, "bot":bot_msg, "text":user_msg}
+    memory.append(entry)
+    try:
+        embedding = embedder.encode([user_msg], convert_to_numpy=True)
+        index.add(embedding)
+        logger.info(f"Embedding adicionado ao índice. Total de vetores: {index.ntotal}")
+        save_state() # Salva após adicionar
+        compact_memory() # Verifica se precisa compactar
+    except Exception as e:
+        logger.error(f"Erro ao adicionar embedding ou salvar estado: {e}")
+
+
+def run_chat_logic(user_msg):
+    """Executa a lógica principal do chat: busca contexto, gera resposta."""
+    logger.info(f"Executando lógica do chat para: {repr(user_msg)}")
+    global memory, index # Garante acesso às variáveis globais
+
+    if not all([tokenizer, model, chat_pipe, embedder, index]):
+         logger.error("Modelos ou índice não foram carregados corretamente.")
+         return "Desculpe, o sistema de IA não está pronto. Tente novamente mais tarde."
+
+    # 1. Embedding e Busca FAISS
+    logger.info("Gerando embedding da mensagem do usuário...")
+    try:
+        emb = embedder.encode([user_msg], convert_to_numpy=True)
+    except Exception as e:
+        logger.error(f"Erro ao gerar embedding: {e}")
+        return "Desculpe, houve um erro ao processar sua mensagem (embedding)."
+
+    context = []
+    logger.info(f"Buscando no índice FAISS ({index.ntotal} vetores)...")
+    if index.ntotal > 0:
+        try:
+            D, I = index.search(emb, k=5) # Busca os 5 vizinhos mais próximos
+            logger.info(f"Índices FAISS encontrados: {I[0]}")
+            for idx in I[0]:
+                if 0 <= idx < len(memory): # Validação crucial
+                    itm = memory[idx]
+                    context.append(
+                        f"Lembrança: {itm['text']}" if itm["type"]=="summary"
+                        else f"Histórico [Usuário: {itm['user']} | IA: {itm['bot']}]"
+                    )
+                else:
+                    logger.warning(f"Índice FAISS inválido encontrado: {idx}")
+        except Exception as e:
+            logger.error(f"Erro durante a busca FAISS: {e}")
+            # Continua sem contexto se a busca falhar
+
+    logger.info(f"Contexto recuperado ({len(context)} itens).")
+
+    # 2. Monta Prompt
+    context_str = "\n".join(context)
+    prompt = (
+        f"Você é um Mestre de RPG experiente e criativo. Continue a história de forma envolvente, "
+        f"considerando o seguinte histórico e lembranças:\n{context_str}\n\n"
+        f"Ação do Jogador: {user_msg}\n\nSua Narração:"
+    )
+    logger.info(f"Prompt enviado ao modelo (primeiros 200 chars):\n{prompt[:200]}...")
+
+    # 3. Chama o Modelo (Pipeline)
+    try:
+        logger.info("Chamando pipeline text-generation...")
+        # return_full_text=False pega só a continuação
+        outputs = chat_pipe(prompt, max_new_tokens=200, return_full_text=False, num_return_sequences=1)
+        logger.info(f"Saída bruta do pipeline: {outputs}")
+
+        if not outputs or not outputs[0] or "generated_text" not in outputs[0]:
+             logger.error("Pipeline não retornou 'generated_text' válido.")
+             return "Desculpe, a IA não conseguiu gerar uma resposta válida."
+
+        bot_msg = outputs[0]["generated_text"].strip()
+        # Limpeza adicional
+        bot_msg = bot_msg.split("<|endoftext|>")[0].strip()
+        # Remover repetições exatas do prompt final se houver
+        if bot_msg.startswith(f"Sua Narração:"): bot_msg = bot_msg[len("Sua Narração:"):].strip()
+
+        logger.info(f"Resposta processada do bot: {repr(bot_msg)}")
+
+        # 4. Adiciona na memória e índice
+        add_to_memory_and_index(user_msg, bot_msg)
+
+        return bot_msg
+
+    except Exception as e:
+        logger.exception("Erro durante a execução do pipeline ou pós-processamento.") # Loga o traceback completo
+        return f"Desculpe, ocorreu um erro interno ao gerar a resposta da IA: {e}"
+
+# --- Carrega tudo na inicialização ---
+# Esta linha será chamada quando o módulo for importado pela primeira vez
+# (Ou podemos chamar explicitamente via lifespan do FastAPI)
+# load_models_and_memory()