services/vector_db_service.py

"""
Vector database service for interacting with Qdrant
"""

from typing import List, Dict, Any

from fastapi import HTTPException # type: ignore
from qdrant_client import QdrantClient # type: ignore
from qdrant_client.models import Distance, PointStruct, VectorParams # type: ignore

class VectorDatabaseClient:
    """Class for interacting with Qdrant vector database"""
    
    def __init__(self, url: str, api_key: str, collection_name: str, embedding_size: int):
        self.url = url
        self.api_key = api_key
        self.collection_name = collection_name
        self.embedding_size = embedding_size
        self.client = QdrantClient(url=url, api_key=api_key)
        
    def ensure_collection_exists(self) -> None:
        """Ensure the Qdrant collection exists"""
        collections = self.client.get_collections()
        collection_names = [c.name for c in collections.collections]
        
        if self.collection_name not in collection_names:
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=self.embedding_size,
                    distance=Distance.COSINE
                )
            )
            print(f"✅ Collection '{self.collection_name}' created.")
        else:
            print(f"ℹ️ Collection '{self.collection_name}' already exists.")
    
    def add_image(self, image_id: str, embedding: List[float], payload: Dict[str, Any]) -> None:
        """Add an image embedding to the database"""
        self.client.upsert(
            collection_name=self.collection_name,
            points=[
                PointStruct(
                    id=image_id,
                    vector=embedding,
                    payload=payload
                )
            ]
        )
    
    def search_by_vector(self, embedding: List[float], limit: int = 1) -> List[Dict[str, Any]]:
        """Search for similar images using an embedding vector"""
        results = self.client.search(
            collection_name=self.collection_name,
            query_vector=embedding,
            limit=limit
        )
        
        return [
            {
                "id": r.id,
                "score": r.score,
                "payload": r.payload
            }
            for r in results
        ]
    
    def list_collections(self) -> List[str]:
        """List all collections in the database"""
        return [c.name for c in self.client.get_collections().collections]