more modularization for vectorestore and retriever

code/modules/chat/llm_tutor.py

@@ -10,7 +10,7 @@ from modules.chat.helpers import get_prompt
 from modules.chat.chat_model_loader import ChatModelLoader
 from modules.vectorstore.store_manager import VectorStoreManager
 
-from modules.retriever import FaissRetriever, ChromaRetriever
+from modules.retriever import Retriever
 
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 from langchain_core.callbacks.manager import AsyncCallbackManagerForChainRun
@@ -159,16 +159,7 @@ class LLMTutor:
     # Retrieval QA Chain
     def retrieval_qa_chain(self, llm, prompt, db):
 
-        if self.config["vectorstore"]["db_option"] == "FAISS":
-            retriever = FaissRetriever().return_retriever(db, self.config)
-
-        elif self.config["vectorstore"]["db_option"] == "Chroma":
-            retriever = ChromaRetriever().return_retriever(db, self.config)
-
-        elif self.config["vectorstore"]["db_option"] == "RAGatouille":
-            retriever = db.as_langchain_retriever(
-                k=self.config["vectorstore"]["search_top_k"]
-            )
+        retriever = Retriever(self.config)._return_retriever(db)
 
         if self.config["llm_params"]["use_history"]:
             memory = ConversationBufferWindowMemory(

code/modules/config/config.yml

@@ -7,7 +7,7 @@ vectorstore:
   data_path: '../storage/data' # str
   url_file_path: '../storage/data/urls.txt' # str
   expand_urls: False # bool
-  db_option : 'RAGatouille' # str [FAISS, Chroma, RAGatouille]
+  db_option : 'Chroma' # str [FAISS, Chroma, RAGatouille]
   db_path : '../vectorstores' # str
   model : 'sentence-transformers/all-MiniLM-L6-v2' # str [sentence-transformers/all-MiniLM-L6-v2, text-embedding-ada-002']
   search_top_k : 3 # int

code/modules/retriever/__init__.py

@@ -1,2 +1,4 @@
 from .faiss_retriever import FaissRetriever
 from .chroma_retriever import ChromaRetriever
+from .colbert_retriever import ColbertRetriever
+from .retriever import Retriever

code/modules/retriever/colbert_retriever.py

@@ -0,0 +1,10 @@
+from .base import BaseRetriever
+
+
+class ColbertRetriever(BaseRetriever):
+    def __init__(self):
+        pass
+
+    def return_retriever(self, db, config):
+        retriever = db.as_retriever()
+        return retriever

code/modules/retriever/retriever.py

@@ -0,0 +1,24 @@
+from modules.retriever.faiss_retriever import FaissRetriever
+from modules.retriever.chroma_retriever import ChromaRetriever
+from modules.retriever.colbert_retriever import ColbertRetriever
+
+
+class Retriever:
+    def __init__(self, config):
+        self.config = config
+        self._create_retriever()
+
+    def _create_retriever(self):
+        if self.config["vectorstore"]["db_option"] == "FAISS":
+            self.retriever = FaissRetriever()
+        elif self.config["vectorstore"]["db_option"] == "Chroma":
+            self.retriever = ChromaRetriever()
+        elif self.config["vectorstore"]["db_option"] == "RAGatouille":
+            self.retriever = ColbertRetriever()
+        else:
+            raise ValueError(
+                "Invalid db_option: {}".format(self.config["vectorstore"]["db_option"])
+            )
+
+    def _return_retriever(self, db):
+        return self.retriever.return_retriever(db, self.config)

code/modules/vectorstore/store_manager.py

@@ -1,6 +1,4 @@
-from modules.vectorstore.faiss import FaissVectorStore
-from modules.vectorstore.chroma import ChromaVectorStore
-from modules.vectorstore.colbert import ColbertVectorStore
+from modules.vectorstore.vectorstore import VectorStore
 from modules.vectorstore.helpers import *
 from modules.dataloader.webpage_crawler import WebpageCrawler
 from modules.dataloader.data_loader import DataLoader
@@ -15,7 +13,6 @@ import asyncio
 class VectorStoreManager:
     def __init__(self, config, logger=None):
         self.config = config
-        self.db_option = config["vectorstore"]["db_option"]
         self.document_names = None
 
         # Set up logging to both console and a file
@@ -47,9 +44,12 @@ class VectorStoreManager:
 
         self.webpage_crawler = WebpageCrawler()
 
+        self.vector_db = VectorStore(self.config)
+
         self.logger.info("VectorDB instance instantiated")
 
     def load_files(self):
+
         files = os.listdir(self.config["vectorstore"]["data_path"])
         files = [
             os.path.join(self.config["vectorstore"]["data_path"], file)
@@ -71,6 +71,7 @@ class VectorStoreManager:
         return files, urls
 
     def create_embedding_model(self):
+
         self.logger.info("Creating embedding function")
         embedding_model_loader = EmbeddingModelLoader(self.config)
         embedding_model = embedding_model_loader.load_embedding_model()
@@ -83,22 +84,23 @@ class VectorStoreManager:
         documents: list,
         document_metadata: list,
     ):
-        if self.db_option in ["FAISS", "Chroma"]:
+        if self.config["vectorstore"]["db_option"] in ["FAISS", "Chroma"]:
             self.embedding_model = self.create_embedding_model()
 
         self.logger.info("Initializing vector_db")
-        self.logger.info("\tUsing {} as db_option".format(self.db_option))
-        if self.db_option == "FAISS":
-            self.vector_db = FaissVectorStore(self.config)
-            self.vector_db.create_database(document_chunks, self.embedding_model)
-        elif self.db_option == "Chroma":
-            self.vector_db = ChromaVectorStore(self.config)
-            self.vector_db.create_database(document_chunks, self.embedding_model)
-        elif self.db_option == "RAGatouille":
-            self.vector_db = ColbertVectorStore(self.config)
-            self.vector_db.create_database(documents, document_names, document_metadata)
+        self.logger.info(
+            "\tUsing {} as db_option".format(self.config["vectorstore"]["db_option"])
+        )
+        self.vector_db._create_database(
+            document_chunks,
+            document_names,
+            documents,
+            document_metadata,
+            self.embedding_model,
+        )
 
     def create_database(self):
+
         start_time = time.time()  # Start time for creating database
         data_loader = DataLoader(self.config, self.logger)
         self.logger.info("Loading data")
@@ -126,18 +128,11 @@ class VectorStoreManager:
         )
 
     def load_database(self):
+
         start_time = time.time()  # Start time for loading database
-        if self.db_option in ["FAISS", "Chroma"]:
+        if self.config["vectorstore"]["db_option"] in ["FAISS", "Chroma"]:
             self.embedding_model = self.create_embedding_model()
-        if self.db_option == "FAISS":
-            self.vector_db = FaissVectorStore(self.config)
-            self.loaded_vector_db = self.vector_db.load_database(self.embedding_model)
-        elif self.db_option == "Chroma":
-            self.vector_db = ChromaVectorStore(self.config)
-            self.loaded_vector_db = self.vector_db.load_database(self.embedding_model)
-        elif self.db_option == "RAGatouille":
-            self.vector_db = ColbertVectorStore(self.config)
-            self.loaded_vector_db = self.vector_db.load_database()
+        self.loaded_vector_db = self.vector_db._load_database(self.embedding_model)
         end_time = time.time()  # End time for loading database
         self.logger.info(
             f"Time taken to load database: {end_time - start_time} seconds"

code/modules/vectorstore/vectorstore.py

@@ -0,0 +1,50 @@
+from modules.vectorstore.faiss import FaissVectorStore
+from modules.vectorstore.chroma import ChromaVectorStore
+from modules.vectorstore.colbert import ColbertVectorStore
+
+
+class VectorStore:
+    def __init__(self, config):
+        self.config = config
+        self.vectorstore = None
+
+    def _create_database(
+        self,
+        document_chunks,
+        document_names,
+        documents,
+        document_metadata,
+        embedding_model,
+    ):
+        if self.config["vectorstore"]["db_option"] == "FAISS":
+            self.vectorstore = FaissVectorStore(self.config)
+            self.vectorstore.create_database(document_chunks, embedding_model)
+        elif self.config["vectorstore"]["db_option"] == "Chroma":
+            self.vectorstore = ChromaVectorStore(self.config)
+            self.vectorstore.create_database(document_chunks, embedding_model)
+        elif self.config["vectorstore"]["db_option"] == "RAGatouille":
+            self.vectorstore = ColbertVectorStore(self.config)
+            self.vectorstore.create_database(
+                documents, document_names, document_metadata
+            )
+        else:
+            raise ValueError(
+                "Invalid db_option: {}".format(self.config["vectorstore"]["db_option"])
+            )
+
+    def _load_database(self, embedding_model):
+        if self.config["vectorstore"]["db_option"] == "FAISS":
+            self.vectorstore = FaissVectorStore(self.config)
+            return self.vectorstore.load_database(embedding_model)
+        elif self.config["vectorstore"]["db_option"] == "Chroma":
+            self.vectorstore = ChromaVectorStore(self.config)
+            return self.vectorstore.load_database(embedding_model)
+        elif self.config["vectorstore"]["db_option"] == "RAGatouille":
+            self.vectorstore = ColbertVectorStore(self.config)
+            return self.vectorstore.load_database()
+
+    def _as_retriever(self):
+        return self.vectorstore.as_retriever()
+
+    def _get_vectorstore(self):
+        return self.vectorstore