Create retriever/vector_store_manager.py

retriever/vector_store_manager.py

@@ -0,0 +1,98 @@
+import os
+import logging
+from config.config import ConfigConstants
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+
+class VectorStoreManager:
+    def __init__(self, embedding_path="embeddings.faiss"):
+        """
+        Initialize the vector store manager.
+        
+        Args:
+            embedding_path (str): Path to save/load the FAISS index.
+        """
+        self.embedding_path = embedding_path
+        self.embedding_model = HuggingFaceEmbeddings(model_name=ConfigConstants.EMBEDDING_MODEL_NAME)
+        self.vector_store = self._initialize_vector_store()
+
+    def _initialize_vector_store(self):
+        """Initialize or load the FAISS vector store."""
+        if os.path.exists(self.embedding_path):
+            logging.info("Loading embeddings from local file")
+            return FAISS.load_local(
+                self.embedding_path,
+                self.embedding_model,
+                allow_dangerous_deserialization=True
+            )
+        else:
+            '''logging.info("Creating new vector store")
+            # Return an empty vector store; it will be populated when documents are added
+            return FAISS.from_texts(
+                texts=[""],  # Dummy text to initialize
+                embedding=self.embedding_model,
+                metadatas=[{"source": "init", "doc_id": "init"}]
+            )'''
+            logging.info("Creating new vector store (unpopulated)")
+            return None 
+
+    def add_documents(self, documents):
+        """
+        Add new documents to the vector store and save it.
+        
+        Args:
+            documents (list): List of dictionaries with 'text', 'source', and 'doc_id'.
+        """
+        if not documents:
+            return
+
+        texts = [doc['text'] for doc in documents]
+        metadatas = [{'source': doc['source'], 'doc_id': doc['doc_id']} for doc in documents]
+
+        logging.info("Adding new documents to vector store")
+        
+        if not self.vector_store:
+            self.vector_store = FAISS.from_texts(
+                texts=texts,
+                embedding=self.embedding_model,
+                metadatas=metadatas
+            )
+        else:
+            self.vector_store.add_texts(texts=texts, metadatas=metadatas)
+
+        self.vector_store.save_local(self.embedding_path)
+        logging.info(f"Vector store updated and saved to {self.embedding_path}")
+
+    def search(self, query, doc_id, k=10):
+        """
+        Search the vector store for relevant chunks, filtered by doc_id.
+        
+        Args:
+            query (str): The user's query.
+            doc_id (str): The document ID to filter by.
+            k (int): Number of results to return.
+        
+        Returns:
+            list: List of relevant document chunks with metadata and scores.
+        """
+        if not self.vector_store:
+            return []
+
+        try:
+            query = " ".join(query.lower().split())
+            # Define a filter function to match doc_id
+            filter_fn = lambda metadata: metadata['doc_id'] == doc_id
+            
+            # Perform similarity search with filter
+            results = self.vector_store.similarity_search_with_score(
+                query=query,
+                k=k,
+                filter=filter_fn
+            )
+            
+            # Format results
+            return [{'text': doc.page_content, 'metadata': doc.metadata, 'score': score} for doc, score in results]
+        
+        except Exception as e:
+            logging.error(f"Error during vector store search: {str(e)}")
+            return []