CoolFace
Apppublic

stvnnnnnn/nl2sql-backend-t5

sourceHugging Facemitupdated 10mo agoView on Hugging Face
1likes
app.py1397 linesDownload Raw Back to root
1import os2import io3import zipfile4import re5import difflib6import tempfile7import uuid8from typing import List, Optional, Dict, Any9 10from fastapi import FastAPI, UploadFile, File, HTTPException, Form, Header11from fastapi.middleware.cors import CORSMiddleware12from pydantic import BaseModel13 14import torch15from transformers import AutoTokenizer, AutoModelForSeq2SeqLM16from langdetect import detect17from transformers import MarianMTModel, MarianTokenizer18from openai import OpenAI19 20# ---- Postgres ----21import psycopg222from psycopg2 import sql as pgsql23 24# ---- Supabase ----25from supabase import create_client, Client26 27SUPABASE_URL = "https://bnvmqgjawtaslczewqyd.supabase.co"28SUPABASE_ANON_KEY = (29    "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6ImJudm1x"30    "Z2phd3Rhc2xjemV3cXlkIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NjQ0NjM5NDAsImV4cCI6MjA4"31    "MDAzOTk0MH0.9zkyqrsm-QOSwMTUPZEWqyFeNpbbuar01rB7pmObkUI"32)33 34supabase: Client = create_client(SUPABASE_URL, SUPABASE_ANON_KEY)35 36# ======================================================37# 0) Configuración general de paths / modelo / OpenAI38# ======================================================39 40MODEL_DIR = os.getenv("MODEL_DIR", "stvnnnnnn/t5-large-nl2sql-spider")41DEVICE = torch.device("cpu")42 43OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")44openai_client = OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None45 46# DSN de Supabase Postgres – EJEMPLO:47# postgresql://postgres:TU_PASSWORD@db.xxx.supabase.co:5432/postgres48POSTGRES_DSN = os.getenv("POSTGRES_DSN")49 50if not POSTGRES_DSN:51    raise RuntimeError(52        "⚠️ POSTGRES_DSN no está definido. "53        "Configúralo en los secrets del Space con la cadena de conexión de Supabase."54    )55 56# ======================================================57# 1) Gestor de conexiones dinámicas: Postgres (Neon)58# ======================================================59 60 61class PostgresManager:62    """63    Cada upload crea un *schema* aislado en Neon.64    connections[connection_id] = {65        "label": str,        # nombre de archivo original66        "engine": "postgres",67        "schema": str        # nombre del schema en Neon68    }69    """70 71    def __init__(self, dsn: str):72        self.dsn = dsn73        self.connections: Dict[str, Dict[str, Any]] = {}74 75    # ---------- utilidades internas ----------76 77    def _new_connection_id(self) -> str:78        return f"db_{uuid.uuid4().hex[:8]}"79 80    def _get_info(self, connection_id: str) -> Dict[str, Any]:81        if connection_id not in self.connections:82            raise KeyError(f"connection_id '{connection_id}' no registrado")83        return self.connections[connection_id]84 85    def _get_conn(self, autocommit: bool = True):86        conn = psycopg2.connect(self.dsn)87        conn.autocommit = autocommit88        return conn89 90    # ---------- helpers de sanitización de dumps ----------91 92    def _rewrite_line_for_schema(self, line: str, schema_name: str) -> str:93        """94        Versión simplificada:95        - Solo elimina líneas que modifican el search_path.96        - NO reescribe public./pagila. → dejamos que el dump use su propio schema.97        """98        if "search_path" in line.lower():99            return ""100        return line101 102    def _should_skip_statement(self, stmt: str) -> bool:103        """104        Devuelve True si el statement NO debe ejecutarse (grants, owner, create db, domains, etc.).105        Filtro universal para dumps PostgreSQL (Neon, Pagila, etc.).106        """107        if not stmt:108            return True109 110        upper = stmt.upper().strip()111 112        # 1) Statements globales / de administración que SIEMPRE ignoramos113        skip_prefixes = (114            "SET ",115            "RESET ",116            "SELECT PG_CATALOG.SET_CONFIG",117            "COMMENT ON EXTENSION",118            "COMMENT ON SCHEMA",119            "COMMENT ON DATABASE",120            "COMMENT ON COLLATION",121            "COMMENT ON CONVERSION",122            "COMMENT ON LANGUAGE",123            "COMMENT ON TEXT SEARCH",124            "COMMENT ON FOREIGN",125            "CREATE DATABASE",126            "ALTER DATABASE",127            "DROP DATABASE",128            "CREATE EXTENSION",129            "ALTER EXTENSION",130            "DROP EXTENSION",131            "REVOKE ",132            "GRANT ",133            "ALTER ROLE",134            "CREATE ROLE",135            "DROP ROLE",136            "CREATE USER",137            "ALTER USER",138            "DROP USER",139            "ALTER DEFAULT PRIVILEGES",140            "SECURITY LABEL",141            "BEGIN",142            "COMMIT",143            "ROLLBACK",144        )145        if upper.startswith(skip_prefixes):146            return True147 148        # 2) Cualquier cosa que toque OWNER / AUTHORIZATION la ignoramos149        owner_markers = (150            " OWNER TO ",151            " OWNER ",152            "AUTHORIZATION POSTGRES",153            "AUTHORIZATION PUBLIC",154            "AUTHORIZATION CURRENT_USER",155            "AUTHORIZATION \"POSTGRES\"",156        )157        if any(marker in upper for marker in owner_markers):158            return True159 160        # 3) Grants / revokes explícitos a postgres o public (aunque no empiecen por GRANT/REVOKE)161        if " TO POSTGRES" in upper or " FROM POSTGRES" in upper:162            return True163        if " TO PUBLIC" in upper or " FROM PUBLIC" in upper:164            return True165 166        return False167 168    def _execute_sanitized_pg_dump(169        self, cur, sql_text: str, schema_name: str170    ) -> None:171        """172        Ejecuta un dump de PostgreSQL dentro de un schema de sesión,173        aplicando sanitización y soportando COPY ... FROM stdin;.174 175        - Reescribe public./pagila. -> schema_name.176        - Respeta funciones con $$...$$ (no corta por ';' internos).177        - Ignora statements peligrosos via _should_skip_statement().178        """179 180        in_copy = False181        copy_sql = ""182        copy_lines: list[str] = []183 184        buffer = ""       # statement acumulado185        in_dollar = False # estamos dentro de $$...$$ ?186        dollar_tag = ""   # por ej. "$func$"187 188        in_domain_block = False  # 👈 estamos dentro de un bloque CREATE DOMAIN ?189        in_function_block = False  # 👈 estamos dentro de un CREATE FUNCTION ?190 191        def flush_statement():192            nonlocal buffer193            stmt = buffer.strip()194            buffer = ""195            if not stmt:196                return197            if self._should_skip_statement(stmt):198                return199            try:200                cur.execute(stmt)201            except Exception as e:202                msg = str(e).lower()203                # Ignoramos errores típicos de dumps que no son fatales204                if "already exists" in msg or "duplicate key value" in msg:205                    print("[WARN] Ignorando error no crítico:", e)206                    return207                raise208 209        # Procesar línea por línea210        for raw_line in sql_text.splitlines():211            line = raw_line.rstrip("\n")212            stripped = line.strip()213 214            # ====== BLOQUE CREATE FUNCTION (lo ignoramos entero) ======215            if in_function_block:216                # Cerramos cuando vemos algo tipo "$_$;" o "$func$;"217                if re.search(r"\$[A-Za-z0-9_]*\$;", stripped):218                    in_function_block = False219                continue220 221            # Comentarios y líneas vacías (fuera de COPY / DOMAIN / FUNCTION)222            if not in_copy and not in_domain_block:223                if not stripped or stripped.startswith("--"):224                    continue225 226                upper_line = stripped.upper()227                if (228                    upper_line.startswith("CREATE FUNCTION")229                    or upper_line.startswith("CREATE OR REPLACE FUNCTION")230                    or upper_line.startswith("ALTER FUNCTION")231                ):232                    # Ignoramos toda la función (cabecera + cuerpo)233                    in_function_block = True234                    continue235 236            # ====== BLOQUE COPY ... FROM stdin ======237            if in_copy:238                if stripped == r"\.":239                    # fin de COPY240                    data = "\n".join(copy_lines) + "\n"241                    cur.copy_expert(copy_sql, io.StringIO(data))242                    in_copy = False243                    copy_sql = ""244                    copy_lines.clear()245                else:246                    copy_lines.append(line)247                continue248 249            # Reescribimos la línea según el schema de sesión250            line = self._rewrite_line_for_schema(line, schema_name)251            stripped = line.strip()252            if not stripped:253                continue254 255            # Detectar inicio de COPY ahora que la línea ya está reescrita256            if stripped.upper().startswith("COPY ") and "FROM stdin" in stripped.upper():257                # Ejecutar lo que haya pendiente antes del COPY258                flush_statement()259                in_copy = True260                copy_sql = stripped  # ya reescrita261                copy_lines = []262                continue263 264            # Escanear la línea caracter a caracter para detectar $tag$ y ';'265            i = 0266            start_seg = 0267            length = len(line)268 269            while i < length:270                ch = line[i]271 272                # Manejo de delimitadores $tag$273                if ch == "$":274                    # ¿Inicio o fin de bloque dollar-quoted?275                    j = i + 1276                    while j < length and (line[j].isalnum() or line[j] == "_"):277                        j += 1278                    if j < length and line[j] == "$":279                        tag = line[i : j + 1]  # incluye ambos '$'280                        if not in_dollar:281                            in_dollar = True282                            dollar_tag = tag283                        else:284                            if tag == dollar_tag:285                                in_dollar = False286                                dollar_tag = ""287                        i = j + 1288                        continue289 290                # Fin de statement: ';' fuera de bloque dollar-quoted291                if ch == ";" and not in_dollar:292                    segment = line[start_seg : i + 1]293                    buffer += segment + "\n"294                    flush_statement()295                    start_seg = i + 1296                    i += 1297                    continue298 299                i += 1300 301            # Resto de la línea (después del último ';' o toda la línea si no hubo ';')302            if start_seg < length:303                buffer += line[start_seg:] + "\n"304 305        # Ejecutar lo que quede pendiente306        flush_statement()307 308        # Por seguridad, aseguramos que no haya COPY abierto sin cerrar309        if in_copy:310            raise RuntimeError("Dump SQL inválido: COPY sin terminación '\\.'")311 312    # ---------- creación de BD desde dump ----------313 314    def create_database_from_dump(self, label: str, sql_text: str) -> str:315        """316        Restaura un dump de Postgres (schema + datos) en la BD Neon.317        NO crea schemas de sesión: deja que el dump use sus propios schemas318        (public, pagila, etc.). Luego detecta el schema con más tablas.319        """320        connection_id = self._new_connection_id()321        schema_name: str | None = None322 323        conn = self._get_conn()324        try:325            with conn.cursor() as cur:326                # 1) Ejecutar el dump tal cual (solo limpiamos search_path)327                self._execute_sanitized_pg_dump(cur, sql_text, schema_name="public")328 329                # 2) Detectar el schema REAL donde quedaron las tablas del dump330                cur.execute(331                    """332                    SELECT table_schema, COUNT(*) AS n333                    FROM information_schema.tables334                    WHERE table_type = 'BASE TABLE'335                      AND table_schema NOT IN ('pg_catalog','information_schema')336                    GROUP BY table_schema337                    ORDER BY n DESC;338                    """339                )340                rows = cur.fetchall()341                if not rows:342                    raise RuntimeError(343                        "El dump se ejecutó pero no se encontraron tablas de usuario."344                    )345 346                # Tomamos el schema con más tablas (pagila, public, etc.)347                schema_name = rows[0][0]348 349        except Exception as e:350            conn.close()351            raise RuntimeError(f"Error ejecutando dump SQL en Postgres: {e}")352        finally:353            conn.close()354 355        self.connections[connection_id] = {356            "label": label,357            "engine": "postgres",358            "schema": schema_name,   # 👈 ahora es el schema REAL con tablas359        }360        return connection_id361 362    # ---------- ejecución segura de SQL ----------363 364    def execute_sql(self, connection_id: str, sql_text: str) -> Dict[str, Any]:365        """366        Ejecuta un SELECT dentro del schema asociado al connection_id.367        Bloquea operaciones destructivas por seguridad.368        """369        info = self._get_info(connection_id)370        schema = info["schema"]371 372        forbidden = ["drop ", "delete ", "update ", "insert ", "alter ", "replace "]373        sql_low = sql_text.lower()374        if any(tok in sql_low for tok in forbidden):375            return {376                "ok": False,377                "error": "Query bloqueada por seguridad (operación destructiva).",378                "rows": None,379                "columns": [],380            }381 382        conn = self._get_conn()383        try:384            with conn.cursor() as cur:385                # usar el schema de la sesión386                cur.execute(387                    pgsql.SQL("SET search_path TO {}").format(388                        pgsql.Identifier(schema)389                    )390                )391                cur.execute(sql_text)392 393                if cur.description:394                    rows = cur.fetchall()395                    cols = [d[0] for d in cur.description]396                else:397                    rows, cols = [], []398 399            return {400                "ok": True,401                "error": None,402                "rows": [list(r) for r in rows],403                "columns": cols,404            }405        except Exception as e:406            return {"ok": False, "error": str(e), "rows": None, "columns": []}407        finally:408            conn.close()409 410    # ---------- introspección de esquema ----------411 412    def get_schema(self, connection_id: str) -> Dict[str, Any]:413        info = self._get_info(connection_id)414        schema = info["schema"]  # schema "ideal" que registramos415 416        conn = self._get_conn()417        try:418            tables_info: Dict[str, Dict[str, Any]] = {}419            foreign_keys: List[Dict[str, Any]] = []420 421            with conn.cursor() as cur:422                # 1) Intentamos solo con el schema registrado423                cur.execute(424                    """425                    SELECT table_name426                    FROM information_schema.tables427                    WHERE table_schema = %s428                      AND table_type = 'BASE TABLE'429                    ORDER BY table_name;430                    """,431                    (schema,),432                )433                tables = [r[0] for r in cur.fetchall()]434 435                # 2) 🔁 Fallback: si no hay tablas en ese schema,436                #    buscamos en TODOS los schemas de usuario437                if not tables:438                    cur.execute(439                        """440                        SELECT table_schema, table_name441                        FROM information_schema.tables442                        WHERE table_type = 'BASE TABLE'443                          AND table_schema NOT IN ('pg_catalog','information_schema')444                        ORDER BY table_schema, table_name;445                        """446                    )447                    rows = cur.fetchall()448 449                    if not rows:450                        # No hay tablas en ningún schema de usuario451                        return {452                            "tables": {},453                            "foreign_keys": [],454                        }455 456                    # Schemas candidatos que sí tienen tablas457                    schemas = sorted({s for (s, _) in rows})458 459                    # Preferimos:460                    # 1) el schema ya registrado (si por alguna razón tiene tablas)461                    # 2) 'pagila'462                    # 3) 'public'463                    # 4) el primero que aparezca464                    target_schema = None465                    if schema in schemas:466                        target_schema = schema467                    elif "pagila" in schemas:468                        target_schema = "pagila"469                    elif "public" in schemas:470                        target_schema = "public"471                    else:472                        target_schema = schemas[0]473 474                    print(475                        f"[WARN] Schema '{schema}' sin tablas; usando schema real '{target_schema}'"476                    )477 478                    # Actualizamos el schema asociado a esta conexión479                    schema = target_schema480                    info["schema"] = schema481 482                    tables = [t for (s, t) in rows if s == schema]483 484                # 3) Columnas por tabla del schema final seleccionado485                for t in tables:486                    cur.execute(487                        """488                        SELECT column_name489                        FROM information_schema.columns490                        WHERE table_schema = %s491                          AND table_name = %s492                        ORDER BY ordinal_position;493                        """,494                        (schema, t),495                    )496                    cols = [r[0] for r in cur.fetchall()]497                    tables_info[t] = {"columns": cols}498 499                # 4) Foreign keys del schema final500                cur.execute(501                    """502                    SELECT503                        tc.table_name      AS from_table,504                        kcu.column_name    AS from_column,505                        ccu.table_name     AS to_table,506                        ccu.column_name    AS to_column507                    FROM information_schema.table_constraints AS tc508                    JOIN information_schema.key_column_usage AS kcu509                      ON tc.constraint_name = kcu.constraint_name510                     AND tc.table_schema   = kcu.table_schema511                    JOIN information_schema.constraint_column_usage AS ccu512                      ON ccu.constraint_name = tc.constraint_name513                     AND ccu.table_schema   = tc.table_schema514                    WHERE tc.constraint_type = 'FOREIGN KEY'515                      AND tc.table_schema    = %s;516                    """,517                    (schema,),518                )519                for ft, fc, tt, tc2 in cur.fetchall():520                    foreign_keys.append(521                        {522                            "from_table": ft,523                            "from_column": fc,524                            "to_table": tt,525                            "to_column": tc2,526                        }527                    )528 529            return {530                "tables": tables_info,531                "foreign_keys": foreign_keys,532            }533        finally:534            conn.close()535 536    # ---------- preview de tabla ----------537 538    def get_preview(539        self, connection_id: str, table: str, limit: int = 20540    ) -> Dict[str, Any]:541        info = self._get_info(connection_id)542        schema = info["schema"]543 544        conn = self._get_conn()545        try:546            with conn.cursor() as cur:547                cur.execute(548                    pgsql.SQL("SET search_path TO {}").format(549                        pgsql.Identifier(schema)550                    )551                )552                query = pgsql.SQL("SELECT * FROM {} LIMIT %s").format(553                    pgsql.Identifier(table)554                )555                cur.execute(query, (int(limit),))556                rows = cur.fetchall()557                cols = [d[0] for d in cur.description] if cur.description else []558 559            return {560                "columns": cols,561                "rows": [list(r) for r in rows],562            }563        finally:564            conn.close()565 566 567# Instancia global de PostgresManager568sql_manager = PostgresManager(POSTGRES_DSN)569 570# ======================================================571# 2) Inicialización de FastAPI572# ======================================================573 574app = FastAPI(575    title="NL2SQL Backend",576    version="3.0.0",577)578 579app.add_middleware(580    CORSMiddleware,581    allow_origins=["*"],582    allow_methods=["*"],583    allow_headers=["*"],584)585 586# ======================================================587# 3) Modelo NL→SQL y traductor ES→EN588# ======================================================589 590t5_tokenizer = None591t5_model = None592mt_tokenizer = None593mt_model = None594 595 596def load_nl2sql_model():597    """Carga el modelo NL→SQL (T5-large fine-tuned en Spider) desde HF Hub."""598    global t5_tokenizer, t5_model599    if t5_model is not None:600        return601    print(f"🔁 Cargando modelo NL→SQL desde: {MODEL_DIR}")602    t5_tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=True)603    t5_model = AutoModelForSeq2SeqLM.from_pretrained(604        MODEL_DIR, torch_dtype=torch.float32605    )606    t5_model.to(DEVICE)607    t5_model.eval()608    print("✅ Modelo NL→SQL listo en memoria.")609 610 611def load_es_en_translator():612    """Carga el modelo Helsinki-NLP para traducción ES→EN (solo una vez)."""613    global mt_tokenizer, mt_model614    if mt_model is not None:615        return616    model_name = "Helsinki-NLP/opus-mt-es-en"617    print(f"🔁 Cargando traductor ES→EN: {model_name}")618    mt_tokenizer = MarianTokenizer.from_pretrained(model_name)619    mt_model = MarianMTModel.from_pretrained(model_name)620    mt_model.to(DEVICE)621    mt_model.eval()622    print("✅ Traductor ES→EN listo.")623 624 625def detect_language(text: str) -> str:626    try:627        return detect(text)628    except Exception:629        return "unknown"630 631 632def translate_es_to_en(text: str) -> str:633    """634    Usa Marian ES→EN solo si el texto se detecta como español ('es').635    Si no, devuelve el texto tal cual.636    """637    lang = detect_language(text)638    if lang != "es":639        return text640    if mt_model is None:641        load_es_en_translator()642    inputs = mt_tokenizer(text, return_tensors="pt", truncation=True).to(DEVICE)643    with torch.no_grad():644        out = mt_model.generate(**inputs, max_length=256)645    return mt_tokenizer.decode(out[0], skip_special_tokens=True)646 647 648# ======================================================649# 4) Capa de reparación de SQL (usa el schema real)650# ======================================================651 652 653def _normalize_name_for_match(name: str) -> str:654    s = name.lower()655    s = s.replace('"', "").replace("`", "")656    s = s.replace("_", "")657    if s.endswith("s") and len(s) > 3:658        s = s[:-1]659    return s660 661 662def _build_schema_indexes(663    tables_info: Dict[str, Dict[str, List[str]]]664) -> Dict[str, Dict[str, List[str]]]:665    table_index: Dict[str, List[str]] = {}666    column_index: Dict[str, List[str]] = {}667 668    for t, info in tables_info.items():669        tn = _normalize_name_for_match(t)670        table_index.setdefault(tn, [])671        if t not in table_index[tn]:672            table_index[tn].append(t)673 674        for c in info.get("columns", []):675            cn = _normalize_name_for_match(c)676            column_index.setdefault(cn, [])677            if c not in column_index[cn]:678                column_index[cn].append(c)679 680    return {"table_index": table_index, "column_index": column_index}681 682 683def _best_match_name(missing: str, index: Dict[str, List[str]]) -> Optional[str]:684    if not index:685        return None686 687    key = _normalize_name_for_match(missing)688    if key in index and index[key]:689        return index[key][0]690 691    candidates = difflib.get_close_matches(key, list(index.keys()), n=1, cutoff=0.7)692    if not candidates:693        return None694    best_key = candidates[0]695    if index[best_key]:696        return index[best_key][0]697    return None698 699 700DOMAIN_SYNONYMS_TABLE = {701    "song": "track",702    "songs": "track",703    "tracks": "track",704    "artist": "artist",705    "artists": "artist",706    "album": "album",707    "albums": "album",708    "order": "invoice",709    "orders": "invoice",710}711 712DOMAIN_SYNONYMS_COLUMN = {713    "song": "name",714    "songs": "name",715    "track": "name",716    "title": "name",717    "length": "milliseconds",718    "duration": "milliseconds",719}720 721 722def try_repair_sql(723    sql: str, error: str, schema_meta: Dict[str, Any]724) -> Optional[str]:725    """726    Intenta reparar nombres de tablas/columnas basándose en el esquema real.727    Compatible con mensajes de Postgres y también con los de SQLite728    (por si algún día reusamos la lógica).729    """730    tables_info = schema_meta["tables"]731    idx = _build_schema_indexes(tables_info)732    table_index = idx["table_index"]733    column_index = idx["column_index"]734 735    repaired_sql = sql736    changed = False737 738    missing_table = None739    missing_column = None740 741    m_t = re.search(r'relation "([\w\.]+)" does not exist', error, re.IGNORECASE)742    if not m_t:743        m_t = re.search(r"no such table: ([\w\.]+)", error)744    if m_t:745        missing_table = m_t.group(1)746 747    m_c = re.search(r'column "([\w\.]+)" does not exist', error, re.IGNORECASE)748    if not m_c:749        m_c = re.search(r"no such column: ([\w\.]+)", error)750    if m_c:751        missing_column = m_c.group(1)752 753    if missing_table:754        short = missing_table.split(".")[-1]755        syn = DOMAIN_SYNONYMS_TABLE.get(short.lower())756        target = None757        if syn:758            target = _best_match_name(syn, table_index) or syn759        if not target:760            target = _best_match_name(short, table_index)761 762        if target:763            pattern = r"\b" + re.escape(short) + r"\b"764            new_sql = re.sub(pattern, target, repaired_sql)765            if new_sql != repaired_sql:766                repaired_sql = new_sql767                changed = True768 769    if missing_column:770        short = missing_column.split(".")[-1]771        syn = DOMAIN_SYNONYMS_COLUMN.get(short.lower())772        target = None773        if syn:774            target = _best_match_name(syn, column_index) or syn775        if not target:776            target = _best_match_name(short, column_index)777 778        if target:779            pattern = r"\b" + re.escape(short) + r"\b"780            new_sql = re.sub(pattern, target, repaired_sql)781            if new_sql != repaired_sql:782                repaired_sql = new_sql783                changed = True784 785    if not changed:786        return None787    return repaired_sql788 789 790# ======================================================791# 5) Prompt NL→SQL + re-ranking792# ======================================================793 794 795def build_prompt(question_en: str, db_id: str, schema_str: str) -> str:796    return (797        f"translate to SQL: {question_en} | "798        f"db: {db_id} | schema: {schema_str} | "799        f"note: use JOIN when foreign keys link tables"800    )801 802 803def normalize_score(raw: float) -> float:804    """Normaliza el score logit del modelo a un porcentaje 0-100."""805    norm = (raw + 20) / 25806    norm = max(0, min(1, norm))807    return round(norm * 100, 2)808 809 810def nl2sql_with_rerank(question: str, conn_id: str) -> Dict[str, Any]:811    if conn_id not in sql_manager.connections:812        raise HTTPException(813            status_code=404, detail=f"connection_id '{conn_id}' no registrado"814        )815 816    meta = sql_manager.get_schema(conn_id)817    tables_info = meta["tables"]818 819    parts = []820    for t, info in tables_info.items():821        cols = info.get("columns", [])822        parts.append(f"{t}(" + ", ".join(cols) + ")")823    schema_str = " ; ".join(parts) if parts else "(empty_schema)"824 825    detected = detect_language(question)826    question_en = translate_es_to_en(question) if detected == "es" else question827 828    prompt = build_prompt(question_en, db_id=conn_id, schema_str=schema_str)829 830    if t5_model is None:831        load_nl2sql_model()832 833    inputs = t5_tokenizer(834        [prompt], return_tensors="pt", truncation=True, max_length=768835    ).to(DEVICE)836    num_beams = 6837    num_return = 6838 839    with torch.no_grad():840        out = t5_model.generate(841            **inputs,842            max_length=220,843            num_beams=num_beams,844            num_return_sequences=num_return,845            return_dict_in_generate=True,846            output_scores=True,847        )848 849    sequences = out.sequences850    scores = out.sequences_scores851    if scores is not None:852        scores = scores.cpu().tolist()853    else:854        scores = [0.0] * sequences.size(0)855 856    candidates: List[Dict[str, Any]] = []857    best = None858    best_exec = False859    best_score = -1e9860 861    for i in range(sequences.size(0)):862        raw_sql = t5_tokenizer.decode(863            sequences[i], skip_special_tokens=True864        ).strip()865        cand: Dict[str, Any] = {866            "sql": raw_sql,867            "score": float(scores[i]),868            "repaired_from": None,869            "repair_note": None,870            "raw_sql_model": raw_sql,871        }872 873        exec_info = sql_manager.execute_sql(conn_id, raw_sql)874 875        err_lower = (exec_info["error"] or "").lower()876        if (not exec_info["ok"]) and (877            "no such table" in err_lower878            or "no such column" in err_lower879            or "does not exist" in err_lower880        ):881            current_sql = raw_sql882            last_error = exec_info["error"] or ""883            for step in range(1, 4):884                repaired_sql = try_repair_sql(current_sql, last_error, meta)885                if not repaired_sql or repaired_sql == current_sql:886                    break887                exec_info2 = sql_manager.execute_sql(conn_id, repaired_sql)888                cand["repaired_from"] = (889                    current_sql890                    if cand["repaired_from"] is None891                    else cand["repaired_from"]892                )893                cand["repair_note"] = (894                    f"auto-repair (table/column name, step {step})"895                )896                cand["sql"] = repaired_sql897                exec_info = exec_info2898                current_sql = repaired_sql899                if exec_info2["ok"]:900                    break901                last_error = exec_info2["error"] or ""902 903        cand["exec_ok"] = exec_info["ok"]904        cand["exec_error"] = exec_info["error"]905        cand["rows_preview"] = (906            exec_info["rows"][:5] if exec_info["ok"] and exec_info["rows"] else None907        )908        cand["columns"] = exec_info["columns"]909 910        candidates.append(cand)911 912        if exec_info["ok"]:913            if (not best_exec) or cand["score"] > best_score:914                best_exec = True915                best_score = cand["score"]916                best = cand917        elif not best_exec and cand["score"] > best_score:918            best_score = cand["score"]919            best = cand920 921    if best is None and candidates:922        best = candidates[0]923 924    return {925        "question_original": question,926        "detected_language": detected,927        "question_en": question_en,928        "connection_id": conn_id,929        "schema_summary": schema_str,930        "best_sql": best["sql"],931        "best_exec_ok": best.get("exec_ok", False),932        "best_exec_error": best.get("exec_error"),933        "best_rows_preview": best.get("rows_preview"),934        "best_columns": best.get("columns", []),935        "candidates": candidates,936        "score_percent": normalize_score(best["score"]),937    }938 939 940# ======================================================941# 6) Schemas Pydantic942# ======================================================943 944 945class UploadResponse(BaseModel):946    connection_id: str947    label: str948    db_path: str949    note: Optional[str] = None950 951 952class ConnectionInfo(BaseModel):953    connection_id: str954    label: str955    engine: Optional[str] = None956    db_name: Optional[str] = None  # ya no usamos archivo, pero mantenemos campo957 958 959class SchemaResponse(BaseModel):960    connection_id: str961    schema_summary: str962    tables: Dict[str, Dict[str, List[str]]]963 964 965class PreviewResponse(BaseModel):966    connection_id: str967    table: str968    columns: List[str]969    rows: List[List[Any]]970 971 972class InferRequest(BaseModel):973    connection_id: str974    question: str975 976 977class InferResponse(BaseModel):978    question_original: str979    detected_language: str980    question_en: str981    connection_id: str982    schema_summary: str983    best_sql: str984    best_exec_ok: bool985    best_exec_error: Optional[str]986    best_rows_preview: Optional[List[List[Any]]]987    best_columns: List[str]988    candidates: List[Dict[str, Any]]989 990 991class SpeechInferResponse(BaseModel):992    transcript: str993    result: InferResponse994 995 996# ======================================================997# 7) Helpers para /upload (.sql y .zip)998# ======================================================999 1000 1001def _combine_sql_files_from_zip(zip_bytes: bytes) -> str:1002    """1003    Lee un ZIP, se queda solo con los .sql y los concatena.1004    Orden:1005      1) archivos con 'schema' o 'structure' en el nombre1006      2) el resto (data, etc.)1007    """1008    try:1009        with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:1010            names = [info.filename for info in zf.infolist() if not info.is_dir()]1011            sql_names = [n for n in names if n.lower().endswith(".sql")]1012 1013            if not sql_names:1014                raise ValueError("El ZIP no contiene archivos .sql utilizables.")1015 1016            def sort_key(name: str) -> int:1017                nl = name.lower()1018                if "schema" in nl or "structure" in nl:1019                    return 01020                return 11021 1022            sql_names_sorted = sorted(sql_names, key=sort_key)1023 1024            parts: List[str] = []1025            for name in sql_names_sorted:1026                with zf.open(name) as f:1027                    text = f.read().decode("utf-8", errors="ignore")1028                parts.append(f"-- FILE: {name}\n{text}\n")1029 1030            return "\n\n".join(parts)1031    except zipfile.BadZipFile:1032        raise ValueError("Archivo ZIP inválido o corrupto.")1033 1034 1035# ======================================================1036# 8) Endpoints FastAPI1037# ======================================================1038 1039 1040@app.on_event("startup")1041async def startup_event():1042    load_nl2sql_model()1043    print("✅ Backend NL2SQL inicializado.")1044    print(f"MODEL_DIR={MODEL_DIR}, DEVICE={DEVICE}")1045    print(f"Conexiones activas al inicio: {len(sql_manager.connections)}")1046 1047 1048@app.post("/upload", response_model=UploadResponse)1049async def upload_database(1050    mode: str = Form("full"),  # "full" | "schema_data" | "zip"1051    db_files: List[UploadFile] = File(...),  # uno o varios archivos1052    authorization: Optional[str] = Header(None),1053):1054    """1055    Sube uno o varios archivos SQL/ZIP según el modo:1056 1057    - mode = "full":1058        * Espera EXACTAMENTE 1 archivo .sql1059        * El .sql trae esquema + datos juntos (dump de PostgreSQL)1060 1061    - mode = "schema_data":1062        * Espera EXACTAMENTE 2 archivos .sql1063        * Uno de esquema y otro de datos (el orden lo resolvemos nosotros)1064 1065    - mode = "zip":1066        * Espera EXACTAMENTE 1 archivo .zip1067        * Dentro del zip buscamos SOLO archivos .sql (ignoramos el resto)1068    """1069    if authorization is None:1070        raise HTTPException(401, "Missing Authorization header")1071 1072    jwt = authorization.replace("Bearer ", "")1073    user = supabase.auth.get_user(jwt)1074    if not user or not user.user:1075        raise HTTPException(401, "Invalid Supabase token")1076 1077    if not db_files:1078        raise HTTPException(400, "No se recibió ningún archivo.")1079 1080    mode = mode.lower().strip()1081 1082    # =======================1083    # MODO 1: FULL (.sql único)1084    # =======================1085    if mode == "full":1086        if len(db_files) != 1:1087            raise HTTPException(1088                400, "Modo FULL requiere exactamente 1 archivo .sql."1089            )1090 1091        file = db_files[0]1092        filename = file.filename or ""1093        if not filename.lower().endswith(".sql"):1094            raise HTTPException(400, "Modo FULL solo acepta archivos .sql.")1095 1096        contents = await file.read()1097        sql_text = contents.decode("utf-8", errors="ignore")1098 1099    # ====================================1100    # MODO 2: ESQUEMA + DATOS (2 archivos)1101    # ====================================1102    elif mode == "schema_data":1103        if len(db_files) != 2:1104            raise HTTPException(1105                400,1106                "Modo esquema+datos requiere exactamente 2 archivos .sql.",1107            )1108 1109        print("FILES RECEIVED:", [f.filename for f in db_files])1110 1111        files_info: List[tuple[str, str]] = []1112        for f in db_files:1113            fname = f.filename or ""1114            if not fname.lower().endswith(".sql"):1115                raise HTTPException(400, "Todos los archivos deben ser .sql.")1116            contents = await f.read()1117            files_info.append(1118                (fname, contents.decode("utf-8", errors="ignore"))1119            )1120 1121        # Intentamos poner primero el esquema y luego los datos1122        def weight(name: str) -> int:1123            nl = name.lower().replace("-", "_").replace(" ", "_")1124        1125            if any(x in nl for x in ["schema", "structure", "ddl"]):1126                return 01127            if any(x in nl for x in ["data", "dml", "insert", "rows"]):1128                return 11129            return 21130 1131        files_info_sorted = sorted(files_info, key=lambda x: weight(x[0]))1132 1133        sql_parts: List[str] = []1134        for fname, text in files_info_sorted:1135            sql_parts.append(f"-- FILE: {fname}\n{text}\n")1136 1137        sql_text = "\n\n".join(sql_parts)1138        # usamos el nombre del primer archivo como label "principal"1139        filename = files_info_sorted[0][0]1140 1141    # ==================1142    # MODO 3: ZIP (.zip)1143    # ==================1144    elif mode == "zip":1145        if len(db_files) != 1:1146            raise HTTPException(1147                400, "Modo ZIP requiere exactamente 1 archivo .zip."1148            )1149 1150        file = db_files[0]1151        filename = file.filename or ""1152        if not filename.lower().endswith(".zip"):1153            raise HTTPException(400, "Modo ZIP solo acepta archivos .zip.")1154 1155        contents = await file.read()1156        # tu helper ya ignora carpetas y solo concatena .sql1157        sql_text = _combine_sql_files_from_zip(contents)1158 1159    else:1160        raise HTTPException(400, f"Modo no soportado: {mode}")1161 1162    # --- crear schema dinámico en Postgres (Neon) ---1163    try:1164        conn_id = sql_manager.create_database_from_dump(1165            label=filename, sql_text=sql_text1166        )1167    except Exception as e:1168        raise HTTPException(400, f"Error creando BD: {e}")1169 1170    meta = sql_manager.connections[conn_id]1171 1172    # --- guardar metadatos en Supabase (sin romper el upload si falla) ---1173    try:1174        supabase.table("databases").insert(1175            {1176                "user_id": user.user.id,1177                "filename": filename,1178                "engine": meta["engine"],1179                "connection_id": conn_id,1180            }1181        ).execute()1182    except Exception as e:1183        # Solo logeamos, pero NO rompemos el endpoint1184        print("[WARN] No se pudieron guardar metadatos en Supabase:", repr(e))1185 1186    return UploadResponse(1187        connection_id=conn_id,1188        label=filename,1189        db_path=f"{meta['engine']}://schema/{meta['schema']}",1190        note="Database schema created in Neon and indexed in Supabase.",1191    )1192 1193 1194@app.get("/connections", response_model=List[ConnectionInfo])1195async def list_connections():1196    return [1197        ConnectionInfo(1198            connection_id=cid,1199            label=meta.get("label", ""),1200            engine=meta.get("engine"),

Showing the first 1,200 of 1397 lines. Download the file for the rest.