app/vector_store.py

"""Qdrant vector store: collection management, upsert, and search."""
from __future__ import annotations

import uuid
from functools import lru_cache
from typing import List, Optional

from qdrant_client import QdrantClient
from qdrant_client.http import models as qm

from .config import get_settings


@lru_cache
def get_client() -> QdrantClient:
    settings = get_settings()
    return QdrantClient(url=settings.qdrant_url, api_key=settings.qdrant_api_key, timeout=60)


def ensure_collection() -> None:
    """Create the collection if it does not exist."""
    settings = get_settings()
    client = get_client()
    existing = {c.name for c in client.get_collections().collections}
    if settings.qdrant_collection not in existing:
        client.create_collection(
            collection_name=settings.qdrant_collection,
            vectors_config=qm.VectorParams(
                size=settings.embedding_dim, distance=qm.Distance.COSINE
            ),
        )
        # Index on document_id so we can filter / delete by document.
        client.create_payload_index(
            collection_name=settings.qdrant_collection,
            field_name="document_id",
            field_schema=qm.PayloadSchemaType.KEYWORD,
        )


def upsert_chunks(
    document_id: str,
    filename: str,
    chunks: List[str],
    vectors: List[List[float]],
) -> None:
    settings = get_settings()
    client = get_client()
    points = [
        qm.PointStruct(
            id=str(uuid.uuid4()),
            vector=vector,
            payload={
                "document_id": document_id,
                "filename": filename,
                "chunk_index": idx,
                "text": chunk,
            },
        )
        for idx, (chunk, vector) in enumerate(zip(chunks, vectors))
    ]
    client.upsert(collection_name=settings.qdrant_collection, points=points)


def search(
    query_vector: List[float],
    top_k: int,
    document_ids: Optional[List[str]] = None,
):
    settings = get_settings()
    client = get_client()

    query_filter = None
    if document_ids:
        query_filter = qm.Filter(
            must=[qm.FieldCondition(key="document_id", match=qm.MatchAny(any=document_ids))]
        )

    # query_points is the current API (works on qdrant-client >=1.10; the old
    # .search() was removed in 1.18). Returns a QueryResponse; .points is the
    # list of ScoredPoint (each has .payload and .score), matching prior usage.
    return client.query_points(
        collection_name=settings.qdrant_collection,
        query=query_vector,
        limit=top_k,
        query_filter=query_filter,
        with_payload=True,
    ).points


def delete_document(document_id: str) -> None:
    settings = get_settings()
    client = get_client()
    client.delete(
        collection_name=settings.qdrant_collection,
        points_selector=qm.FilterSelector(
            filter=qm.Filter(
                must=[
                    qm.FieldCondition(
                        key="document_id", match=qm.MatchValue(value=document_id)
                    )
                ]
            )
        ),
    )