QdrantManager.py

import httpx
import json
import uuid
from typing import List, Dict, Optional
from Configurator import AppConfig


class QdrantManager:
    """
    Manages Qdrant collection operations including creation, deletion, and data insertion.
    """
    
    def __init__(self, config: AppConfig = None):
        """
        Initialize the Qdrant manager.
        
        Args:
            config (AppConfig, optional): Application configuration. If None, loads from Configurator.
        """
        if config is None:
            from Configurator import get_config
            config = get_config()
        
        self.config = config
        self.qdrant_url = config.qdrant_url
        self.qdrant_api_key = config.qdrant_api_key
        self.collection_name = config.qdrant_collection
        
        # Setup headers for Qdrant API
        self.headers = {
            "Content-Type": "application/json",
            "api-key": self.qdrant_api_key
        }
    
    async def create_collection(self, vector_size: int = 3072) -> Dict:
        """
        Create a new Qdrant collection.
        
        Args:
            vector_size (int): Size of the embedding vectors
            
        Returns:
            Dict: Creation result
        """
        create_url = f"{self.qdrant_url}/collections/{self.collection_name}"
        
        collection_config = {
            "vectors": {
                "size": vector_size,
                "distance": "Cosine"
            }
        }
        
        try:
            async with httpx.AsyncClient() as client:
                response = await client.put(
                    create_url,
                    headers=self.headers,
                    json=collection_config,
                    timeout=30.0
                )
                
                if response.status_code in [200, 201]:
                    return {
                        "success": True,
                        "message": f"Collection '{self.collection_name}' created successfully",
                        "vector_size": vector_size
                    }
                elif response.status_code in [400, 409]:
                    # Collection already exists (400 or 409)
                    return {
                        "success": True,
                        "message": f"Collection '{self.collection_name}' already exists",
                        "vector_size": vector_size
                    }
                else:
                    return {
                        "success": False,
                        "error": f"Failed to create collection: HTTP {response.status_code} - {response.text}"
                    }
                    
        except Exception as e:
            return {
                "success": False,
                "error": f"Collection creation failed: {str(e)}"
            }
    
    async def delete_collection(self) -> Dict:
        """
        Delete the Qdrant collection.
        
        Returns:
            Dict: Deletion result
        """
        delete_url = f"{self.qdrant_url}/collections/{self.collection_name}"
        
        try:
            async with httpx.AsyncClient() as client:
                response = await client.delete(
                    delete_url,
                    headers=self.headers,
                    timeout=30.0
                )
                
                if response.status_code in [200, 202]:
                    return {
                        "success": True,
                        "message": f"Collection '{self.collection_name}' deleted successfully"
                    }
                else:
                    return {
                        "success": False,
                        "error": f"Failed to delete collection: HTTP {response.status_code} - {response.text}"
                    }
                    
        except Exception as e:
            return {
                "success": False,
                "error": f"Collection deletion failed: {str(e)}"
            }
    
    async def upsert_points(self, points: List[Dict]) -> Dict:
        """
        Insert or update points in the collection.
        
        Args:
            points (List[Dict]): List of points to insert/update
            
        Returns:
            Dict: Upsert result
        """
        upsert_url = f"{self.qdrant_url}/collections/{self.collection_name}/points"
        
        upsert_payload = {
            "points": points
        }
        
        try:
            async with httpx.AsyncClient() as client:
                response = await client.put(
                    upsert_url,
                    headers=self.headers,
                    json=upsert_payload,
                    timeout=60.0
                )
                
                if response.status_code in [200, 202]:
                    return {
                        "success": True,
                        "message": f"Successfully upserted {len(points)} points",
                        "points_count": len(points)
                    }
                else:
                    return {
                        "success": False,
                        "error": f"Failed to upsert points: HTTP {response.status_code} - {response.text}"
                    }
                    
        except Exception as e:
            return {
                "success": False,
                "error": f"Points upsert failed: {str(e)}"
            }
    
    async def get_collection_info(self) -> Dict:
        """
        Get information about the collection.
        
        Returns:
            Dict: Collection information
        """
        info_url = f"{self.qdrant_url}/collections/{self.collection_name}"
        
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    info_url,
                    headers=self.headers,
                    timeout=10.0
                )
                
                if response.status_code == 200:
                    data = response.json()
                    # Extract points count from the nested structure
                    result = data.get('result', {})
                    points_count = result.get('points_count', 0)
                    return {
                        "success": True,
                        "points_count": points_count,
                        "status": result.get('status', 'unknown'),
                        "vector_size": result.get('config', {}).get('params', {}).get('vectors', {}).get('size', 0),
                        "distance": result.get('config', {}).get('params', {}).get('vectors', {}).get('distance', 'unknown'),
                        "raw_data": data
                    }
                else:
                    return {
                        "success": False,
                        "error": f"Failed to get collection info: HTTP {response.status_code}"
                    }
                    
        except Exception as e:
            return {
                "success": False,
                "error": f"Collection info retrieval failed: {str(e)}"
            }
    
    async def clear_collection(self) -> Dict:
        """
        Clear all points from the collection.
        
        Returns:
            Dict: Clear result
        """
        clear_url = f"{self.qdrant_url}/collections/{self.collection_name}/points/delete"
        
        clear_payload = {
            "filter": {
                "must": [
                    {
                        "key": "id",
                        "match": {
                            "any": []
                        }
                    }
                ]
            }
        }
        
        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    clear_url,
                    headers=self.headers,
                    json=clear_payload,
                    timeout=30.0
                )
                
                if response.status_code in [200, 202]:
                    return {
                        "success": True,
                        "message": f"Collection '{self.collection_name}' cleared successfully"
                    }
                else:
                    return {
                        "success": False,
                        "error": f"Failed to clear collection: HTTP {response.status_code} - {response.text}"
                    }
                    
        except Exception as e:
            return {
                "success": False,
                "error": f"Collection clear failed: {str(e)}"
            }
    
    def create_point(self, vector: List[float], payload: Dict, point_id: Optional[str] = None) -> Dict:
        """
        Create a point for insertion into Qdrant.
        
        Args:
            vector (List[float]): Embedding vector
            payload (Dict): Point payload (metadata)
            point_id (str, optional): Point ID (generated if not provided)
            
        Returns:
            Dict: Point data
        """
        if point_id is None:
            point_id = str(uuid.uuid4())
        
        return {
            "id": point_id,
            "vector": vector,
            "payload": payload
        }
    
    async def validate_connection(self) -> Dict:
        """
        Validate connection to Qdrant.
        
        Returns:
            Dict: Validation results
        """
        validation_results = {
            "valid": True,
            "errors": [],
            "warnings": [],
            "collection_info": None
        }
        
        try:
            # Test connection by getting collection info
            collection_info = await self.get_collection_info()
            
            if collection_info.get("success") is False:
                validation_results["valid"] = False
                validation_results["errors"].append(f"Connection failed: {collection_info.get('error')}")
            else:
                validation_results["collection_info"] = collection_info
                
                # Check if collection has points
                points_count = collection_info.get("points_count", 0)
                if points_count == 0:
                    validation_results["warnings"].append("Collection is empty")
                
        except Exception as e:
            validation_results["valid"] = False
            validation_results["errors"].append(f"Connection validation failed: {str(e)}")
        
        return validation_results