Update app.py

app.py

@@ -40,7 +40,7 @@ YOUTUBE_URL_3 = "https://www.youtube.com/watch?v=vw-KWfKwvTQ"
 
 MODEL_NAME  = "gpt-4"
 
-def document_retrieval_chroma():
+def document_retrieval_chroma(llm):
     vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
                        persist_directory = CHROMA_DIR)
     rag_chain = RetrievalQA.from_chain_type(llm, 
@@ -50,7 +50,7 @@ def document_retrieval_chroma():
     result = rag_chain({"query": prompt})
     return result["result"]
 
-def invoke(openai_api_key, use_rag, rag_db, prompt):
+def invoke(openai_api_key, rag, prompt):
     if (openai_api_key == ""):
         raise gr.Error("OpenAI API Key is required.")
     if (use_rag is None):
@@ -61,7 +61,7 @@ def invoke(openai_api_key, use_rag, rag_db, prompt):
         llm = ChatOpenAI(model_name = MODEL_NAME, 
                          openai_api_key = openai_api_key, 
                          temperature = 0)
-        if (use_rag):
+        if (rag != "None"):
             # Document loading
             #docs = []
             # Load PDF
@@ -93,7 +93,7 @@ def invoke(openai_api_key, use_rag, rag_db, prompt):
             ##                                        return_source_documents = True)
             ##result = rag_chain({"query": prompt})
             ##result = result["result"]
-            result = document_retrieval_chroma()
+            result = document_retrieval_chroma(llm)
         else:
             chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)
             result = chain.run({"question": prompt})
@@ -120,8 +120,7 @@ description = """<strong>Overview:</strong> Reasoning application that demonstra
 gr.close_all()
 demo = gr.Interface(fn=invoke, 
                     inputs = [gr.Textbox(label = "OpenAI API Key", value = "sk-", lines = 1), 
-                              gr.Radio([True, False], label="Retrieval Augmented Generation (RAG)", value = False),
-                              gr.Radio(["Chroma", "MongoDB"], label="RAG Database", value = "Chroma"),
+                              gr.Radio(["None", "Chroma", "MongoDB"], label="Retrieval Augmented Generation", value = "None"),
                               gr.Textbox(label = "Prompt", value = "What is GPT-4?", lines = 1)],
                     outputs = [gr.Textbox(label = "Completion", lines = 1)],
                     title = "Generative AI - LLM & RAG",