Spaces:
Sleeping
Sleeping
| """ | |
| Qdrant Vector Database Client | |
| ============================= | |
| Async client for vector operations on image embeddings. | |
| Uses AsyncQdrantClient for non-blocking I/O. | |
| """ | |
| import logging | |
| from typing import Optional, List, Dict, Any, Tuple | |
| from qdrant_client import AsyncQdrantClient | |
| from qdrant_client.http import models as qdrant_models | |
| from config import settings | |
| logger = logging.getLogger(__name__) | |
| class QdrantVectorClient: | |
| """ | |
| Qdrant client for image embedding operations. | |
| Handles: | |
| - Collection creation with proper vector config | |
| - Embedding upsert/retrieval | |
| - Similarity search (text-to-image and image-to-image) | |
| """ | |
| def __init__(self): | |
| self._client: Optional[AsyncQdrantClient] = None | |
| async def connect(self) -> None: | |
| """Initialize Qdrant connection.""" | |
| try: | |
| # Support both local and cloud Qdrant | |
| qdrant_secret = settings.QDRANT_API_KEY | |
| if qdrant_secret is not None: | |
| # Qdrant Cloud - use url for cloud endpoints | |
| api_key = qdrant_secret.get_secret_value() # pylint: disable=no-member | |
| # Check if host looks like a cloud endpoint (contains 'cloud' or 'qdrant.io') | |
| if 'cloud' in settings.QDRANT_HOST or 'qdrant.io' in settings.QDRANT_HOST: | |
| self._client = AsyncQdrantClient( | |
| url=f"https://{settings.QDRANT_HOST}:{settings.QDRANT_PORT}", | |
| api_key=api_key | |
| ) | |
| else: | |
| self._client = AsyncQdrantClient( | |
| host=settings.QDRANT_HOST, | |
| port=settings.QDRANT_PORT, | |
| api_key=api_key, | |
| https=True | |
| ) | |
| else: | |
| # Local/self-hosted | |
| self._client = AsyncQdrantClient( | |
| host=settings.QDRANT_HOST, | |
| port=settings.QDRANT_PORT | |
| ) | |
| # Ensure collection exists | |
| await self._ensure_collection() | |
| logger.info(f"Connected to Qdrant at {settings.QDRANT_HOST}:{settings.QDRANT_PORT}") | |
| except Exception as e: | |
| logger.error(f"Failed to connect to Qdrant: {e}") | |
| raise | |
| async def disconnect(self) -> None: | |
| """Close Qdrant connection.""" | |
| if self._client: | |
| await self._client.close() | |
| logger.info("Disconnected from Qdrant") | |
| async def _ensure_collection(self) -> None: | |
| """Create collection if it doesn't exist.""" | |
| collections = await self._client.get_collections() | |
| collection_names = [c.name for c in collections.collections] | |
| if settings.QDRANT_COLLECTION not in collection_names: | |
| await self._client.create_collection( | |
| collection_name=settings.QDRANT_COLLECTION, | |
| vectors_config=qdrant_models.VectorParams( | |
| size=settings.EMBEDDING_DIMENSION, | |
| distance=qdrant_models.Distance.COSINE | |
| ) | |
| ) | |
| logger.info(f"Created Qdrant collection: {settings.QDRANT_COLLECTION}") | |
| else: | |
| logger.info(f"Qdrant collection exists: {settings.QDRANT_COLLECTION}") | |
| # Ensure payload index on user_id for fast filtered search | |
| try: | |
| await self._client.create_payload_index( | |
| collection_name=settings.QDRANT_COLLECTION, | |
| field_name="user_id", | |
| field_schema=qdrant_models.PayloadSchemaType.KEYWORD, | |
| ) | |
| logger.info("Qdrant payload index on user_id ensured") | |
| except Exception: | |
| pass # index already exists | |
| def client(self) -> AsyncQdrantClient: | |
| """Get the Qdrant client.""" | |
| if not self._client: | |
| raise RuntimeError("Qdrant not connected. Call connect() first.") | |
| return self._client | |
| def is_connected(self) -> bool: | |
| """Check if client is connected.""" | |
| return self._client is not None | |
| # ========================================================================= | |
| # Vector Operations | |
| # ========================================================================= | |
| async def upsert_embedding( | |
| self, | |
| image_id: str, | |
| embedding: List[float], | |
| payload: Optional[Dict[str, Any]] = None | |
| ) -> bool: | |
| """ | |
| Insert or update an image embedding. | |
| Args: | |
| image_id: Unique image identifier | |
| embedding: Vector embedding (dimension must match collection config) | |
| payload: Additional metadata to store with the vector | |
| Returns: | |
| True if operation succeeded. | |
| """ | |
| try: | |
| # Validate embedding dimension | |
| if len(embedding) != settings.EMBEDDING_DIMENSION: | |
| logger.error( | |
| f"Embedding dimension mismatch: got {len(embedding)}, " | |
| f"expected {settings.EMBEDDING_DIMENSION}" | |
| ) | |
| return False | |
| point = qdrant_models.PointStruct( | |
| id=image_id, # Use image_id as point ID | |
| vector=embedding, | |
| payload=payload or {"image_id": image_id} | |
| ) | |
| await self.client.upsert( | |
| collection_name=settings.QDRANT_COLLECTION, | |
| points=[point] | |
| ) | |
| logger.debug(f"Upserted embedding for image: {image_id}") | |
| return True | |
| except Exception as e: | |
| logger.error(f"Failed to upsert embedding: {e}") | |
| return False | |
| async def get_embedding(self, image_id: str) -> Optional[List[float]]: | |
| """Retrieve embedding by image_id.""" | |
| try: | |
| result = await self.client.retrieve( | |
| collection_name=settings.QDRANT_COLLECTION, | |
| ids=[image_id], | |
| with_vectors=True | |
| ) | |
| if result: | |
| return result[0].vector | |
| return None | |
| except Exception as e: | |
| logger.warning(f"Failed to get embedding: {e}") | |
| return None | |
| async def delete_embedding(self, image_id: str) -> bool: | |
| """Delete embedding by image_id.""" | |
| try: | |
| await self.client.delete( | |
| collection_name=settings.QDRANT_COLLECTION, | |
| points_selector=qdrant_models.PointIdsList( | |
| points=[image_id] | |
| ) | |
| ) | |
| return True | |
| except Exception as e: | |
| logger.warning(f"Failed to delete embedding: {e}") | |
| return False | |
| async def search_similar( | |
| self, | |
| query_vector: List[float], | |
| limit: int = 10, | |
| score_threshold: Optional[float] = None, | |
| filter_conditions: Optional[Dict[str, Any]] = None, | |
| user_id: Optional[str] = None, | |
| ) -> List[Tuple[str, float, Dict[str, Any]]]: | |
| """ | |
| Search for similar images by embedding vector. | |
| Args: | |
| query_vector: Query embedding vector | |
| limit: Maximum results to return | |
| score_threshold: Minimum similarity score (0-1 for cosine) | |
| filter_conditions: Optional payload filters | |
| Returns: | |
| List of (image_id, score, payload) tuples. | |
| """ | |
| try: | |
| # Build filter if provided | |
| query_filter = None | |
| must_conditions = [] | |
| if filter_conditions: | |
| for key, value in filter_conditions.items(): | |
| must_conditions.append( | |
| qdrant_models.FieldCondition( | |
| key=key, | |
| match=qdrant_models.MatchValue(value=value) | |
| ) | |
| ) | |
| if user_id: | |
| must_conditions.append( | |
| qdrant_models.FieldCondition( | |
| key="user_id", | |
| match=qdrant_models.MatchValue(value=user_id) | |
| ) | |
| ) | |
| if must_conditions: | |
| query_filter = qdrant_models.Filter(must=must_conditions) | |
| # Use new Qdrant query_points API | |
| threshold = score_threshold or settings.SIMILARITY_THRESHOLD | |
| logger.info(f"Querying Qdrant with threshold={threshold}, limit={limit}") | |
| results = await self.client.query_points( | |
| collection_name=settings.QDRANT_COLLECTION, | |
| query=query_vector, | |
| limit=limit, | |
| score_threshold=threshold, | |
| query_filter=query_filter, | |
| with_payload=True | |
| ) | |
| logger.info(f"Qdrant returned {len(results.points)} points") | |
| return [ | |
| (point.id, point.score, point.payload or {}) | |
| for point in results.points | |
| ] | |
| except Exception as e: | |
| logger.error(f"Search failed: {e}") | |
| return [] | |
| async def search_by_image_id( | |
| self, | |
| image_id: str, | |
| limit: int = 10, | |
| exclude_self: bool = True, | |
| user_id: Optional[str] = None, | |
| ) -> List[Tuple[str, float, Dict[str, Any]]]: | |
| """ | |
| Find images similar to a given image. | |
| Args: | |
| image_id: Source image to find similar images for | |
| limit: Maximum results | |
| exclude_self: Whether to exclude the source image from results | |
| Returns: | |
| List of (image_id, score, payload) tuples. | |
| """ | |
| # Get the source embedding | |
| embedding = await self.get_embedding(image_id) | |
| if not embedding: | |
| logger.warning(f"No embedding found for image: {image_id}") | |
| return [] | |
| # Search for similar | |
| results = await self.search_similar( | |
| query_vector=embedding, | |
| limit=limit + 1 if exclude_self else limit, | |
| user_id=user_id, | |
| ) | |
| # Filter out self if needed | |
| if exclude_self: | |
| results = [(id, score, payload) for id, score, payload in results if id != image_id] | |
| results = results[:limit] | |
| return results | |
| async def batch_upsert( | |
| self, | |
| embeddings: List[Tuple[str, List[float], Dict[str, Any]]] | |
| ) -> int: | |
| """ | |
| Batch upsert multiple embeddings. | |
| Args: | |
| embeddings: List of (image_id, vector, payload) tuples | |
| Returns: | |
| Number of successfully upserted embeddings. | |
| """ | |
| if not embeddings: | |
| return 0 # Nothing to upsert | |
| try: | |
| points = [ | |
| qdrant_models.PointStruct( | |
| id=image_id, | |
| vector=vector, | |
| payload=payload or {"image_id": image_id} | |
| ) | |
| for image_id, vector, payload in embeddings | |
| ] | |
| await self.client.upsert( | |
| collection_name=settings.QDRANT_COLLECTION, | |
| points=points | |
| ) | |
| logger.info(f"Batch upserted {len(points)} embeddings") | |
| return len(points) | |
| except Exception as e: | |
| logger.error(f"Batch upsert failed: {e}") | |
| return 0 | |
| async def get_collection_info(self) -> Dict[str, Any]: | |
| """Get collection statistics.""" | |
| try: | |
| info = await self.client.get_collection(settings.QDRANT_COLLECTION) | |
| return { | |
| "name": settings.QDRANT_COLLECTION, | |
| "points_count": info.points_count, | |
| "status": info.status.name if hasattr(info.status, 'name') else str(info.status) | |
| } | |
| except Exception as e: | |
| logger.error(f"Failed to get collection info: {e}") | |
| return {} | |
| # Global client instance | |
| qdrant_client = QdrantVectorClient() | |