Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import logging, os, sys
+import logging, os, sys, time
 
 from agent_langchain import agent_langchain
 from agent_llamaindex import agent_llamaindex
@@ -33,6 +33,8 @@ def invoke(openai_api_key, prompt, agent_option):
     output = ""
     
     try:
+        start_time_ms = round(time.time() * 1000)
+        
         if (agent_option == AGENT_LANGCHAIN):
             completion = agent_langchain(
                 config,
@@ -59,6 +61,20 @@ def invoke(openai_api_key, prompt, agent_option):
         err_msg = e
 
         raise gr.Error(e)
+    finally:
+        end_time_ms = round(time.time() * 1000)
+        
+        trace_wandb(
+            config,
+            rag_option,
+            prompt, 
+            completion, 
+            result, 
+            callback, 
+            err_msg, 
+            start_time_ms, 
+            end_time_ms
+        )
 
     return output