Update app.py

app.py

@@ -30,7 +30,6 @@ MONGODB_INDEX_NAME = "default"
 
 config = {
     "model": "gpt-4",
-    "rag_option": "Off",
     "temperature": 0,
 }
 
@@ -106,7 +105,6 @@ def document_retrieval_mongodb(llm, prompt):
 def llm_chain(llm, prompt):
     llm_chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)
     completion = llm_chain.run({"question": prompt})
-    wandb.log({"prompt": prompt, "completion": completion})
     return completion
 
 def rag_chain(llm, prompt, db):
@@ -115,9 +113,7 @@ def rag_chain(llm, prompt, db):
                                             retriever = db.as_retriever(search_kwargs = {"k": 3}), 
                                             return_source_documents = True)
     completion = rag_chain({"query": prompt})
-    completion_result = completion["result"]
-    wandb.log({"prompt": prompt, "completion": completion_result})
-    return completion_result
+    return completion["result"]
 
 def invoke(openai_api_key, rag_option, prompt):
     if (openai_api_key == ""):
@@ -126,7 +122,6 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error("Retrieval Augmented Generation is required.")
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
-    #wandb.config[rag_option] = rag_option
     try:
         llm = ChatOpenAI(model_name = config.model, 
                          openai_api_key = openai_api_key, 
@@ -135,17 +130,19 @@ def invoke(openai_api_key, rag_option, prompt):
             #splits = document_loading_splitting()
             #document_storage_chroma(splits)
             db = document_retrieval_chroma(llm, prompt)
-            result = rag_chain(llm, prompt, db)
+            completion = rag_chain(llm, prompt, db)
         elif (rag_option == "MongoDB"):
             #splits = document_loading_splitting()
             #document_storage_mongodb(splits)
             db = document_retrieval_mongodb(llm, prompt)
-            result = rag_chain(llm, prompt, db)
+            completion = rag_chain(llm, prompt, db)
         else:
             result = llm_chain(llm, prompt)
     except Exception as e:
         raise gr.Error(e)
-    return result
+    wandb.config[rag_option] = rag_option
+    wandb.log({"prompt": prompt, "completion": completion})
+    return completion
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application that demonstrates a <strong>large language model (LLM)</strong> with 
                  <strong>retrieval augmented generation (RAG)</strong>.