Update app/memory.py

app/memory.py

@@ -1,82 +1,268 @@
-import csv
 import os
-import tempfile
-from datetime import datetime, timezone
-from pathlib import Path
 
+import numpy as np
+import pandas as pd
+import torch
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
 
-DEFAULT_LOGS_DIR = Path(os.getenv("LOGS_DIR", "/data/logs"))
-FALLBACK_LOGS_DIR = Path(tempfile.gettempdir()) / "tcc2_agent" / "logs"
+from .memory import caminho_memoria_negativa
+from .utils import (
+    bonus_lexical,
+    inferir_categoria_consulta,
+    limpar_texto,
+    mapear_categoria,
+)
 
 
-def _resolve_logs_dir():
-    preferred_dir = os.environ.get("LOGS_DIR", "").strip()
-    candidates = [Path(preferred_dir)] if preferred_dir else []
-    candidates.extend([DEFAULT_LOGS_DIR, FALLBACK_LOGS_DIR])
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(BASE_DIR, "data")
+PATH_PRODUCTS = os.path.join(DATA_DIR, "produtos_finetunado.csv")
+PATH_EMBEDDINGS = os.path.join(DATA_DIR, "embeddings_produtos_finetunado.npy")
+
+MODEL_NAME = os.getenv("HF_MODEL_REPO", "Ana2012/bertimbau-buscador").strip()
+HF_API_TOKEN = os.getenv("HF_API_TOKEN", "").strip()
+
+
+class SearchEngine:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = None
+        self.df_produtos = None
+        self.emb_produtos = None
+        self.df_negative_memory = pd.DataFrame()
+        self.negative_memory_mtime = None
+
+    def load(self):
+        self._load_products()
+        self._load_model()
+        self._load_embeddings()
+        self._refresh_negative_memory(force=True)
+
+    def _load_products(self):
+        df = pd.read_csv(PATH_PRODUCTS)
+        df.columns = df.columns.str.strip().str.lower()
+
+        df["product_name"] = df["product_name"].fillna("").astype(str)
+        df["description"] = df["description"].fillna("").astype(str)
+        df["categoria_principal"] = df["categoria_principal"].fillna("").astype(str)
+        df["category_names_text"] = df["category_names_text"].fillna("").astype(str)
+        df["region"] = df["region"].fillna("").astype(str)
+        df["neighborhood"] = df["neighborhood"].fillna("").astype(str)
+
+        df["product_name_limpo"] = df["product_name"].apply(limpar_texto)
+        df["description_limpa"] = df["description"].apply(limpar_texto)
+        df["categoria_principal_limpa"] = df["categoria_principal"].apply(limpar_texto)
+        df["category_names_text_limpo"] = df["category_names_text"].apply(limpar_texto)
+        df["region_limpa"] = df["region"].apply(limpar_texto)
+        df["neighborhood_limpo"] = df["neighborhood"].apply(limpar_texto)
+
+        df["texto_busca_reforcado"] = (
+            "produto " + df["product_name_limpo"] + " "
+            + "categoria " + df["categoria_principal_limpa"] + " "
+            + "categorias " + df["category_names_text_limpo"] + " "
+            + "bairro " + df["neighborhood_limpo"] + " "
+            + "regiao " + df["region_limpa"] + " "
+            + "descricao " + df["description_limpa"]
+        ).str.strip()
+
+        df["categoria_grupo"] = df["categoria_principal"].apply(mapear_categoria)
+
+        self.df_produtos = df
+
+    def _load_model(self):
+        kwargs = {"device": self.device}
+        if HF_API_TOKEN:
+            kwargs["token"] = HF_API_TOKEN
+
+        self.model = SentenceTransformer(MODEL_NAME, **kwargs)
+
+    def _load_embeddings(self):
+        self.emb_produtos = np.load(PATH_EMBEDDINGS)
+
+        if self.emb_produtos.ndim != 2:
+            raise RuntimeError("O arquivo de embeddings precisa conter uma matriz 2D.")
+
+    def runtime_info(self):
+        return {
+            "model_repo": MODEL_NAME,
+            "device": self.device,
+            "products_loaded": 0 if self.df_produtos is None else int(len(self.df_produtos)),
+            "embeddings_loaded": 0 if self.emb_produtos is None else int(len(self.emb_produtos)),
+            "embedding_dim": 0 if self.emb_produtos is None else int(self.emb_produtos.shape[1]),
+        }
+
+    def _refresh_negative_memory(self, force=False):
+        memory_path = caminho_memoria_negativa()
+
+        if not os.path.exists(memory_path):
+            self.df_negative_memory = pd.DataFrame()
+            self.negative_memory_mtime = None
+            return
+
+        current_mtime = os.path.getmtime(memory_path)
+        if not force and self.negative_memory_mtime == current_mtime:
+            return
 
-    for directory in candidates:
         try:
-            directory.mkdir(parents=True, exist_ok=True)
-            return directory
-        except OSError:
-            continue
-
-    raise OSError(
-        "Nao foi possivel criar um diretorio para armazenar a memoria negativa. "
-        "Defina LOGS_DIR para um caminho gravavel."
-    )
-
-
-def caminho_memoria_negativa():
-    return str(_resolve_logs_dir() / "negative_memory.csv")
-
-
-def inicializar_memoria_negativa():
-    memory_file = caminho_memoria_negativa()
-
-    if not os.path.exists(memory_file):
-        with open(memory_file, "w", newline="", encoding="utf-8") as f:
-            writer = csv.writer(f)
-            writer.writerow([
-                "timestamp",
-                "search_id",
-                "rank",
-                "query",
-                "product_id",
-                "product_name",
-                "rating",
-                "motivo",
-            ])
-
-
-def salvar_memoria_negativa(
-    query,
-    product_id,
-    product_name,
-    rating,
-    motivo="feedback_negativo",
-    search_id=None,
-    rank=None,
-):
-    inicializar_memoria_negativa()
-    memory_file = caminho_memoria_negativa()
-
-    row = [
-        datetime.now(timezone.utc).isoformat(),  # padrão UTC
-        search_id or "",
-        rank if rank is not None else "",
-        query or "",
-        product_id or "",
-        product_name or "",
-        rating if rating is not None else "",
-        motivo or "feedback_negativo",
-    ]
-
-    try:
-        with open(memory_file, "a", newline="", encoding="utf-8") as f:
-            writer = csv.writer(f)
-            writer.writerow(row)
-    except Exception as e:
-        print("Erro ao salvar memoria negativa:", e)
-
-    return {"status": "ok"}
+            df = pd.read_csv(memory_path)
+        except Exception:
+            self.df_negative_memory = pd.DataFrame()
+            self.negative_memory_mtime = current_mtime
+            return
+
+        df.columns = df.columns.str.strip().str.lower()
+
+        for col in ["query", "product_id", "product_name", "motivo", "rating"]:
+            if col not in df.columns:
+                df[col] = ""
+
+        df["query"] = df["query"].fillna("").astype(str)
+        df["query_limpa"] = df["query"].apply(limpar_texto)
+        df["product_id"] = df["product_id"].fillna("").astype(str)
+        df["product_name"] = df["product_name"].fillna("").astype(str)
+        df["motivo"] = df["motivo"].fillna("").astype(str)
+        df["rating_num"] = pd.to_numeric(df["rating"], errors="coerce")
+
+        self.df_negative_memory = df
+        self.negative_memory_mtime = current_mtime
+
+    def _similaridade_consulta(self, query_atual, query_memoria):
+        if not query_atual or not query_memoria:
+            return 0.0
+
+        if query_atual == query_memoria:
+            return 1.0
+
+        termos_atuais = set(query_atual.split())
+        termos_memoria = set(query_memoria.split())
+
+        if not termos_atuais or not termos_memoria:
+            return 0.0
+
+        intersecao = len(termos_atuais & termos_memoria)
+        if intersecao == 0:
+            return 0.0
+
+        return intersecao / max(len(termos_atuais), len(termos_memoria))
+
+    def _calcular_penalidade_feedback(self, query_text, df_filtrado):
+        self._refresh_negative_memory()
+
+        if self.df_negative_memory.empty:
+            return np.zeros(len(df_filtrado))
+
+        query_limpa = limpar_texto(query_text)
+        memorias = self.df_negative_memory[
+            self.df_negative_memory["product_id"].isin(df_filtrado["product_id"].astype(str))
+        ]
+
+        if memorias.empty:
+            return np.zeros(len(df_filtrado))
+
+        penalidades = {}
+
+        for _, memoria in memorias.iterrows():
+            similaridade = self._similaridade_consulta(query_limpa, memoria["query_limpa"])
+            if similaridade <= 0:
+                continue
+
+            penalidade = 0.08 + (0.12 * similaridade)
+
+            if memoria["motivo"] == "nao_foi_util":
+                penalidade += 0.04
+
+            if pd.notna(memoria["rating_num"]) and memoria["rating_num"] <= 2:
+                penalidade += 0.04
+
+            product_id = memoria["product_id"]
+            penalidades[product_id] = min(
+                penalidades.get(product_id, 0.0) + penalidade,
+                0.45,
+            )
+
+        return df_filtrado["product_id"].astype(str).map(
+            lambda x: penalidades.get(x, 0.0)
+        ).values
+
+    def gerar_embedding_unico(self, texto):
+        embedding = self.model.encode(
+            texto,
+            convert_to_numpy=True,
+            normalize_embeddings=False,
+            show_progress_bar=False,
+        )
+        return np.asarray(embedding, dtype=np.float32)
+
+    def buscar(self, query_text, top_k=5):
+        top_k = max(1, int(top_k or 5))
+
+        query_limpa = limpar_texto(query_text)
+        categoria = inferir_categoria_consulta(query_limpa)
+
+        if categoria is not None:
+            mask = self.df_produtos["categoria_grupo"] == categoria
+            df_filtrado = self.df_produtos[mask].copy()
+            idx_filtrado = df_filtrado.index.tolist()
+        else:
+            df_filtrado = self.df_produtos.copy()
+            idx_filtrado = df_filtrado.index.tolist()
+
+        if len(df_filtrado) == 0:
+            df_filtrado = self.df_produtos.copy()
+            idx_filtrado = df_filtrado.index.tolist()
+
+        emb_query = self.gerar_embedding_unico(query_text).reshape(1, -1)
+        emb_base = self.emb_produtos[idx_filtrado]
+
+        if emb_base.shape[1] != emb_query.shape[1]:
+            raise RuntimeError(
+                "Dimensao incompatível entre os embeddings salvos e o embedding da consulta. "
+                "Regenere o arquivo .npy com o mesmo modelo SentenceTransformer."
+            )
+
+        sims = cosine_similarity(emb_query, emb_base)[0]
+
+        bonus = df_filtrado.apply(
+            lambda row: bonus_lexical(
+                query_text,
+                row["product_name"],
+                row["categoria_principal"],
+                row["neighborhood"],
+                row["region"],
+                row["description"],
+                row["texto_busca_reforcado"],
+            ),
+            axis=1,
+        ).values
+
+        penalidade_feedback = self._calcular_penalidade_feedback(query_text, df_filtrado)
+        score_final = sims + bonus - penalidade_feedback
+
+        top_idx_local = np.argsort(score_final)[::-1][:top_k]
+
+        resultados = []
+        for rank, idx_local in enumerate(top_idx_local, start=1):
+            idx_global = idx_filtrado[idx_local]
+            prod = self.df_produtos.iloc[idx_global]
+
+            resultados.append({
+                "rank": rank,
+                "establishment_id": str(prod["establishment_id"]),
+                "product_id": str(prod["product_id"]),
+                "product_name": prod["product_name"],
+                "categoria_principal": prod["categoria_principal"],
+                "categoria_grupo": prod["categoria_grupo"],
+                "region": prod["region"],
+                "neighborhood": prod["neighborhood"],
+                "score_semantico": float(sims[idx_local]),
+                "bonus_lexical": float(bonus[idx_local]),
+                "penalidade_feedback": float(penalidade_feedback[idx_local]),
+                "score_final": float(score_final[idx_local]),
+            })
+
+        return {
+            "query": query_text,
+            "categoria_inferida": categoria,
+            "resultados": resultados,
+        }