Update rag.py

rag.py

@@ -95,10 +95,13 @@ def llm_chain(config, openai_api_key, prompt):
     
     return completion, llm_chain
 
-def rag_chain(config, openai_api_key, prompt):
+def rag_chain(config, openai_api_key, rag_option, prompt):
     llm = get_llm(config, openai_api_key)
     
-    db = document_retrieval_chroma(llm, prompt)
+    if (rag_option == RAG_CHROMA):
+        db = document_retrieval_chroma(llm, prompt)
+    else:
+        db = document_retrieval_mongodb(llm, prompt)
     
     rag_chain = RetrievalQA.from_chain_type(llm, 
                                             chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT},