Update app.py

app.py

@@ -50,6 +50,16 @@ def document_retrieval_chroma(llm, prompt):
     result = rag_chain({"query": prompt})
     return result["result"]
 
+def document_retrieval_mongodb(llm, prompt):
+    vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
+                       persist_directory = CHROMA_DIR)
+    rag_chain = RetrievalQA.from_chain_type(llm, 
+                                            chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
+                                            retriever = vector_db.as_retriever(search_kwargs = {"k": 3}), 
+                                            return_source_documents = True)
+    result = rag_chain({"query": prompt})
+    return result["result"]
+
 def invoke(openai_api_key, rag, prompt):
     if (openai_api_key == ""):
         raise gr.Error("OpenAI API Key is required.")
@@ -57,11 +67,13 @@ def invoke(openai_api_key, rag, prompt):
         raise gr.Error("Retrieval Augmented Generation is required.")
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
+    
     try:
         llm = ChatOpenAI(model_name = MODEL_NAME, 
                          openai_api_key = openai_api_key, 
                          temperature = 0)
-        if (rag != "None"):
+
+        if (rag == "Chroma"):
             # Document loading
             #docs = []
             # Load PDF
@@ -85,15 +97,10 @@ def invoke(openai_api_key, rag, prompt):
             #                                  embedding = OpenAIEmbeddings(disallowed_special = ()), 
             #                                  persist_directory = CHROMA_DIR)
             # Document retrieval
-            ##vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
-            ##                   persist_directory = CHROMA_DIR)
-            ##rag_chain = RetrievalQA.from_chain_type(llm, 
-            ##                                        chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
-            ##                                        retriever = vector_db.as_retriever(search_kwargs = {"k": 3}), 
-            ##                                        return_source_documents = True)
-            ##result = rag_chain({"query": prompt})
-            ##result = result["result"]
             result = document_retrieval_chroma(llm, prompt)
+        else if (rag == "MongoDB"):
+            # Document retrieval
+            result = document_retrieval_mongodb(llm, prompt)
         else:
             chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)
             result = chain.run({"question": prompt})
@@ -107,8 +114,8 @@ description = """<strong>Overview:</strong> Reasoning application that demonstra
                  <a href='""" + YOUTUBE_URL_1 + """'>YouTube</a>, <a href='""" + PDF_URL + """'>PDF</a>, and <a href='""" + WEB_URL + """'>Web</a> 
                  <strong>data on GPT-4</strong> (published after LLM knowledge cutoff).
                  <ul style="list-style-type:square;">
-                 <li>Set "Retrieval Augmented Generation" to "<strong>False</strong>" and submit prompt "What is GPT-4?" The LLM <strong>without</strong> RAG does not know the answer.</li>
-                 <li>Set "Retrieval Augmented Generation" to "<strong>True</strong>" and submit prompt "What is GPT-4?" The LLM <strong>with</strong> RAG knows the answer.</li>
+                 <li>Set "Retrieval Augmented Generation" to "<strong>None</strong>" and submit prompt "What is GPT-4?" The LLM <strong>without</strong> RAG does not know the answer.</li>
+                 <li>Set "Retrieval Augmented Generation" to "<strong>Chroma</strong>" or "<strong>MongoDB</strong>" and submit prompt "What is GPT-4?" The LLM <strong>with</strong> RAG knows the answer.</li>
                  <li>Experiment with prompts, e.g. "What are GPT-4's media capabilities in 3 emojis and 1 sentence?", "List GPT-4's exam scores and benchmark results.", or "Compare GPT-4 to GPT-3.5 in markdown table format."</li>
                  <li>Experiment some more, for example "What is the GPT-4 API's cost and rate limit? Answer in English, Arabic, Chinese, Hindi, and Russian in JSON format." or "Write a Python program that calls the GPT-4 API."</li>
                  </ul>\n\n