stvnnnnnn/nl2sql-backend-t5
1
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"),