Multimedika/Bot_Development

1from llama_index.core import VectorStoreIndex2from llama_index.core import StorageContext3# from llama_index.core import Settings4from pinecone import Pinecone, ServerlessSpec5from llama_index.vector_stores.pinecone import PineconeVectorStore6from fastapi import HTTPException, status7from fastapi.responses import JSONResponse8from config import PINECONE_CONFIG9from math import ceil10import numpy as np11import logging12 13 14class IndexManager:15    # def __init__(self, index_name: str = "chatbook-bot"):16    def __init__(self, index_name: str = "multimodal-index"):17        self.vector_index = None18        self.index_name = index_name19        self.client = self._get_pinecone_client()20        self.pinecone_index = self._create_pinecone_index()21 22    def _get_pinecone_client(self):23        """Initialize and return the Pinecone client."""24        # api_key = os.getenv("PINECONE_API_KEY")25        api_key = PINECONE_CONFIG.PINECONE_API_KEY26        if not api_key:27            raise ValueError(28                "Pinecone API key is missing. Please set it in environment variables."29            )30        return Pinecone(api_key=api_key)31 32    def _create_pinecone_index(self):33        """Create Pinecone index if it doesn't already exist."""34        if self.index_name not in self.client.list_indexes().names():35            self.client.create_index(36                name=self.index_name,37                dimension=3072,38                metric="cosine",39                spec=ServerlessSpec(cloud="aws", region="us-east-1"),40            )41        return self.client.Index(self.index_name)42 43    def _initialize_vector_store(self) -> StorageContext:44        """Initialize and return the vector store with the Pinecone index."""45        vector_store = PineconeVectorStore(pinecone_index=self.pinecone_index)46        return StorageContext.from_defaults(vector_store=vector_store)47 48 49    def build_indexes(self, nodes):50        """Build vector and tree indexes from nodes."""51        try:52            storage_context = self._initialize_vector_store()53            self.vector_index = VectorStoreIndex(nodes, storage_context=storage_context)54            55        except HTTPException as http_exc:56            raise http_exc  # Re-return JSONResponses to ensure FastAPI handles them57        58        except Exception as e:59            print("Error building index : ",e)60            raise JSONResponse(61                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,62                content=f"Error loading existing indexes: {str(e)}"63            )64            65    def get_ids_from_query(self, input_vector, title):66        print("Searching Pinecone...")67        print(title)68 69        new_ids = set()  # Initialize new_ids outside the loop70 71        while True:72            results = self.pinecone_index.query(73                vector=input_vector,74                top_k=10000,75                filter={76                    "title": {"$eq": f"{title}"},77                },78            )79            80            ids = set()81            for result in results['matches']:82                ids.add(result['id'])83            # Check if there's any overlap between ids and new_ids84            if ids.issubset(new_ids):85                break86            else:87                new_ids.update(ids)  # Add all new ids to new_ids88 89        return new_ids90 91           92    def get_all_ids_from_index(self, title):93        num_dimensions = 153694 95        num_vectors = self.pinecone_index.describe_index_stats(96        )["total_vector_count"]97        98        input_vector = np.random.rand(num_dimensions).tolist()99        ids = self.get_ids_from_query(input_vector, title)100 101        return ids102    103    def delete_vector_database(self, title):104        try :105            batch_size = 1000106            all_ids = self.get_all_ids_from_index(title)107            all_ids = list(all_ids)108            109            # Split ids into chunks of batch_size110            num_batches = ceil(len(all_ids) / batch_size)111            112            for i in range(num_batches):113                # Fetch a batch of IDs114                batch_ids = all_ids[i * batch_size: (i + 1) * batch_size]115                self.pinecone_index.delete(ids=batch_ids)116                logging.info(f"delete from id {i * batch_size} to {(i + 1) * batch_size} successful")117        except Exception as e:118            return JSONResponse(status_code=500, content="An error occurred while delete metadata")    119 120    def update_vector_database(self, current_reference, new_reference):121        122        reference = new_reference123        124        all_ids = self.get_all_ids_from_index(current_reference['title'])125        all_ids = list(all_ids)126        127        for id in all_ids:128            self.pinecone_index.update(129                id=id,130                set_metadata=reference131            )132 133    def load_existing_indexes(self):134        """Load existing indexes from Pinecone."""135        try:136            client = self._get_pinecone_client()137            pinecone_index = client.Index(self.index_name)138            139            vector_store = PineconeVectorStore(pinecone_index=pinecone_index)140            retriever = VectorStoreIndex.from_vector_store(vector_store)141            142            return retriever143        except Exception as e:144            return JSONResponse(145                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,146                content=f"Error loading existing indexes: {str(e)}"147            )148