chatbot_v3/src/infrastructure/qdrant/collection.py

"""The `chunks` collection schema and its provisioning (ADR-0001).

**This runs as a deployment step, never at FastAPI startup.** Creating a
collection is DDL, and the project already rules DDL out of the request/boot
path: ADR-0009 requires Alembic for Postgres schema, and ADR-0012 makes
LangGraph's `.setup()` a deployment step. Doing it in the lifespan would also
couple boot to Qdrant being reachable (that is `/readyz`'s job), race across
replicas, and hide a misconfigured collection until traffic arrives.

Entry point for operators: `uv run python -m src.cli.qdrant_bootstrap`.

Two properties of this schema are load-bearing and fail *silently* if wrong,
which is why `ensure_chunks_collection` verifies rather than skips:

- **`sparse` must carry `modifier=IDF`.** `src/infrastructure/embedding/bm25.py`
  computes only BM25's term-frequency saturation; Qdrant supplies IDF from
  collection-wide statistics. Without the modifier there is no error and no
  warning — lexical retrieval just quietly loses its IDF term (ADR-0005).
- **The dense dimensions are pinned**: `dense_nomic` 768,
  `dense_openai` 3072 (native, `dimensions` deliberately unset). They are
  constants here rather than settings because they are model facts; changing
  one is a re-embedding migration, not a config tweak (ADR-0001).

All four named vectors are defined at creation even though `late_interaction`
stays unpopulated until ADR-0003's rerank work (ADR-0017 does not compute it at
ingest). Sparse and multivector fields cannot be added to an existing
collection without recreating it, so deferring them is the one thing this
schema cannot afford.
"""

import structlog
from qdrant_client import AsyncQdrantClient, models

logger = structlog.get_logger(__name__)

DENSE_NOMIC_VECTOR = "dense_nomic"
DENSE_OPENAI_VECTOR = "dense_openai"
SPARSE_VECTOR = "sparse"
LATE_INTERACTION_VECTOR = "late_interaction"

DENSE_NOMIC_DIMENSIONS = 768
DENSE_OPENAI_DIMENSIONS = 3072
# jina-colbert-v2's per-token output dimension (ADR-0005).
LATE_INTERACTION_DIMENSIONS = 128


class CollectionSchemaMismatchError(RuntimeError):
    """An existing collection does not match the schema this code expects.

    Raised loudly instead of returning: silently accepting a collection whose
    dense size or sparse modifier differs is precisely the failure this
    explicit bootstrap step exists to prevent.
    """


def _vectors_config() -> dict[str, models.VectorParams]:
    return {
        DENSE_NOMIC_VECTOR: models.VectorParams(
            size=DENSE_NOMIC_DIMENSIONS, distance=models.Distance.COSINE
        ),
        DENSE_OPENAI_VECTOR: models.VectorParams(
            size=DENSE_OPENAI_DIMENSIONS, distance=models.Distance.COSINE
        ),
        # Rerank-only: never independently ANN-searched, so its HNSW graph is
        # disabled (m=0), and stored on disk so its larger footprint does not
        # degrade dense/sparse query latency (ADR-0001).
        LATE_INTERACTION_VECTOR: models.VectorParams(
            size=LATE_INTERACTION_DIMENSIONS,
            distance=models.Distance.COSINE,
            multivector_config=models.MultiVectorConfig(
                comparator=models.MultiVectorComparator.MAX_SIM
            ),
            hnsw_config=models.HnswConfigDiff(m=0),
            on_disk=True,
        ),
    }


def _sparse_vectors_config() -> dict[str, models.SparseVectorParams]:
    return {SPARSE_VECTOR: models.SparseVectorParams(modifier=models.Modifier.IDF)}


# `content`'s full-text index backs ADR-0002's keyword search. The
# `multilingual` tokenizer is the one that segments Persian correctly; `word`
# splits on non-alphanumerics, which mis-handles ZWNJ-joined compounds. No
# stemmer or stopword list is configured: `content` is already letter-folded by
# `normalize_persian_text` at ingest (ADR-0018), and the *ranked* Farsi lexical
# path is the benchmarked BM25 sparse vector, not this index. This one exists
# for exact keyword/filter matching, which ADR-0002 keeps deliberately distinct
# from retrieval.
_CONTENT_INDEX = models.TextIndexParams(
    type=models.TextIndexType.TEXT,
    tokenizer=models.TokenizerType.MULTILINGUAL,
    lowercase=True,
)

# (field name, schema). `order_id` is float because Qdrant's `Range` conditions
# and `order_by` only support numeric/datetime payloads -- a keyword key could
# only be sorted client-side after fetching every chunk (ADR-0001).
_PayloadIndexSchema = models.PayloadSchemaType | models.KeywordIndexParams | models.TextIndexParams
_PAYLOAD_INDEXES: tuple[tuple[str, _PayloadIndexSchema], ...] = (
    # `is_tenant` co-locates a tenant's vectors on disk for sequential reads,
    # which is the whole point of payload-partitioned multitenancy.
    (
        "tenant_id",
        models.KeywordIndexParams(type=models.KeywordIndexType.KEYWORD, is_tenant=True),
    ),
    ("domain", models.PayloadSchemaType.KEYWORD),
    ("file_id", models.PayloadSchemaType.KEYWORD),
    ("order_id", models.PayloadSchemaType.FLOAT),
    ("previous_chunk_id", models.PayloadSchemaType.KEYWORD),
    ("next_chunk_id", models.PayloadSchemaType.KEYWORD),
    ("content", _CONTENT_INDEX),
    # Every read path filters on `is_active` (ADR-0002 implies `is_active: true`
    # unless the caller opts in), and the re-ingestion sweep and `/v1/points`
    # both range over `chunk_index`. Both were unindexed while ingestion was the
    # only reader; plan 002 makes them hot.
    ("is_active", models.PayloadSchemaType.BOOL),
    ("chunk_index", models.PayloadSchemaType.INTEGER),
)


def _verify_existing(collection: str, info: models.CollectionInfo) -> None:
    params = info.config.params
    vectors = params.vectors
    if not isinstance(vectors, dict):
        raise CollectionSchemaMismatchError(
            f"collection {collection!r} has an unnamed dense vector; ADR-0001 requires "
            f"named vectors and this cannot be fixed without recreating the collection"
        )

    for name, expected_size in (
        (DENSE_NOMIC_VECTOR, DENSE_NOMIC_DIMENSIONS),
        (DENSE_OPENAI_VECTOR, DENSE_OPENAI_DIMENSIONS),
        (LATE_INTERACTION_VECTOR, LATE_INTERACTION_DIMENSIONS),
    ):
        existing = vectors.get(name)
        if existing is None:
            raise CollectionSchemaMismatchError(
                f"collection {collection!r} is missing the {name!r} vector"
            )
        if existing.size != expected_size:
            raise CollectionSchemaMismatchError(
                f"collection {collection!r} has {name!r} at {existing.size} dimensions, "
                f"expected {expected_size}; re-dimensioning is a re-embedding migration"
            )

    sparse = (params.sparse_vectors or {}).get(SPARSE_VECTOR)
    if sparse is None:
        raise CollectionSchemaMismatchError(
            f"collection {collection!r} is missing the {SPARSE_VECTOR!r} vector; sparse "
            f"vectors cannot be added without recreating the collection"
        )
    if sparse.modifier != models.Modifier.IDF:
        raise CollectionSchemaMismatchError(
            f"collection {collection!r} has {SPARSE_VECTOR!r} with modifier "
            f"{sparse.modifier!r}, expected 'idf'; without it Qdrant applies no IDF "
            f"and lexical retrieval silently degrades (ADR-0005)"
        )


async def ensure_chunks_collection(client: AsyncQdrantClient, *, collection: str) -> bool:
    """Create the `chunks` collection and its payload indexes if absent.

    Idempotent: an existing collection is verified against the pinned schema
    and left alone. Returns whether it created the collection.

    Raises `CollectionSchemaMismatchError` if an existing collection diverges.
    """
    if await client.collection_exists(collection):
        _verify_existing(collection, await client.get_collection(collection))
        logger.info("qdrant.collection.verified", collection=collection)
        # Payload indexes are additive and idempotent, so (re)creating them
        # here is what lets an index be added to an already-live collection.
        await _create_payload_indexes(client, collection=collection)
        return False

    await client.create_collection(
        collection_name=collection,
        vectors_config=_vectors_config(),
        sparse_vectors_config=_sparse_vectors_config(),
        # m=0 disables the global index; payload_m builds per-tenant graphs
        # instead, per Qdrant's multitenant guidance (ADR-0001).
        hnsw_config=models.HnswConfigDiff(m=0, payload_m=16),
    )
    logger.info("qdrant.collection.created", collection=collection)
    await _create_payload_indexes(client, collection=collection)
    return True


async def _create_payload_indexes(client: AsyncQdrantClient, *, collection: str) -> None:
    for field_name, field_schema in _PAYLOAD_INDEXES:
        await client.create_payload_index(
            collection_name=collection, field_name=field_name, field_schema=field_schema
        )
    logger.info(
        "qdrant.collection.payload_indexes.ensured",
        collection=collection,
        fields=[name for name, _ in _PAYLOAD_INDEXES],
    )