Update app.py

app.py

@@ -132,7 +132,7 @@ def wandb_trace(prompt, completion, rag_option):
         },
         #start_time_ms=start_time_ms,
         #end_time_ms=end_time_ms,
-        inputs={"prompt": prompt},
+        inputs={"prompt": prompt, "template": LLM_CHAIN_PROMPT},
         outputs={"completion": completion},
     )
     trace.log("test")
@@ -155,21 +155,21 @@ def invoke(openai_api_key, rag_option, prompt):
             #document_storage_chroma(splits)
             db = document_retrieval_chroma(llm, prompt)
             completion = rag_chain(llm, prompt, db)
-            completion = completion["result"]
+            result = completion["result"]
         elif (rag_option == "MongoDB"):
             #splits = document_loading_splitting()
             #document_storage_mongodb(splits)
             db = document_retrieval_mongodb(llm, prompt)
             completion = rag_chain(llm, prompt, db)
-            completion = completion["result"]
+            result = completion["result"]
         else:
-            completion = llm_chain(llm, prompt)
+            result = llm_chain(llm, prompt)
     except Exception as e:
         completion = e
         raise gr.Error(e)
     finally:
         wandb_trace(prompt, completion, rag_option)
-    return completion
+    return result
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application using a <strong>large language model (LLM)</strong> with 
                  <strong>retrieval augmented generation (RAG)</strong>.