done more refactoring

config.py

@@ -2,7 +2,8 @@ import os
 from dotenv import load_dotenv
 from rag_app.database.db_handler import DataBaseHandler
 from langchain_huggingface import HuggingFaceEndpoint
-from langchain_huggingface import HuggingFaceHubEmbeddings
+# from langchain_huggingface import HuggingFaceHubEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings
 
 load_dotenv()
 
@@ -16,7 +17,15 @@ HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
 
 
-embeddings = HuggingFaceHubEmbeddings(repo_id=EMBEDDING_MODEL)
+# embeddings = HuggingFaceHubEmbeddings(repo_id=EMBEDDING_MODEL)
+
+model_kwargs = {'device': 'cpu'}
+encode_kwargs = {'normalize_embeddings': False}
+embeddings = HuggingFaceEmbeddings(
+    model_name=EMBEDDING_MODEL,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs
+)
 
 db = DataBaseHandler()
 

rag_app/knowledge_base/utils.py

@@ -1,38 +1,30 @@
 from langchain_core.documents import Document
 from chains import generate_document_summary_prompt
-from config import SEVEN_B_LLM_MODEL
 # embeddings functions
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.embeddings.sentence_transformer import (
     SentenceTransformerEmbeddings,
 )
 import time
-from langchain_core.documents import Document
-from config import EMBEDDING_MODEL
+from langchain_core.language_models import BaseChatModel
 from langchain.retrievers import VectorStoreRetriever
 from langchain_core.vectorstores import VectorStoreRetriever
 # vectorization functions
 from langchain_community.vectorstores import FAISS
 from langchain_community.vectorstores import Chroma
 from langchain_community.retrievers import BM25Retriever
+from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
 
-from rag_app.knowledge_base.utils import create_embeddings
-from rag_app.utils.generate_summary import generate_description, generate_keywords
-
-import time
-import os
-
-from config import FAISS_INDEX_PATH
 
 from pathlib import Path
 from langchain_community.vectorstores import FAISS
 from dotenv import load_dotenv
 import os
-from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
 import requests
-from langchain_community.vectorstores import Chroma
-
 
+from rag_app.knowledge_base.utils import create_embeddings
+from rag_app.utils.generate_summary import generate_description, generate_keywords
+from config import EMBEDDING_MODEL, FAISS_INDEX_PATH, SEVEN_B_LLM_MODEL
 
 def create_embeddings(
         docs: list[Document], 
@@ -76,7 +68,8 @@ def create_embeddings(
 
 
 def generate_document_summaries(
-        docs: list[Document]
+        docs: list[Document],
+        llm:BaseChatModel= SEVEN_B_LLM_MODEL,
     ) -> list[Document]:
     """
     Generates summaries for a list of Document objects and updates their metadata with the summaries.
@@ -99,7 +92,7 @@ def generate_document_summaries(
     
     for doc in new_docs:
         
-        genrate_summary_chain = generate_document_summary_prompt | SEVEN_B_LLM_MODEL
+        genrate_summary_chain = generate_document_summary_prompt | llm
         summary = genrate_summary_chain.invoke(
             {"document":str(doc.metadata)}
         )        
@@ -155,105 +148,4 @@ def build_vector_store(
     result = f"built vectore store at {FAISS_INDEX_PATH}"
     return result
 
-def get_reranked_docs_faiss(
-    query:str, 
-    path_to_db:str, 
-    embedding_model:str,
-    hf_api_key:str, 
-    num_docs:int=5
-    ) -> list:
-    """ Re-ranks the similarity search results and returns top-k highest ranked docs
-
-    Args:
-        query (str): The search query
-        path_to_db (str): Path to the vectorstore database
-        embedding_model (str): Embedding model used in the vector store
-        num_docs (int): Number of documents to return
-    
-    Returns: A list of documents with the highest rank
-    """
-    assert num_docs <= 10, "num_docs should be less than similarity search results"
-    
-    embeddings = HuggingFaceInferenceAPIEmbeddings(
-        api_key=hf_api_key,
-        model_name=embedding_model
-        )
-    
-    # Load the vectorstore database
-    db = FAISS.load_local(
-        folder_path=path_to_db,
-        embeddings=embeddings,
-        allow_dangerous_deserialization=True
-        )
-    
-    # Get 10 documents based on similarity search
-    docs =  db.similarity_search(query=query, k=10)
-
-    # Add the page_content, description and title together
-    passages = [doc.page_content + "\n" + doc.metadata.get('title', "") +"\n"+ doc.metadata.get('description', "") 
-                for doc in docs]
-    
-    # Prepare the payload
-    inputs = [{"text": query, "text_pair": passage} for passage in passages]
-
-    API_URL = "https://api-inference.huggingface.co/models/deepset/gbert-base-germandpr-reranking"
-    headers = {"Authorization": f"Bearer {hf_api_key}"}
-
-    response = requests.post(API_URL, headers=headers, json=inputs)
-    scores = response.json()
-    
-    try:
-        relevance_scores = [item[1]['score'] for item in scores]
-    except ValueError as e:
-        print('Could not get the relevance_scores -> something might be wrong with the json output')
-        return 
-    
-    if relevance_scores:
-        ranked_results = sorted(zip(docs, passages, relevance_scores), key=lambda x: x[2], reverse=True)
-        top_k_results = ranked_results[:num_docs]
-        return [doc for doc, _, _ in top_k_results]
-    
-    
-def get_reranked_docs_chroma(query:str, 
-                      path_to_db:str, 
-                      embedding_model:str,
-                      hf_api_key:str,
-                      reranking_hf_url:str = "https://api-inference.huggingface.co/models/sentence-transformers/all-mpnet-base-v2", 
-                      num_docs:int=5) -> list:
-    """ Re-ranks the similarity search results and returns top-k highest ranked docs
-
-        Args:
-            query (str): The search query
-            path_to_db (str): Path to the vectorstore database
-            embedding_model (str): Embedding model used in the vector store
-            num_docs (int): Number of documents to return
-        
-        Returns: A list of documents with the highest rank
-    """
-    embeddings = HuggingFaceInferenceAPIEmbeddings(api_key=hf_api_key,
-                                                   model_name=embedding_model)
-    # Load the vectorstore database
-    db = Chroma(persist_directory=path_to_db, embedding_function=embeddings)
-    
-    # Get k documents based on similarity search
-    sim_docs =  db.similarity_search(query=query, k=10)
-
-    passages = [doc.page_content for doc in sim_docs]
-    
-    # Prepare the payload
-    payload = {"inputs": 
-               {"source_sentence": query,
-	            "sentences": passages}}
-    
-    headers = {"Authorization": f"Bearer {hf_api_key}"}
-
-    response = requests.post(url=reranking_hf_url, headers=headers, json=payload)
-    print(f'{response = }')
-    if response.status_code != 200:
-        print('Something went wrong with the response')
-        return
-    
-    similarity_scores = response.json()
-    ranked_results = sorted(zip(sim_docs, passages, similarity_scores), key=lambda x: x[2], reverse=True)
-    top_k_results = ranked_results[:num_docs]
-    return [doc for doc, _, _ in top_k_results]
+    

rag_app/vector_store_handler/vectorstores.py

@@ -146,32 +146,43 @@ class ChromaVectorStore(BaseVectorStore):
         query:str,
         num_docs:int=5
         ):
+        """ Re-ranks the similarity search results and returns top-k highest ranked docs
+
+        Args:
+            query (str): The search query
+            path_to_db (str): Path to the vectorstore database
+            embedding_model (str): Embedding model used in the vector store
+            num_docs (int): Number of documents to return
         
-        # Get 10 documents based on similarity search
-        docs = self.vectorstore.similarity_search(query=query, k=10)
+        Returns: A list of documents with the highest rank
+        """
+        
+        # Get k documents based on similarity search
+        sim_docs =  self.vectorstore.similarity_search(query=query, k=10)
         
         # Add the page_content, description and title together
-        passages = [doc.page_content + "\n" + doc.metadata.get('title', "") +"\n"+ doc.metadata.get('description', "") 
-                for doc in docs]
+        passages = [doc.page_content for doc in sim_docs]
+        
         # Prepare the payload
-        inputs = [{"text": query, "text_pair": passage} for passage in passages]
-
-        API_URL = "https://api-inference.huggingface.co/models/deepset/gbert-base-germandpr-reranking"
+        payload = {"inputs": 
+                {"source_sentence": query,
+                    "sentences": passages}}
+        
         headers = {"Authorization": f"Bearer {HUGGINGFACEHUB_API_TOKEN}"}
+        reranking_hf_url:str = "https://api-inference.huggingface.co/models/sentence-transformers/all-mpnet-base-v2"
 
-        response = requests.post(API_URL, headers=headers, json=inputs)
-        scores = response.json()
+        response = requests.post(url=reranking_hf_url, headers=headers, json=payload)
+        print(f'{response = }')
+        if response.status_code != 200:
+            print('Something went wrong with the response')
+            return
         
-        try:
-            relevance_scores = [item[1]['score'] for item in scores]
-        except ValueError as e:
-            print('Could not get the relevance_scores -> something might be wrong with the json output')
-            return 
+        similarity_scores = response.json()
+        ranked_results = sorted(zip(sim_docs, passages, similarity_scores), key=lambda x: x[2], reverse=True)
+        top_k_results = ranked_results[:num_docs]
+        return [doc for doc, _, _ in top_k_results]
         
-        if relevance_scores:
-            ranked_results = sorted(zip(docs, passages, relevance_scores), key=lambda x: x[2], reverse=True)
-            top_k_results = ranked_results[:num_docs]
-            return [doc for doc, _, _ in top_k_results]
+
 
 class FAISSVectorStore(BaseVectorStore):
     """
@@ -187,7 +198,7 @@ class FAISSVectorStore(BaseVectorStore):
         """
         self.vectorstore = FAISS.from_documents(texts, self.embeddings)
 
-    def load_existing_vectorstore(self):
+    def load_existing_vectorstore(self,allow_dangerous_deserialization:bool=False):
         """
         Load an existing FAISS vector store from the persist directory.
 
@@ -195,7 +206,7 @@ class FAISSVectorStore(BaseVectorStore):
             ValueError: If persist_directory is not set.
         """
         if self.persist_directory:
-            self.vectorstore = FAISS.load_local(self.persist_directory, self.embeddings, allow_dangerous_deserialization=True)
+            self.vectorstore = FAISS.load_local(self.persist_directory, self.embeddings, allow_dangerous_deserialization)
         else:
             raise ValueError("Persist directory is required for loading FAISS.")
 

tests/vector_store_handler/test_vectorstores.py

@@ -1,14 +1,16 @@
 import unittest
 from unittest.mock import MagicMock, patch
-from langchain.embeddings import OpenAIEmbeddings
-from langchain.schema import Document
+# from langchain.embeddings import OpenAIEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings
+# from langchain.schema import Document
+from langchain_core.documents import Document
 
 # Update the import to reflect your project structure
 from rag_app.vector_store_handler.vectorstores import BaseVectorStore, ChromaVectorStore, FAISSVectorStore
 
 class TestBaseVectorStore(unittest.TestCase):
     def setUp(self):
-        self.embedding_model = MagicMock(spec=OpenAIEmbeddings)
+        self.embedding_model = MagicMock(spec=HuggingFaceEmbeddings)
         self.base_store = BaseVectorStore(self.embedding_model, "test_dir")
 
     def test_init(self):
@@ -34,7 +36,7 @@ class TestBaseVectorStore(unittest.TestCase):
 
 class TestChromaVectorStore(unittest.TestCase):
     def setUp(self):
-        self.embedding_model = MagicMock(spec=OpenAIEmbeddings)
+        self.embedding_model = MagicMock(spec=HuggingFaceEmbeddings)
         self.chroma_store = ChromaVectorStore(self.embedding_model, "test_dir")
 
     @patch('rag_app.vector_store_handler.vectorstores.Chroma')
@@ -62,7 +64,7 @@ class TestChromaVectorStore(unittest.TestCase):
 
 class TestFAISSVectorStore(unittest.TestCase):
     def setUp(self):
-        self.embedding_model = MagicMock(spec=OpenAIEmbeddings)
+        self.embedding_model = MagicMock(spec=HuggingFaceEmbeddings)
         self.faiss_store = FAISSVectorStore(self.embedding_model, "test_dir")
 
     @patch('rag_app.vector_store_handler.vectorstores.FAISS')