app/core/lifecycle.py

"""FastAPI startup and shutdown lifecycle hooks."""

from contextlib import asynccontextmanager
from typing import AsyncIterator
import os
from pathlib import Path

from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles

from app.core.logging import get_logger, setup_logging
from app.core.paths import ensure_storage_dirs
from app.database.connection import engine, verify_connection, async_session_factory
from app.database.migrations import apply_migrations
from app.database.pgvector import create_vector_index, ensure_vector_dimension

logger = get_logger(__name__)


async def _requeue_all_embeddings(session) -> None:
    """Dispatch embed_document Celery jobs for every document that has chunks.

    Called after a vector-dimension migration wiped chunk_embeddings. Summaries
    and figure descriptions are prompt-hash cached, so only the (cheap)
    embedding pass re-runs. If Redis/Celery is down the jobs are simply not
    queued — searches return empty until a re-embed happens, so warn loudly.
    """
    from sqlalchemy import text

    result = await session.execute(
        text("SELECT DISTINCT document_id FROM chunks")
    )
    doc_ids = [row[0] for row in result.fetchall()]
    if not doc_ids:
        return
    try:
        from app.workers.tasks import embed_document
        for doc_id in doc_ids:
            embed_document.delay(str(doc_id))
        logger.warning(
            "Re-queued embedding jobs for %d document(s) after the vector-dimension "
            "migration. Vector search returns sparse results until they finish.",
            len(doc_ids),
        )
    except Exception:
        logger.exception(
            "Could not queue re-embedding jobs (is Redis running?). Vector search "
            "will return nothing until documents are re-embedded."
        )


async def _report_ai_backend() -> None:
    """Log which backend auto-detection picked — or, if nothing is usable,
    the exact instructions the user asked for ("put your API key or your
    Ollama connection"). Never fatal: the app still serves stored papers."""
    from app.api.errors import ModelUnavailable, NoLLMConfigured
    from app.llm.resolver import resolve_llm

    try:
        await resolve_llm()  # logs "LLM backend: ..." on success
    except NoLLMConfigured as e:
        logger.error(str(e.model))
    except ModelUnavailable as e:
        logger.error("LLM backend misconfigured: %s", e.model)


async def _check_embedding_model_switch(session) -> None:
    """Detect stored embeddings made by a different model than the active one.

    Vectors from different models are not comparable, so search quality
    silently degrades after a backend switch. When EMBEDDING_PROVIDER is
    pinned explicitly the switch is deliberate: wipe the stale vectors and
    re-embed (summaries/figure descriptions are prompt-hash cached and don't
    re-run). In auto mode only warn — a temporarily unreachable Ollama must
    not trigger a destructive re-embed loop.
    """
    from sqlalchemy import text

    from app.api.errors import ModelUnavailable
    from app.core.config import settings
    from app.llm.resolver import resolve_embedding

    try:
        target = await resolve_embedding()
    except ModelUnavailable as e:
        logger.error("Embedding backend misconfigured: %s", e.model)
        return

    result = await session.execute(
        text("SELECT DISTINCT embedding_model FROM chunk_embeddings")
    )
    stored = [row[0] for row in result.fetchall()]
    stale = [m for m in stored if m != target.model]
    if not stale:
        return

    pinned = (settings.embedding_provider or "auto").strip().lower() not in ("", "auto")
    if not pinned:
        logger.warning(
            "Stored embeddings were generated by %s but the active embedding "
            "model is %s (provider: %s) — these vectors are NOT comparable, so "
            "search quality is degraded. If this switch is permanent, pin "
            "EMBEDDING_PROVIDER in backend/.env and restart: the library will "
            "be re-embedded automatically. Otherwise restore the old backend.",
            ", ".join(stale), target.model, target.provider,
        )
        return

    logger.warning(
        "EMBEDDING_PROVIDER is pinned to %s (model: %s) but stored embeddings "
        "were generated by %s. Re-embedding the library with the new model.",
        target.provider, target.model, ", ".join(stale),
    )
    await session.execute(
        text("DELETE FROM chunk_embeddings WHERE embedding_model != :model"),
        {"model": target.model},
    )
    await session.commit()
    await _requeue_all_embeddings(session)


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    """Application lifespan: startup and shutdown.

    Long-running work (PDF ingestion, embedding) is dispatched to Celery workers.
    """
    setup_logging()
    logger.info("Starting 9XAIPal backend")

    from app.core.config import settings
    if settings.postgres_password == "9xaipal_dev_password":
        logger.warning(
            "PostgreSQL is using the default development password. "
            "Set POSTGRES_PASSWORD in backend/.env before exposing this server to a network."
        )

    # Ensure storage directories exist
    ensure_storage_dirs()

    # Optional single-origin SPA serving for "your machine = server" mode.
    # Done here (in lifespan) rather than at pure import time so the volume
    # state is guaranteed stable when the decision is made (critical for
    # Docker named volumes + uvicorn workers).
    try:
        _frontend_dist = Path("/app/frontend/dist")
        _serve_frontend = os.getenv("SERVE_FRONTEND", "false").lower() in ("1", "true", "yes")
        _has_dist = _frontend_dist.exists() and (_frontend_dist / "index.html").exists()
        if _serve_frontend or _has_dist:
            if _has_dist:
                app.mount(
                    "/",
                    StaticFiles(directory=str(_frontend_dist), html=True, check_dir=False),
                    name="frontend-spa",
                )
                logger.info("SPA frontend mounted at / (single-port server mode active)")
    except Exception as e:
        logger.warning("Frontend SPA mount skipped (non-fatal): %s", e)

    # Verify database connectivity and apply migrations
    await verify_connection()
    await apply_migrations()

    # Report which AI backend auto-detection picked (Ollama → cloud API keys
    # → clear configure-me message).
    await _report_ai_backend()

    # Sync the embedding column to the configured dimension, then ensure the
    # search indexes exist (idempotent, cheap when nothing changed).
    async with async_session_factory() as session:
        dimension_migrated = await ensure_vector_dimension(session)
        await create_vector_index(session)
        await session.commit()
        if dimension_migrated:
            await _requeue_all_embeddings(session)
        else:
            # Dimension unchanged — but the embedding MODEL may have switched
            # (e.g. moving the library from Ollama to a cloud embedder).
            await _check_embedding_model_switch(session)

    logger.info("9XAIPal backend ready")
    yield

    # Shutdown
    logger.info("Shutting down 9XAIPal backend")
    await engine.dispose()