Update app.py

app.py

@@ -107,7 +107,6 @@ def document_retrieval_mongodb(llm, prompt):
 
 def llm_chain(llm, prompt):
     llm_chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT, verbose = False)
-    #completion = llm_chain.run({"question": prompt})
     completion = llm_chain.generate([{"question": prompt}])
     return completion, llm_chain
 
@@ -121,9 +120,7 @@ def rag_chain(llm, prompt, db):
     return completion, rag_chain
 
 def wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, end_time_ms):
-    if (str(err_msg) != ""):
-        result = ""
-    elif (rag_option == RAG_OFF):
+    if (rag_option == RAG_OFF):
         result = completion.generations[0][0].text if (completion.generations[0] != None and 
                                                        completion.generations[0][0] != None) else ""
     else:
@@ -143,8 +140,8 @@ def wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, e
         },
         inputs = {"rag_option": rag_option if (str(err_msg) == "") else "",
                   "prompt": str(prompt if (str(err_msg) == "") else ""),
-                  "prompt_template": str((llm_template if (rag_option == RAG_OFF) else rag_template) if (str(status_msg) == "") else "")},
-        outputs = {"result": result,
+                  "prompt_template": str((llm_template if (rag_option == RAG_OFF) else rag_template) if (str(err_msg) == "") else "")},
+        outputs = {"result": result if (str(err_msg) == "") else "",
                    "completion": str(completion)},
         start_time_ms = start_time_ms,
         end_time_ms = end_time_ms