Update app.py

app.py

@@ -16,6 +16,7 @@ from langchain.vectorstores import MongoDBAtlasVectorSearch
 
 from pymongo import MongoClient
 
+from wandb.integration.langchain import WandbTracer
 from wandb.sdk.data_types.trace_tree import Trace
 
 from dotenv import load_dotenv, find_dotenv
@@ -151,6 +152,7 @@ def invoke(openai_api_key, rag_option, prompt):
     chain_name = ""
     status_msg = ""
     try:
+        WandbTracer.init()        
         start_time_ms = round(time.time() * 1000)
         llm = ChatOpenAI(model_name = config["model"], 
                          openai_api_key = openai_api_key, 
@@ -180,7 +182,8 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error(e)
     finally:
         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)
+        #wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_name, status_msg, start_time_ms, end_time_ms)
+        WandbTracer.finish()
     return result
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application using a <strong>large language model (LLM)</strong> with