ruby2210/rag-chatbot

1"""2Qdrant vector database client and collection management for the RAG Chatbot application.3"""4from qdrant_client import QdrantClient5from qdrant_client.http import models6from typing import List, Dict, Any, Optional7from .config import settings8import uuid9 10 11class VectorDB:12    """13    Wrapper class for Qdrant client with collection management functions.14    """15    def __init__(self):16        self.client = QdrantClient(17            url=settings.QDRANT_URL,18            api_key=settings.QDRANT_API_KEY,19            prefer_grpc=False  # Use HTTP for better compatibility20        )21        self.collection_name = settings.QDRANT_COLLECTION_NAME22 23    def create_collection(self, vector_size: int = 1024) -> bool:24        """25        Create the collection for storing book content embeddings.26        Vector size is set to 1024 which is the default for Cohere embeddings.27        """28        try:29            # Check if collection already exists30            collections = self.client.get_collections()31            collection_names = [c.name for c in collections.collections]32 33            if self.collection_name not in collection_names:34                self.client.create_collection(35                    collection_name=self.collection_name,36                    vectors_config=models.VectorParams(37                        size=vector_size,38                        distance=models.Distance.COSINE39                    )40                )41                print(f"Created collection: {self.collection_name}")42            else:43                print(f"Collection {self.collection_name} already exists")44 45            return True46        except Exception as e:47            print(f"Error creating collection: {e}")48            return False49 50    def add_embeddings(self, texts: List[str], metadata: List[Dict[str, Any]], embeddings: List[List[float]] = None) -> bool:51        """52        Add embeddings to the collection.53        Each text is stored with its corresponding metadata.54        If embeddings are provided, they are used; otherwise, placeholder vectors are used.55        """56        try:57            # Generate unique IDs for each text58            ids = [str(uuid.uuid4()) for _ in texts]59 60            # If embeddings are not provided, create placeholder embeddings61            if embeddings is None:62                embeddings = [[0.0] * 1536 for _ in texts]  # Placeholder embeddings63 64            # Validate that embeddings and texts have the same length65            if len(embeddings) != len(texts):66                raise ValueError(f"Number of embeddings ({len(embeddings)}) does not match number of texts ({len(texts)})")67 68            # Add points to the collection69            self.client.upsert(70                collection_name=self.collection_name,71                points=[72                    models.PointStruct(73                        id=ids[i],74                        vector=embeddings[i],75                        payload={76                            "text": texts[i],77                            "metadata": metadata[i]78                        }79                    ) for i in range(len(texts))80                ]81            )82            return True83        except ValueError as ve:84            print(f"Value error adding embeddings: {ve}")85            return False86        except Exception as e:87            print(f"Error adding embeddings: {e}")88            return False89 90    def search_similar(self, query_vector: List[float], limit: int = 5) -> List[Dict[str, Any]]:91        """92        Search for similar vectors to the query vector.93        Returns the most similar text chunks with their metadata.94        """95        try:96            results = self.client.search(97                collection_name=self.collection_name,98                query_vector=query_vector,99                limit=limit100            )101 102            return [103                {104                    "text": result.payload["text"],105                    "metadata": result.payload["metadata"],106                    "score": result.score107                }108                for result in results109            ]110        except Exception as e:111            print(f"Error searching vectors: {e}")112            return []113 114    def delete_collection(self) -> bool:115        """116        Delete the collection (useful for re-indexing).117        """118        try:119            self.client.delete_collection(collection_name=self.collection_name)120            return True121        except Exception as e:122            print(f"Error deleting collection: {e}")123            return False124 125 126# Create a singleton instance127vector_db = VectorDB()