Update app.py

app.py

@@ -119,12 +119,7 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
-def wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, end_time_ms):
-    if (rag_option == RAG_OFF):
-        result = completion.generations[0][0].text if (completion.generations[0] != None and 
-                                                       completion.generations[0][0] != None) else ""
-    else:
-        result = completion["result"]
+def wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_time_ms, end_time_ms):
     wandb.init(project = "openai-llm-rag")
     trace = Trace(
         kind = "chain",
@@ -187,7 +182,7 @@ def invoke(openai_api_key, rag_option, prompt):
     finally:
         end_time_ms = round(time.time() * 1000)
         print(chain)
-        wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()