attempt to refactor module

config.py

@@ -2,6 +2,7 @@ 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
 
 load_dotenv()
 
@@ -11,8 +12,12 @@ EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
 SEVEN_B_LLM_MODEL = os.getenv("SEVEN_B_LLM_MODEL")
 BERT_MODEL = os.getenv("BERT_MODEL")
 FAISS_INDEX_PATH = os.getenv("FAISS_INDEX_PATH")
+HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
 
+
+embeddings = HuggingFaceHubEmbeddings(repo_id=EMBEDDING_MODEL)
+
 db = DataBaseHandler()
 
 db.create_all_tables()

rag_app/knowledge_base/build_vector_store.py

@@ -1,57 +0,0 @@
-# vectorization functions
-from langchain_community.vectorstores import FAISS
-from langchain_community.vectorstores import Chroma
-from langchain_community.retrievers import BM25Retriever
-
-from rag_app.knowledge_base.create_embedding import create_embeddings
-from rag_app.utils.generate_summary import generate_description, generate_keywords
-
-import time
-import os
-
-from config import FAISS_INDEX_PATH
-
-def build_vector_store(
-        docs: list, 
-        embedding_model: str, 
-        new_db:bool=False, 
-        chunk_size:int=500, 
-        chunk_overlap:int=50,
-        ):
-    """
-
-    """
-
-    embeddings,chunks = create_embeddings(
-        docs, 
-        chunk_size, 
-        chunk_overlap, 
-        embedding_model
-        )
-
-    #load chunks into vector store
-    print(f'Loading chunks into faiss vector store ...')
-    
-    st = time.time()
-    if new_db:
-        db_faiss = FAISS.from_documents(chunks, embeddings)
-        bm25_retriever = BM25Retriever.from_documents(chunks)
-    else:
-        db_faiss = FAISS.add_documents(chunks, embeddings)
-        bm25_retriever = BM25Retriever.add_documents(chunks)
-        
-    db_faiss.save_local(FAISS_INDEX_PATH)
-    et = time.time() - st
-    print(f'Time taken: {et} seconds.')
-
-    print(f'Loading chunks into chroma vector store ...')
-    
-    st = time.time()
-    persist_directory='./vectorstore/chroma-insurance-agent-1500'
-    db_chroma = Chroma.from_documents(chunks, embeddings, persist_directory=persist_directory)
-    et = time.time() - st
-    
-    print(f'Time taken: {et} seconds.')
-    result = f"built vectore store at {FAISS_INDEX_PATH}"
-    return result
-

rag_app/knowledge_base/get_db_retriever.py

@@ -1,61 +0,0 @@
-# retriever and qa_chain function
-
-# HF libraries
-from langchain.llms import HuggingFaceHub
-from langchain_huggingface import HuggingFaceHubEmbeddings
-# vectorestore
-from langchain_community.vectorstores import FAISS
-# retrieval chain
-from langchain.chains import RetrievalQA
-# prompt template
-from langchain.prompts import PromptTemplate
-from langchain.memory import ConversationBufferMemory
-from config import EMBEDDING_MODEL, VECTOR_DATABASE_LOCATION
-
-
-def get_db_retriever():
-    """
-    Creates and returns a retriever object based on a FAISS vector database.
-
-    This function initializes an embedding model and loads a pre-existing FAISS
-    vector database from a local location. It then creates a retriever from this
-    database.
-
-    Returns:
-    --------
-    retriever : langchain.vectorstores.FAISS.VectorStoreRetriever
-        A retriever object that can be used to fetch relevant documents from the
-        vector database.
-
-    Global Variables Used:
-    ----------------------
-    EMBEDDING_MODEL : str
-        The identifier for the Hugging Face Hub embedding model to be used.
-    VECTOR_DATABASE_LOCATION : str
-        The local path where the FAISS vector database is stored.
-
-    Dependencies:
-    -------------
-    - langchain_huggingface.HuggingFaceHubEmbeddings
-    - langchain_community.vectorstores.FAISS
-
-    Note:
-    -----
-    This function assumes that a FAISS vector database has already been created
-    and saved at the location specified by VECTOR_DATABASE_LOCATION.
-    """
-
-    # Initialize the embedding model
-    embeddings = HuggingFaceHubEmbeddings(repo_id=EMBEDDING_MODEL)
-    
-    # Load the FAISS vector database from the local storage
-    db = FAISS.load_local(
-        VECTOR_DATABASE_LOCATION,
-        embeddings,
-    )
-
-    # Create and return a retriever from the loaded database
-    retriever = db.as_retriever()
-    
-    return retriever
-

rag_app/knowledge_base/hybrid_search.py

@@ -1,63 +0,0 @@
-from pathlib import Path
-from langchain_community.vectorstores import FAISS
-from dotenv import load_dotenv
-import os
-from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
-from langchain.retrievers import EnsembleRetriever
-from langchain_community.retrievers import BM25Retriever
-
-
-def get_hybrid_search_results(query:str,
-                              path_to_db:str,
-                              embedding_model:str,
-                              hf_api_key:str,
-                              num_docs:int=5) -> list:
-    """ Uses an ensemble retriever of BM25 and FAISS to return k num documents
-
-        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
-            List of documents
-    
-    """
-    
-    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)
-
-    all_docs = db.similarity_search("", k=db.index.ntotal)
-
-    bm25_retriever = BM25Retriever.from_documents(all_docs)
-    bm25_retriever.k = num_docs  # How many results you want
-
-    faiss_retriever = db.as_retriever(search_kwargs={'k': num_docs})
-
-    ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, faiss_retriever],
-                                           weights=[0.5,0.5])
-    
-    results = ensemble_retriever.invoke(input=query) 
-    return results
-
-
-if __name__ == "__main__":
-    query = "Haustierversicherung"
-    HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
-    EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
-    
-    path_to_vector_db = Path("..")/'vectorstore/faiss-insurance-agent-500'
-
-    results = get_hybrid_search_results(query=query, 
-                                    path_to_db=path_to_vector_db, 
-                                    embedding_model=EMBEDDING_MODEL, 
-                                    hf_api_key=HUGGINGFACEHUB_API_TOKEN)
-    
-    for doc in results:
-        print(doc)
-        print()

rag_app/knowledge_base/reranking.py

@@ -1,137 +0,0 @@
-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
-
-
-load_dotenv()
-
-
-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]
-    
-
-
-if __name__ == "__main__":
-
-   
-    HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
-    EMBEDDING_MODEL = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
-
-    project_dir = Path().cwd().parent
-    path_to_vector_db = str(project_dir/'vectorstore/chroma-zurich-mpnet-1500')
-    assert Path(path_to_vector_db).exists(), "Cannot access path_to_vector_db "
-
-    query = "I'm looking for student insurance"
-
-    re_ranked_docs = get_reranked_docs_chroma(query=query,
-                                              path_to_db= path_to_vector_db,
-                                              embedding_model=EMBEDDING_MODEL,
-                                              hf_api_key=HUGGINGFACEHUB_API_TOKEN)
-
-
-    print(f"{re_ranked_docs=}")

rag_app/knowledge_base/utils.py

@@ -11,6 +11,28 @@ from langchain_core.documents import Document
 from config import EMBEDDING_MODEL
 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 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
+
+
 
 def create_embeddings(
         docs: list[Document], 
@@ -86,4 +108,152 @@ def generate_document_summaries(
             {"summary":summary}
         )
     
-    return new_docs
+    return new_docs
+
+
+def build_vector_store(
+        docs: list, 
+        embedding_model: str, 
+        new_db:bool=False, 
+        chunk_size:int=500, 
+        chunk_overlap:int=50,
+        ):
+    """
+
+    """
+
+    embeddings,chunks = create_embeddings(
+        docs, 
+        chunk_size, 
+        chunk_overlap, 
+        embedding_model
+        )
+
+    #load chunks into vector store
+    print(f'Loading chunks into faiss vector store ...')
+    
+    st = time.time()
+    if new_db:
+        db_faiss = FAISS.from_documents(chunks, embeddings)
+        bm25_retriever = BM25Retriever.from_documents(chunks)
+    else:
+        db_faiss = FAISS.add_documents(chunks, embeddings)
+        bm25_retriever = BM25Retriever.add_documents(chunks)
+        
+    db_faiss.save_local(FAISS_INDEX_PATH)
+    et = time.time() - st
+    print(f'Time taken: {et} seconds.')
+
+    print(f'Loading chunks into chroma vector store ...')
+    
+    st = time.time()
+    persist_directory='./vectorstore/chroma-insurance-agent-1500'
+    db_chroma = Chroma.from_documents(chunks, embeddings, persist_directory=persist_directory)
+    et = time.time() - st
+    
+    print(f'Time taken: {et} seconds.')
+    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

@@ -10,7 +10,10 @@ from langchain_community.embeddings.sentence_transformer import (
 )
 import time
 from langchain_core.documents import Document
-from config import EMBEDDING_MODEL
+from config import EMBEDDING_MODEL, HUGGINGFACEHUB_API_TOKEN
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
+import requests
 
 class BaseVectorStore(ABC):
     """
@@ -48,7 +51,10 @@ class BaseVectorStore(ABC):
         documents = loader.load()
         text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
         return text_splitter.split_documents(documents)
-
+    
+    def get_hybrid_search_result(self,query:str):
+        pass
+        
     @abstractmethod
     def create_vectorstore(self, texts):
         """
@@ -89,6 +95,7 @@ class BaseVectorStore(ABC):
         Save the current state of the vector store.
         """
         pass
+    
 
 class ChromaVectorStore(BaseVectorStore):
     """
@@ -133,6 +140,38 @@ class ChromaVectorStore(BaseVectorStore):
         if not self.vectorstore:
             raise ValueError("Vector store not initialized. Nothing to save.")
         self.vectorstore.persist()
+        
+    def get_reranked_docs(
+        self,
+        query:str,
+        num_docs:int=5
+        ):
+        
+        # Get 10 documents based on similarity search
+        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]
+        # 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 {HUGGINGFACEHUB_API_TOKEN}"}
+
+        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]
 
 class FAISSVectorStore(BaseVectorStore):
     """
@@ -170,6 +209,67 @@ class FAISSVectorStore(BaseVectorStore):
         if self.vectorstore is None:
             raise ValueError("Vector store not initialized. Nothing to save.")
         self.vectorstore.save_local(self.persist_directory)
+        
+    def get_hybrid_search_result(
+        self,
+        query:str,
+        num_docs:int=5
+        )-> list[Document]:
+        """ Uses an ensemble retriever of BM25 and FAISS to return k num documents
+
+        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
+            List of documents
+
+        """
+        all_docs = self.vectorstore.similarity_search("", k=self.vectorstore.index.ntotal)
+        bm25_retriever = BM25Retriever.from_documents(all_docs)
+        bm25_retriever.k = num_docs  # How many results you want
+
+        faiss_retriever = self.vectorstore.as_retriever(search_kwargs={'k': num_docs})
+
+        ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, faiss_retriever],
+                                            weights=[0.5,0.5])
+        
+        results = ensemble_retriever.invoke(input=query) 
+        return results
+    
+    def get_reranked_docs(
+        self,
+        query:str,
+        num_docs:int=5
+        ):
+        
+        # Get 10 documents based on similarity search
+        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]
+        # 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 {HUGGINGFACEHUB_API_TOKEN}"}
+
+        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]
 
 # Usage example:
 def main():