backend/server.py

"""Punto de entrada FastAPI del backend (Semantic Engine).

Tras Hardening 3, este modulo se limita a la composicion de la app: el ciclo de vida
(`lifespan`, que carga el modelo y el vocabulario), el middleware de seguridad (CORS de
Hardening 1) y el montaje de los APIRouters por dominio (backend/routers/*). Los endpoints
y el estado global NO viven aca: las rutas estan en `routers/` y el estado/los singletons
en `state.py`.

`lifespan` es quien POBLA el estado de `state.py` al arranque (carga del SentenceTransformer
y precomputo del vocabulario). Importar `state` no dispara esa carga.
"""

import asyncio
import logging
import os
import time
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

from routers import core, dimension, models, relief, sae, telemetry

# Estado global compartido y helpers. Importar `state` NO carga ningun modelo: el
# SentenceTransformer se carga en `lifespan` (abajo) o al cambiar de modelo. `sae_manager`
# y `get_relief_pipeline` se reexportan aca por compatibilidad con importadores externos.
from state import (
    APP_MODE,
    CORS_ORIGINS,
    DEFAULT_MODEL_KEY,
    HOST,
    PORT,
    VOCAB_PATH,
    get_optimal_device,
    get_relief_pipeline,
    load_demo_vocab_bundle,
    sae_manager,
    set_active_model,
    state,
)

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)


# --- Inicialización ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Set the running event loop on telemetry_manager
    from telemetry import telemetry_manager

    telemetry_manager.loop = asyncio.get_running_loop()

    env_device = os.getenv("DEVICE", "AUTO").upper().strip()
    device = get_optimal_device(env_device) if APP_MODE != "demo" else "cpu"
    if APP_MODE != "demo":
        logger.info(
            f"Resolved device configuration: {env_device} -> target device: {device}"
        )
    else:
        logger.info(
            "APP_MODE=demo: omitiendo resolución GPU; embeddings vía serverless HF."
        )

    start_time = time.time()
    await asyncio.to_thread(set_active_model, DEFAULT_MODEL_KEY)
    logger.info(
        "Modelo activo '%s' listo en %.2fs (APP_MODE=%s, %dD).",
        state.active_model_key,
        time.time() - start_time,
        APP_MODE,
        state.embedding_dim,
    )

    # Vocabulario: demo carga npz empaquetado; local precalcula vía modelo
    if APP_MODE == "demo":
        load_demo_vocab_bundle()
    elif os.path.exists(VOCAB_PATH):
        logger.info("🧠 Pre-calculando vocabulario para interpretación...")
        with open(VOCAB_PATH, encoding="utf-8") as f:
            state.vocab_words = [line.strip() for line in f if line.strip()]

        # Inferencia masiva (Esto es lo que mataba a Chrome)
        t0 = time.time()
        state.vocab_embeddings = state.model.encode(
            state.vocab_words, convert_to_tensor=True
        )
        # Convertir a numpy para búsquedas rápidas
        state.vocab_embeddings = state.vocab_embeddings.cpu().numpy()

        # Pre-calculate Token IDs
        # Use the tokenizer to get IDs. We take the first meaningful ID (skipping CLS if present).
        # encode returns [CLS, ID, SEP] usually.
        logger.info("🔢 Calculating Token IDs...")
        all_ids = state.model.tokenizer(state.vocab_words, add_special_tokens=False)[
            "input_ids"
        ]
        # all_ids is a list of lists. We take the first ID of each word.
        state.vocab_ids = [ids[0] if ids else 0 for ids in all_ids]

        logger.info(
            f"✅ {len(state.vocab_words)} palabras procesadas en {time.time() - t0:.2f}s"
        )
    else:
        logger.warning(
            f"⚠️ No se encontró {VOCAB_PATH}. El modo Interpretación estará desactivado."
        )

    # Relief Semantico: verificar la tabla persistida contra la dimensionalidad del modelo
    # activo. Si quedo una tabla heredada de otro ancho (p. ej. 1024D de BGE-M3), se descarta
    # para recrearla limpia en la proxima ingesta. No bloquea el arranque ante errores.
    try:
        get_relief_pipeline().ensure_storage_ready()
    except Exception as relief_err:
        logger.error(
            "No se pudo verificar la tabla de relieve al arranque: %s", relief_err
        )

    # Configure and lazily load SAE
    sae_manager.device = device
    if sae_manager.is_trained():
        logger.info("SAE weights found, lazily initializing...")
        sae_manager.load_model()

    yield


app = FastAPI(title="Semantic Engine v2.0", lifespan=lifespan)

# CORS restringido: solo los origenes declarados en CORS_ORIGINS (default: frontend Vite).
# allow_credentials=False porque no se intercambian cookies/credenciales cross-origin.
app.add_middleware(
    CORSMiddleware,
    allow_origins=CORS_ORIGINS,
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Montaje de routers por dominio (Hardening 3). Cada router define sus endpoints e importa
# el estado compartido de state.py. Las rutas y shapes son identicas al monolito previo.
app.include_router(core.router)
app.include_router(relief.router)
app.include_router(models.router)
app.include_router(dimension.router)
app.include_router(sae.router)
app.include_router(telemetry.router, prefix="/api", tags=["telemetry"])

dist_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../dist"))
if os.path.exists(dist_path):
    logger.info(f"Montando estáticos del frontend desde {dist_path}")
    app.mount("/", StaticFiles(directory=dist_path, html=True), name="static")


if __name__ == "__main__":
    import uvicorn

    # Bind loopback por default. Para escuchar en la red, exportar HOST=0.0.0.0.
    uvicorn.run(app, host=HOST, port=PORT)