"""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) ) ] ) ), )