condnl checks replaved with dict lookups

code/modules/retriever/retriever.py

@@ -6,19 +6,19 @@ from modules.retriever.colbert_retriever import ColbertRetriever
 class Retriever:
     def __init__(self, config):
         self.config = config
+        self.retriever_classes = {
+            "FAISS": FaissRetriever,
+            "Chroma": ChromaRetriever,
+            "RAGatouille": ColbertRetriever,
+        }
         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"])
-            )
+        db_option = self.config["vectorstore"]["db_option"]
+        retriever_class = self.retriever_classes.get(db_option)
+        if not retriever_class:
+            raise ValueError(f"Invalid db_option: {db_option}")
+        self.retriever = retriever_class()
 
     def _return_retriever(self, db):
         return self.retriever.return_retriever(db, self.config)

code/modules/vectorstore/store_manager.py

@@ -86,6 +86,8 @@ class VectorStoreManager:
     ):
         if self.config["vectorstore"]["db_option"] in ["FAISS", "Chroma"]:
             self.embedding_model = self.create_embedding_model()
+        else:
+            self.embedding_model = None
 
         self.logger.info("Initializing vector_db")
         self.logger.info(
@@ -132,6 +134,8 @@ class VectorStoreManager:
         start_time = time.time()  # Start time for loading database
         if self.config["vectorstore"]["db_option"] in ["FAISS", "Chroma"]:
             self.embedding_model = self.create_embedding_model()
+        else:
+            self.embedding_model = None
         self.loaded_vector_db = self.vector_db._load_database(self.embedding_model)
         end_time = time.time()  # End time for loading database
         self.logger.info(

code/modules/vectorstore/vectorstore.py

@@ -7,6 +7,11 @@ class VectorStore:
     def __init__(self, config):
         self.config = config
         self.vectorstore = None
+        self.vectorstore_classes = {
+            "FAISS": FaissVectorStore,
+            "Chroma": ChromaVectorStore,
+            "RAGatouille": ColbertVectorStore,
+        }
 
     def _create_database(
         self,
@@ -16,32 +21,32 @@ class VectorStore:
         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)
+        db_option = self.config["vectorstore"]["db_option"]
+        vectorstore_class = self.vectorstore_classes.get(db_option)
+        if not vectorstore_class:
+            raise ValueError(f"Invalid db_option: {db_option}")
+
+        self.vectorstore = vectorstore_class(self.config)
+
+        if db_option == "RAGatouille":
             self.vectorstore.create_database(
                 documents, document_names, document_metadata
             )
         else:
-            raise ValueError(
-                "Invalid db_option: {}".format(self.config["vectorstore"]["db_option"])
-            )
+            self.vectorstore.create_database(document_chunks, embedding_model)
 
     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)
+        db_option = self.config["vectorstore"]["db_option"]
+        vectorstore_class = self.vectorstore_classes.get(db_option)
+        if not vectorstore_class:
+            raise ValueError(f"Invalid db_option: {db_option}")
+
+        self.vectorstore = vectorstore_class(self.config)
+
+        if db_option == "RAGatouille":
             return self.vectorstore.load_database()
+        else:
+            return self.vectorstore.load_database(embedding_model)
 
     def _as_retriever(self):
         return self.vectorstore.as_retriever()