Update rag.py

rag.py

@@ -14,6 +14,9 @@ from langchain.vectorstores import MongoDBAtlasVectorSearch
 
 from pymongo import MongoClient
 
+RAG_CHROMA  = "Chroma"
+RAG_MONGODB = "MongoDB"
+
 PDF_URL       = "https://arxiv.org/pdf/2303.08774.pdf"
 WEB_URL       = "https://openai.com/research/gpt-4"
 YOUTUBE_URL_1 = "https://www.youtube.com/watch?v=--khbXchTeE"
@@ -71,11 +74,11 @@ def document_storage_mongodb(documents):
                                             collection = collection,
                                             index_name = MONGODB_INDEX_NAME)
 
-def document_retrieval_chroma(llm, prompt):
+def document_retrieval_chroma():
     return Chroma(embedding_function = OpenAIEmbeddings(),
                   persist_directory = CHROMA_DIR)
 
-def document_retrieval_mongodb(llm, prompt):
+def document_retrieval_mongodb():
     return MongoDBAtlasVectorSearch.from_connection_string(MONGODB_ATLAS_CLUSTER_URI,
                                                            MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
                                                            OpenAIEmbeddings(disallowed_special = ()),
@@ -99,9 +102,9 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
     llm = get_llm(config, openai_api_key)
     
     if (rag_option == RAG_CHROMA):
-        db = document_retrieval_chroma(llm, prompt)
-    else:
-        db = document_retrieval_mongodb(llm, prompt)
+        db = document_retrieval_chroma()
+    elif (rag_option == RAG_MONGODB):
+        db = document_retrieval_mongodb()
     
     rag_chain = RetrievalQA.from_chain_type(llm, 
                                             chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT},