chatbot_v3/src/bootstrap/lifespan.py

from collections.abc import AsyncIterator, Callable, Sequence
from contextlib import AbstractAsyncContextManager, asynccontextmanager

import httpx
import structlog
from anyio import CapacityLimiter, Semaphore, to_thread
from fastapi import FastAPI

from src.application.ingestion import get_encoder
from src.application.ports.embedding import DenseEmbedder
from src.bootstrap.dependencies import AppResources
from src.config import Settings
from src.infrastructure.embedding.bm25 import Bm25SparseEmbedder
from src.infrastructure.embedding.openai_compatible import (
    OpenAICompatibleEmbedder,
    is_ollama_base_url,
)
from src.infrastructure.minio.client import create_client as create_minio_client
from src.infrastructure.minio.storage import MinioObjectStorage
from src.infrastructure.observability.logging import configure_logging
from src.infrastructure.postgres.database import create_engine, create_sessionmaker
from src.infrastructure.qdrant.client import create_client as create_qdrant_client
from src.infrastructure.qdrant.points import QdrantPointStorage

logger = structlog.get_logger(__name__)


def _auth_headers(api_key: str | None) -> dict[str, str]:
    """Bearer header, or none at all when no key is configured.

    Sending an empty `Bearer ` is worse than sending nothing: some gateways
    treat a malformed credential as an auth failure rather than as anonymous.
    """
    return {"Authorization": f"Bearer {api_key}"} if api_key else {}


async def _warm_dense_embedders(embedders: Sequence[DenseEmbedder]) -> None:
    """Force each dense model to load before the first upload needs it.

    Same rationale as the tiktoken warm-up above, but with the opposite
    failure policy. A self-hosted embedder that has unloaded the model takes
    minutes to serve its first request — longer than
    `INGESTION_TIMEOUT_SECONDS` — so paying that once at boot keeps it off a
    user's upload. Unlike the tokenizer this is best-effort: an embedder that
    is merely *down* must not stop the process from booting and reporting its
    own health, and `/readyz` is where that condition belongs.
    """
    for embedder in embedders:
        try:
            await embedder.embed_batch(["warmup"])
            logger.info("lifespan.embedder.warmed", embedder=embedder.name)
        except Exception:
            logger.warning("lifespan.embedder.warm_failed", embedder=embedder.name, exc_info=True)


def create_lifespan(
    settings: Settings | None = None,
) -> Callable[[FastAPI], AbstractAsyncContextManager[None, bool | None]]:
    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
        resolved_settings = settings or Settings()
        configure_logging(resolved_settings.logging, resolved_settings.app)

        # tiktoken fetches its vocabulary over the network on first use, so warm
        # it here: a missing vocabulary should fail the process at boot, not the
        # first upload. Blocking, hence the thread.
        await to_thread.run_sync(get_encoder, resolved_settings.chunking.encoding_name)
        logger.info(
            "lifespan.tokenizer.loaded",
            encoding=resolved_settings.chunking.encoding_name,
        )

        db_engine = create_engine(resolved_settings.postgres)
        db_sessionmaker = create_sessionmaker(db_engine)
        logger.info("lifespan.postgres.engine.created")

        minio_client = create_minio_client(resolved_settings.minio)
        logger.info("lifespan.minio.client.created")

        qdrant_client = create_qdrant_client(resolved_settings.qdrant)
        # No collection DDL here: `ensure_chunks_collection` is a deployment
        # step (`python -m src.cli.qdrant_bootstrap`), for the same reason
        # ADR-0009 keeps Alembic out of startup and ADR-0012 makes LangGraph's
        # `.setup()` a deployment step.
        point_storage = QdrantPointStorage(
            qdrant_client, collection=resolved_settings.qdrant.collection
        )
        logger.info("lifespan.qdrant.client.created")

        nomic_settings = resolved_settings.embedding.nomic
        nomic_http_client = httpx.AsyncClient(
            base_url=nomic_settings.base_url,
            timeout=nomic_settings.timeout_seconds,
            headers=_auth_headers(nomic_settings.api_key),
        )
        openai_settings = resolved_settings.embedding.openai
        openai_http_client = httpx.AsyncClient(
            base_url=openai_settings.base_url,
            timeout=openai_settings.timeout_seconds,
            headers=_auth_headers(openai_settings.api_key),
        )
        dense_embedders = (
            OpenAICompatibleEmbedder(
                nomic_http_client,
                name="dense_nomic",
                model=nomic_settings.model,
                document_prefix=nomic_settings.document_prefix,
                keep_alive=(
                    nomic_settings.keep_alive
                    if is_ollama_base_url(nomic_settings.base_url)
                    else None
                ),
            ),
            OpenAICompatibleEmbedder(
                openai_http_client,
                name="dense_openai",
                model=openai_settings.model,
                dimensions=openai_settings.dimensions,
                document_prefix=openai_settings.document_prefix,
            ),
        )
        sparse_embedder = Bm25SparseEmbedder(resolved_settings.embedding.sparse)
        logger.info("lifespan.embedders.created")

        await _warm_dense_embedders(dense_embedders)

        # Bounds how many ingestions run in this process at once (ADR-0017);
        # a distinct resource from ingestion_limiter, which bounds threads
        # spent on blocking work within a single ingestion.
        ingestion_concurrency_limiter = Semaphore(resolved_settings.ingestion.max_concurrency)

        # Bounds threads spent on blocking ingestion work (parsing, chunking,
        # hashing, the sync minio SDK) so it cannot exhaust Starlette's own
        # thread pool (ADR-0017).
        ingestion_limiter = CapacityLimiter(resolved_settings.ingestion.thread_pool_size)
        object_storage = MinioObjectStorage(
            minio_client, bucket=resolved_settings.minio.bucket, limiter=ingestion_limiter
        )

        app.state.resources = AppResources(
            settings=resolved_settings,
            db_engine=db_engine,
            db_sessionmaker=db_sessionmaker,
            minio_client=minio_client,
            qdrant_client=qdrant_client,
            object_storage=object_storage,
            point_storage=point_storage,
            ingestion_limiter=ingestion_limiter,
            dense_embedders=dense_embedders,
            sparse_embedder=sparse_embedder,
            ingestion_concurrency_limiter=ingestion_concurrency_limiter,
        )

        try:
            yield
        finally:
            try:
                await db_engine.dispose()
            except Exception:
                logger.exception("lifespan.postgres.dispose.failed")

            try:
                await qdrant_client.close()
            except Exception:
                logger.exception("lifespan.qdrant.close.failed")

            try:
                await nomic_http_client.aclose()
            except Exception:
                logger.exception("lifespan.embedding.nomic_client.close.failed")

            try:
                await openai_http_client.aclose()
            except Exception:
                logger.exception("lifespan.embedding.openai_client.close.failed")

    return lifespan