""" utils.py — Funciones utilitarias reutilizables para el buscador de jurisprudencia. Contiene toda la lógica de negocio desacoplada de la interfaz: • Carga de datos (descarga HF, procesamiento de DataFrames, conexión ChromaDB) • Búsqueda semántica genérica • Renderizado de resultados en HTML (con título y campos extra configurables) """ from __future__ import annotations import os from typing import Any import chromadb import pandas as pd from huggingface_hub import snapshot_download # ───────────────────────────────────────────────────────────────────────────── # 1. CARGA DE DATOS # ───────────────────────────────────────────────────────────────────────────── def load_database_assets(db_config: dict) -> dict: """ Descarga el dataset desde Hugging Face, procesa los DataFrames de textos y metadata, y conecta con las colecciones de ChromaDB. Parámetros ---------- db_config : dict Diccionario de configuración de una base de datos individual (una entrada del diccionario DATABASES de config.py). Retorna ------- dict con las claves: - "metadata" : pd.DataFrame con la metadata indexada por ID compuesto. - "collections" : lista de objetos chromadb.Collection. - "field_map" : dict con el mapeo de campos internos → columnas reales. """ print(f" Descargando dataset: {db_config['repo_id']}...") dataset_path = snapshot_download( repo_id=db_config["repo_id"], repo_type=db_config.get("repo_type", "dataset"), ) # --- Cargar textos y metadata --- print(f" Cargando archivos: {db_config['pickle_file']}, {db_config['csv_file']}...") textos = pd.read_pickle(os.path.join(dataset_path, db_config["pickle_file"])) metadata = pd.read_csv(os.path.join(dataset_path, db_config["csv_file"])) # Eliminar columnas innecesarias (si las hay) drop_cols = db_config.get("drop_columns", []) existing_drop_cols = [c for c in drop_cols if c in metadata.columns] if existing_drop_cols: metadata = metadata.drop(columns=existing_drop_cols) # Merge de textos con metadata merge_key = db_config.get("merge_key", "ID") metadata = pd.merge(textos, metadata, on=merge_key, how="inner") # Crear ID compuesto: ID_Page id_col = db_config.get("id_column", "ID") page_col = db_config.get("page_column", "Page") metadata[id_col] = metadata[id_col].astype(str) + "_" + textos[page_col].astype(str) # Indexar para búsqueda rápida metadata.set_index(id_col, inplace=True) # --- Conectar a ChromaDB --- chroma_dir = os.path.join(dataset_path, db_config["chroma_subdir"]) print(f" Conectando a ChromaDB en: {chroma_dir}...") client = chromadb.PersistentClient(path=chroma_dir) collection_names = db_config.get("collection_names", []) collections = [client.get_collection(name) for name in collection_names] print(f" ✓ Base de datos '{db_config.get('display_name', db_config['repo_id'])}' lista.") return { "metadata": metadata, "collections": collections, "field_map": db_config.get("field_map", {}), } # ───────────────────────────────────────────────────────────────────────────── # 2. BÚSQUEDA SEMÁNTICA # ───────────────────────────────────────────────────────────────────────────── def _safe_get(row: Any, column: str | None, default: str = "N/A") -> Any: """ Obtiene de forma segura el valor de una columna en una fila de DataFrame, manejando tanto Series como filas individuales. Retorna *default* si column es None o no existe. """ if column is None: return default try: if hasattr(row, "get"): return row.get(column, default) return row[column] except (KeyError, TypeError): return default def perform_search( query_embedding: list[float], collections: list, metadata_df: pd.DataFrame, field_map: dict, n_results: int = 15, ) -> list[dict]: """ Ejecuta la búsqueda semántica en las colecciones de ChromaDB y enriquece los resultados con la metadata del DataFrame. Parámetros ---------- query_embedding : list[float] Vector de embedding de la consulta del usuario. collections : list Lista de colecciones de ChromaDB donde buscar. metadata_df : pd.DataFrame DataFrame de metadata indexado por ID compuesto. field_map : dict Mapeo de campos internos estandarizados → nombres reales de columnas. Claves esperadas: "enlace" → columna con el link "text" → columna con el texto principal "date" → columna con la fecha "page" → columna con la página "titulo" → columna con el título (str o None) "extra_fields" → lista de tuplas (etiqueta, columna) n_results : int Número máximo de resultados a retornar. Retorna ------- list[dict] Lista de diccionarios con los resultados ordenados por similitud descendente. Cada diccionario contiene las claves estandarizadas: id, similitud, enlace, text, date, page, titulo, extras. """ # Nombres reales de las columnas según el field_map col_enlace = field_map.get("enlace", "Enlace") col_text = field_map.get("text", "Text") col_date = field_map.get("date", "Date") col_page = field_map.get("page", "Page") col_titulo = field_map.get("titulo") # puede ser None extra_fields = field_map.get("extra_fields", []) # lista de (etiqueta, columna) all_results: list[dict] = [] for collection in collections: results = collection.query( query_embeddings=[query_embedding], n_results=n_results, include=["distances"], ) cosine_similarities = [1 - dist for dist in results["distances"][0]] for i, chroma_id in enumerate(results["ids"][0]): try: if chroma_id in metadata_df.index: row = metadata_df.loc[chroma_id] # Campos base res: dict[str, Any] = { "id": chroma_id, "similitud": cosine_similarities[i], "enlace": _safe_get(row, col_enlace, "#"), "text": _safe_get(row, col_text, "N/A"), "date": _safe_get(row, col_date, "N/A"), "page": _safe_get(row, col_page, "N/A"), } # Título (puede no existir) res["titulo"] = _safe_get(row, col_titulo, None) # Campos extra configurables: lista de (etiqueta, valor) extras: list[tuple[str, str]] = [] for label, col_name in extra_fields: value = str(_safe_get(row, col_name, "N/A")) extras.append((label, value)) res["extras"] = extras all_results.append(res) else: print(f" Warning: ID {chroma_id} no encontrado en metadata.") except Exception as e: print(f" Error al recuperar metadata para ID {chroma_id}: {e}") # Ordenar por similitud descendente y limitar resultados all_results = sorted( all_results, key=lambda x: x["similitud"], reverse=True )[:n_results] return all_results # ───────────────────────────────────────────────────────────────────────────── # 3. RENDERIZADO HTML # ───────────────────────────────────────────────────────────────────────────── def render_results_html(search_results: list[dict]) -> str: """ Genera el HTML de las tarjetas (cards) de resultados a partir de una lista de diccionarios estandarizados. Estructura de cada card: ┌─ card-header ──────────────────────────────────────────────┐ │ Página · Fecha · Relevancia │ ├─ card-title (solo si hay título) ──────────────────────────┤ │ 📌 Título de la decisión │ ├─ card-extras (solo si hay extra_fields) ──────────────────┤ │ Etiqueta1: valor1 | Etiqueta2: valor2 | ... │ ├─ card-body ────────────────────────────────────────────────┤ │ Texto principal del fragmento │ ├─ card-footer ──────────────────────────────────────────────┤ │ 🔗 Ver Documento Completo │ └────────────────────────────────────────────────────────────┘ Parámetros ---------- search_results : list[dict] Lista de resultados proveniente de `perform_search`. Retorna ------- str Cadena HTML con las tarjetas de resultados, lista para inyectar en un componente gr.HTML de Gradio. """ if not search_results: return ( '
' "No se encontraron resultados relevantes.
" ) html = '
' for item in search_results: similitud_pct = f"{item['similitud'] * 100:.1f}%" enlace = item.get("enlace", "#") fecha = item.get("date", "N/A") texto = str(item.get("text", "N/A")).replace("\n", " ") pagina = item.get("page", "N/A") titulo = item.get("titulo") extras = item.get("extras", []) # ── Sección: título (condicional) ── titulo_html = "" if titulo: titulo_safe = str(titulo).replace("\n", " ") titulo_html = f"""
📌 {titulo_safe}
""" # ── Sección: campos extra concatenados (condicional) ── extras_html = "" if extras: items_html = " | ".join( f"{label}: {value}" for label, value in extras ) extras_html = f"""
{items_html}
""" # ── Card completa ── html += f"""
⚖️ Página: {pagina} 📅 {fecha} 🎯 Relevancia: {similitud_pct}
{titulo_html}{extras_html}

{texto}

🔗 Ver Documento Completo
""" html += "
" return html # ───────────────────────────────────────────────────────────────────────────── # 4. CSS (lectura desde archivo estático) # ───────────────────────────────────────────────────────────────────────────── def load_css(css_path: str | None = None) -> str: """ Lee y retorna el contenido del archivo CSS. Busca el archivo en el siguiente orden de prioridad: 1. La ruta explícita proporcionada en `css_path`. 2. `static/custom.css` (relativo al directorio de este módulo). 3. `custom.css` en la raíz del proyecto (mismo directorio que este módulo). Parámetros ---------- css_path : str | None Ruta absoluta o relativa al archivo CSS. Si es None, se busca automáticamente en las ubicaciones por defecto. Retorna ------- str Contenido del archivo CSS como cadena de texto. Raises ------ FileNotFoundError Si no se encuentra el archivo CSS en ninguna ubicación. """ base_dir = os.path.dirname(os.path.abspath(__file__)) # Lista de rutas candidatas en orden de prioridad if css_path is not None: candidates = [css_path] else: candidates = [ os.path.join(base_dir, "static", "custom.css"), os.path.join(base_dir, "custom.css"), ] for path in candidates: if os.path.isfile(path): with open(path, "r", encoding="utf-8") as f: return f.read() searched = ", ".join(candidates) raise FileNotFoundError( f"No se encontró custom.css en ninguna de estas ubicaciones: {searched}" )