import uuid from qdrant_client import QdrantClient, models from app.services.vector_store_contract import ( EmbeddingMetadata, SearchResult, VectorDocument, ) class QdrantVectorStore: HEALTH_TENANT_ID = "qdrant-health" HEALTH_ENTRY_ID = "qdrant-health-probe" HEALTH_TEXT = "qdrant health probe" def __init__( self, url: str, collection_name: str, vector_size: int, timeout: int = 10, client: QdrantClient | None = None, ): self.client = client or QdrantClient(url=url, timeout=timeout) self.collection_name = collection_name self.vector_size = vector_size @staticmethod def _point_id(tenant_id: str, vector_id: str) -> str: return str(uuid.uuid5(uuid.NAMESPACE_URL, f"{tenant_id}:{vector_id}")) @staticmethod def _tenant_filter(tenant_id: str) -> models.Filter: return models.Filter( must=[ models.FieldCondition( key="tenant_id", match=models.MatchValue(value=tenant_id), ) ] ) @staticmethod def _docs_filter(tenant_id: str, entry_ids: list[str]) -> models.Filter: return models.Filter( must=[ models.FieldCondition( key="tenant_id", match=models.MatchValue(value=tenant_id), ), models.FieldCondition( key="entry_id", match=models.MatchAny(any=entry_ids), ), ] ) @staticmethod def _normalize_cosine_score(score: float) -> float: return max(0.0, min(1.0, (score + 1.0) / 2.0)) def _ensure_collection(self) -> None: if self.client.collection_exists(collection_name=self.collection_name): return try: self.client.create_collection( collection_name=self.collection_name, vectors_config=models.VectorParams( size=self.vector_size, distance=models.Distance.COSINE, ), ) except Exception: if self.client.collection_exists(collection_name=self.collection_name): return raise self.client.create_payload_index( collection_name=self.collection_name, field_name="tenant_id", field_schema=models.PayloadSchemaType.KEYWORD, wait=True, ) self.client.create_payload_index( collection_name=self.collection_name, field_name="entry_id", field_schema=models.PayloadSchemaType.KEYWORD, wait=True, ) def upsert(self, document: VectorDocument) -> None: self._ensure_collection() payload = { key: value for key, value in document.metadata.items() if value is not None } payload.update( { "vector_id": document.vector_id or document.entry_id, "tenant_id": document.tenant_id, "entry_id": document.entry_id, "text": document.text, "embedding_model": document.embedding_metadata.embedding_model, "embedding_dim": document.embedding_metadata.embedding_dim, } ) self.client.upsert( collection_name=self.collection_name, points=[ models.PointStruct( id=self._point_id( document.tenant_id, document.vector_id or document.entry_id ), vector=document.embedding, payload=payload, ) ], wait=True, ) def query( self, tenant_id: str, query_embedding: list[float], n_results: int ) -> list[SearchResult]: self._ensure_collection() response = self.client.query_points( collection_name=self.collection_name, query=query_embedding, query_filter=self._tenant_filter(tenant_id), limit=n_results, with_payload=True, with_vectors=False, ) return [ SearchResult( entry_id=str(point.payload.get("entry_id", point.id)), text=str(point.payload.get("text", "")), similarity=self._normalize_cosine_score(point.score), metadata={ key: value for key, value in point.payload.items() if key not in {"text", "vector_id"} and value is not None }, ) for point in response.points ] def delete_docs(self, tenant_id: str, entry_ids: list[str]) -> None: if not entry_ids: return self._ensure_collection() self.client.delete( collection_name=self.collection_name, points_selector=models.FilterSelector( filter=self._docs_filter(tenant_id, entry_ids) ), wait=True, ) def delete_tenant(self, tenant_id: str) -> None: self._ensure_collection() self.client.delete( collection_name=self.collection_name, points_selector=models.FilterSelector(filter=self._tenant_filter(tenant_id)), wait=True, ) def get_collection_dimension(self) -> int | None: """Return the actual vector size of the collection, or None if it does not exist.""" if not self.client.collection_exists(collection_name=self.collection_name): return None info = self.client.get_collection(collection_name=self.collection_name) return info.config.params.vectors.size def health_check(self) -> None: probe_vector = [1.0] + [0.0] * (self.vector_size - 1) self.upsert( VectorDocument( tenant_id=self.HEALTH_TENANT_ID, entry_id=self.HEALTH_ENTRY_ID, text=self.HEALTH_TEXT, embedding=probe_vector, embedding_metadata=EmbeddingMetadata( embedding_model="qdrant-health-check", embedding_dim=self.vector_size, ), ) ) results = self.query(self.HEALTH_TENANT_ID, probe_vector, 1) if not results or results[0].entry_id != self.HEALTH_ENTRY_ID: raise RuntimeError("Qdrant health probe query did not return probe document") self.delete_docs(self.HEALTH_TENANT_ID, [self.HEALTH_ENTRY_ID])