Update app.py

app.py

@@ -115,7 +115,7 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion
 
-def wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_name, status_msg):
+def wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_name, status_msg, start_time_ms, end_time_ms):
     wandb.init(project = "openai-llm-rag")
     trace = Trace(
         name = chain_name,
@@ -129,8 +129,8 @@ def wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_n
             "model": config["model"],
             "temperature": config["temperature"],
         },
-        #start_time_ms = 123,
-        #end_time_ms = 456,
+        start_time_ms = start_time_ms,
+        end_time_ms = end_time_ms,
         inputs = {"rag_option": rag_option, "prompt": prompt, "prompt_template": prompt_template},
         outputs = {"result": str(result), "completion": str(completion)},
     )
@@ -145,7 +145,11 @@ def invoke(openai_api_key, rag_option, prompt):
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
     completion = ""
+    result = ""
+    prompt_template = ""
+    chain_name = ""
     status_msg = ""
+    start_time_ms = round(time.time() * 1000)
     try:
         llm = ChatOpenAI(model_name = config["model"], 
                          openai_api_key = openai_api_key, 
@@ -175,7 +179,8 @@ def invoke(openai_api_key, rag_option, prompt):
         status_msg = e
         raise gr.Error(e)
     finally:
-        wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_name, status_msg)
+        end_time_ms = round(time.time() * 1000)
+        wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_name, status_msg, start_time_ms, end_time_ms)
     return result
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application using a <strong>large language model (LLM)</strong> with