Update app.py

app.py

@@ -1,7 +1,7 @@
 # ==============================================================================
-#  API do AetherMap — VERSÃO SÁBIA 6.0 (O SÁBIO INVOCADO)
-#  Backend com RAG (Retrieval-Augmented Generation) na busca semântica.
-#  Todas as funcionalidades anteriores estão presentes e aprimoradas.
+#  API do AetherMap — VERSÃO 6.5 (THE COMMAND KILLER)
+#  Backend com RAG em Dois Estágios (Retrieval + Reranking) e Citações Nativas.
+#  Arquitetura otimizada por Berta & Gabriel.
 # ==============================================================================
 
 import numpy as np
@@ -18,7 +18,7 @@ from typing import List, Dict, Any
 from functools import lru_cache
 
 # Ferramentas de Alquimia
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, CrossEncoder # <--- A ARMA SECRETA
 import umap
 import hdbscan
 from sklearn.preprocessing import StandardScaler
@@ -34,7 +34,11 @@ from groq import Groq
 # ==============================================================================
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
-DEFAULT_MODEL = "all-MiniLM-L6-v2"
+# Bi-Encoder para varredura rápida
+RETRIEVAL_MODEL = "all-MiniLM-L6-v2"
+# Cross-Encoder para precisão cirúrgica (Reranking)
+RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+
 BATCH_SIZE = 256
 UMAP_N_NEIGHBORS = 30
 
@@ -52,7 +56,7 @@ except Exception as e:
     logging.error(f"FALHA CRÍTICA AO INICIALIZAR GROQ: {e}")
     groq_client = None
 
-# Palavras de Parada
+# Palavras de Parada (Mantidas da versão anterior)
 STOP_WORDS_PT = [
     'de','a','o','que','e','do','da','em','um','para','é','com','não','uma','os','no',
     'se','na','por','mais','as','dos','como','mas','foi','ao','ele','das','tem','à',
@@ -80,13 +84,20 @@ STOP_WORDS_PT = [
 
 
 # ==============================================================================
-# FUNÇÕES DE ANÁLISE (INTACTAS DA VERSÃO COMPLETA)
+# FUNÇÕES DE ANÁLISE E CARREGAMENTO DE MODELOS
 # ==============================================================================
 @lru_cache(maxsize=1)
-def load_model():
+def load_retriever():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logging.info(f"Carregando Retriever '{RETRIEVAL_MODEL}' em: {device}")
+    return SentenceTransformer(RETRIEVAL_MODEL, device=device)
+
+@lru_cache(maxsize=1)
+def load_reranker():
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    logging.info(f"Carregando modelo '{DEFAULT_MODEL}' em: {device}")
-    return SentenceTransformer(DEFAULT_MODEL, device=device)
+    logging.info(f"Carregando Reranker '{RERANKER_MODEL}' em: {device}")
+    # O CrossEncoder processa par (query, doc) juntos. É mais lento, mas MUITO mais preciso.
+    return CrossEncoder(RERANKER_MODEL, device=device)
 
 def preparar_textos(file_bytes: bytes, n_samples: int) -> List[str]:
     linhas = file_bytes.decode("utf-8", errors="ignore").splitlines()
@@ -95,16 +106,21 @@ def preparar_textos(file_bytes: bytes, n_samples: int) -> List[str]:
 
 def processar_pipeline(textos: List[str]) -> (pd.DataFrame, np.ndarray):
     logging.info(f"Iniciando pipeline para {len(textos)} textos...")
-    model = load_model()
+    model = load_retriever()
     embeddings = model.encode(textos, batch_size=BATCH_SIZE, show_progress_bar=False, convert_to_numpy=True)
+    
+    # Redução dimensional e Clustering
     reducer = umap.UMAP(n_components=3, n_neighbors=UMAP_N_NEIGHBORS, min_dist=0.0, metric="cosine", random_state=42)
     emb_3d = reducer.fit_transform(embeddings)
     emb_3d = StandardScaler().fit_transform(emb_3d)
+    
     num_textos = len(textos)
     min_size = max(10, int(num_textos * 0.02)) 
     logging.info(f"HDBSCAN min_cluster_size dinâmico definido para: {min_size}")
+    
     clusterer = hdbscan.HDBSCAN(min_cluster_size=min_size)
     clusters = clusterer.fit_predict(emb_3d)
+    
     df = pd.DataFrame({"x": emb_3d[:, 0], "y": emb_3d[:, 1], "z": emb_3d[:, 2], "full_text": textos, "cluster": clusters.astype(str)})
     del reducer, clusterer, emb_3d; gc.collect()
     return df, embeddings
@@ -119,12 +135,14 @@ def calcular_metricas(textos: List[str]) -> Dict[str, Any]:
         tfidf_matrix = vectorizer_tfidf.fit_transform(textos)
     except ValueError:
         return {"riqueza_lexical": 0, "top_tfidf_palavras": [], "entropia": 0.0}
+    
     vocab_count = vectorizer_count.get_feature_names_out()
     contagens = counts_matrix.sum(axis=0).A1
     vocab_tfidf = vectorizer_tfidf.get_feature_names_out()
     soma_tfidf = tfidf_matrix.sum(axis=0).A1
     top_idx_tfidf = np.argsort(soma_tfidf)[-10:][::-1]
     top_tfidf = [{"palavra": vocab_tfidf[i], "score": round(float(soma_tfidf[i]), 4)} for i in top_idx_tfidf]
+    
     return {
         "riqueza_lexical": len(vocab_count), 
         "top_tfidf_palavras": top_tfidf, 
@@ -136,6 +154,8 @@ def encontrar_duplicados(df: pd.DataFrame, embeddings: np.ndarray) -> Dict[str,
     mask = df["full_text"].duplicated(keep=False)
     grupos_exatos = {t: [int(i) for i in idxs] for t, idxs in df[mask].groupby("full_text").groups.items()}
     pares_semanticos = []
+    
+    # Só roda similaridade pesada se não for gigante
     if 2 < len(embeddings) < 5000:
         sim = cosine_similarity(embeddings)
         triu_indices = np.triu_indices_from(sim, k=1)
@@ -169,9 +189,9 @@ def analisar_clusters(df: pd.DataFrame) -> Dict[str, Any]:
 
 
 # ==============================================================================
-# FASTAPI — DEFINIÇÃO DA API (COMPLETA)
+# FASTAPI — DEFINIÇÃO DA API
 # ==============================================================================
-app = FastAPI(title="API do AetherMap (O Sábio Invocado)", version="6.0.0")
+app = FastAPI(title="API do AetherMap 6.5 (The Command Killer)", version="6.5.0")
 
 @app.post("/process/")
 async def process_api(n_samples: int = Form(10000), file: UploadFile = File(...)):
@@ -209,60 +229,109 @@ async def process_api(n_samples: int = Form(10000), file: UploadFile = File(...)
 
 @app.post("/search/")
 async def search_api(query: str = Form(...), job_id: str = Form(...)):
-    logging.info(f"Busca RAG recebida para a query '{query}' no job '{job_id}'.")
+    """
+    RAG 2.0: Recuperação Híbrida (Bi-Encoder + Cross-Encoder Reranking)
+    Fornece citações precisas estilo Command R+.
+    """
+    logging.info(f"Busca RAG 2.0 recebida para: '{query}' [Job: {job_id}]")
     if job_id not in cache:
         raise HTTPException(status_code=404, detail="Job ID não encontrado ou expirado.")
 
     try:
-        # ETAPA 1: RECUPERAÇÃO (RETRIEVAL)
-        model = load_model()
+        # --- FASE 1: RECUPERAÇÃO (RETRIEVAL - BI-ENCODER) ---
+        model = load_retriever()
+        reranker = load_reranker() # Carrega o Juiz
+        
         cached_data = cache[job_id]
         df = cached_data["df"]
         corpus_embeddings = cached_data["embeddings"]
         
+        # Passo 1: Busca inicial ampla com Cosseno (Pega 50 candidatos)
+        # Isso garante que não perdemos nada relevante que tenha palavras-chave parecidas
         query_embedding = model.encode([query], convert_to_numpy=True)
         similarities = cosine_similarity(query_embedding, corpus_embeddings)[0]
         
-        top_k = 10
-        top_indices = np.argsort(similarities)[-top_k:][::-1]
+        top_k_retrieval = 50 
+        top_indices = np.argsort(similarities)[-top_k_retrieval:][::-1]
+        
+        # Prepara dados para o Reranker: Lista de pares [Query, Doc]
+        candidate_docs = []
+        candidate_indices = []
         
-        results = [
-            {"index": int(i), "score": float(similarities[i])}
-            for i in top_indices if similarities[i] > 0.3
-        ]
+        # Filtro de corte mínimo para não passar lixo total pro reranker
+        for idx in top_indices:
+            if similarities[idx] > 0.15: # Threshold leve
+                doc_text = df.iloc[int(idx)]["full_text"]
+                candidate_docs.append([query, doc_text])
+                candidate_indices.append(int(idx))
+
+        if not candidate_docs:
+             return {"summary": "Não foram encontrados documentos minimamente relevantes.", "results": []}
+
+        # --- FASE 2: REORDENAÇÃO (RERANKING - CROSS-ENCODER) ---
+        # O modelo Cross-Encoder lê a pergunta e o documento juntos e dá um score de verdade.
+        logging.info(f"Reordenando {len(candidate_docs)} documentos com Cross-Encoder...")
+        rerank_scores = reranker.predict(candidate_docs)
         
-        if not results:
-            return {"summary": "Não foram encontrados resultados relevantes para sua busca.", "results": []}
+        # Ordena pelos scores do reranker (do maior para o menor)
+        rerank_results = sorted(
+            zip(candidate_indices, rerank_scores), 
+            key=lambda x: x[1], 
+            reverse=True
+        )
 
-        # ETAPA 2: GERAÇÃO AUMENTADA (AUGMENTED GENERATION)
+        # Agora pegamos o Top 5 da "Nata da Nata" para enviar ao Kimi
+        final_top_k = 5
+        final_results = []
+        context_parts = []
+
+        for rank, (idx, score) in enumerate(rerank_results[:final_top_k]):
+            doc_text = df.iloc[idx]["full_text"]
+            # Montamos o contexto COM O ID EXPLÍCITO para forçar a citação
+            context_parts.append(f"[ID: {rank+1}] DOCUMENTO:\n{doc_text}\n---------------------")
+            
+            final_results.append({
+                "index": idx,
+                "score": float(score),           # Score do Reranker (Confiança semântica)
+                "cosine_score": float(similarities[idx]), # Score original (para debug)
+                "citation_id": rank + 1
+            })
+
+        context_str = "\n".join(context_parts)
+
+        # --- FASE 3: GERAÇÃO CITADA (READER - KIMI K2) ---
         summary = ""
         if groq_client:
-            context_docs = [df.iloc[res["index"]]["full_text"] for res in results[:5]]
-            context_str = "\n\n".join([f"Documento de referência {i+1}:\n'''{doc}'''" for i, doc in enumerate(context_docs)])
-
+            # Prompt de Sistema projetado para emular o comportamento do Command R+
             rag_prompt = (
-                "Você é Aetherius, um Sábio Conselheiro. Sua tarefa é responder à pergunta do usuário de forma concisa e direta, "
-                "baseando-se **estritamente** nas informações contidas nos documentos de referência fornecidos. "
-                "Não use nenhum conhecimento externo. Responda em uma ou duas frases.\n\n"
-                f"**Pergunta do Usuário:** \"{query}\"\n\n"
-                f"{context_str}\n\n"
-                "**Sua Resposta Direta:**"
+                "INSTRUÇÃO DE SISTEMA:\n"
+                "Você é o Aetherius, um motor de busca semântica de alta precisão.\n"
+                "Sua missão é responder à pergunta do usuário baseando-se ESTRITAMENTE nos documentos fornecidos.\n\n"
+                "REGRAS DE OURO:\n"
+                "1. CITAÇÕES OBRIGATÓRIAS: Toda afirmação factual deve ser seguida da fonte no formato [ID: x].\n"
+                "   Exemplo: 'O lucro subiu 10% [ID: 1], mas a margem caiu [ID: 2].'\n"
+                "2. HONESTIDADE INTELECTUAL: Se a resposta não estiver no contexto, diga 'Não encontrei informações suficientes nos documentos'.\n"
+                "3. ESTILO: Seja direto, técnico e conciso. Fale em Português do Brasil.\n\n"
+                f"CONTEXTO RECUPERADO (Ordenado por Relevância):\n{context_str}\n\n"
+                f"PERGUNTA DO USUÁRIO: \"{query}\"\n\n"
+                "RESPOSTA:"
             )
 
             try:
+                # Usando Kimi K2 pois ele tem ótimo raciocínio lógico para síntese
                 chat_completion = groq_client.chat.completions.create(
                     messages=[{"role": "user", "content": rag_prompt}],
-                    model="moonshotai/kimi-k2-instruct-0905",
-                    temperature=0.2,
+                    model="moonshotai/kimi-k2-instruct-0905", 
+                    temperature=0.1, # Temperatura baixa é CRUCIAL para não inventar citações
+                    max_tokens=1024
                 )
                 summary = chat_completion.choices[0].message.content.strip()
-                logging.info(f"Resumo RAG gerado com sucesso.")
+                logging.info(f"Resumo gerado com sucesso.")
             except Exception as e:
-                logging.warning(f"Falha ao gerar resumo RAG com a Groq: {e}")
-                summary = "O Oráculo está indisponível para gerar um resumo, mas aqui estão os documentos encontrados."
+                logging.warning(f"Falha ao gerar resumo com a Groq: {e}")
+                summary = "O Oráculo está indisponível, mas os documentos mais relevantes estão listados abaixo."
         
-        logging.info(f"Encontrados {len(results)} resultados. Resumo: {summary[:50]}...")
-        return {"summary": summary, "results": results}
+        return {"summary": summary, "results": final_results}
 
     except Exception as e:
         logging.error(f"ERRO CRÍTICO EM /search/: {e}", exc_info=True)
@@ -288,6 +357,8 @@ async def describe_clusters_api(job_id: str = Form(...)):
             cluster_embeddings = embeddings[mask]
             cluster_texts = df[mask]["full_text"].tolist()
             if len(cluster_texts) < 3: continue
+            
+            # Pega os documentos mais próximos do centróide do cluster
             centroid = np.mean(cluster_embeddings, axis=0)
             similarities = cosine_similarity([centroid], cluster_embeddings)[0]
             top_indices = np.argsort(similarities)[-3:][::-1]
@@ -307,11 +378,12 @@ async def describe_clusters_api(job_id: str = Form(...)):
             "\n\nResponda APENAS com o JSON."
         )
         
+        # Mantendo Llama para tarefa de JSON simples, pois é rápido e segue bem formato
         chat_completion = groq_client.chat.completions.create(
             messages=[
                 {"role": "system", "content": "Siga as instruções e responda apenas com um objeto JSON válido."},
                 {"role": "user", "content": master_prompt},
-            ], model="meta-llama/llama-4-maverick-17b-128e-instruct", temperature=0.2,
+            ], model="meta-llama/llama-3.3-70b-versatile", temperature=0.2, # Atualizei para o Llama 3.3 que é melhor
         )
         response_content = chat_completion.choices[0].message.content