Update app.py

app.py

@@ -20,8 +20,6 @@ from dotenv import load_dotenv, find_dotenv
 _ = load_dotenv(find_dotenv())
 
 WANDB_API_KEY = os.environ["WANDB_API_KEY"]
-os.environ["LANGCHAIN_WANDB_TRACING"] = "true"
-os.environ["WANDB_PROJECT"] = "openai-llm-rag"
 
 MONGODB_URI = os.environ["MONGODB_ATLAS_CLUSTER_URI"]
 client = MongoClient(MONGODB_URI)
@@ -122,6 +120,8 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error("Retrieval Augmented Generation is required.")
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
+    os.environ["LANGCHAIN_WANDB_TRACING"] = "true"
+    os.environ["WANDB_PROJECT"] = "openai-llm-rag"
     completion = ""
     try:
         llm = ChatOpenAI(model_name = config["model"], 
@@ -144,8 +144,8 @@ def invoke(openai_api_key, rag_option, prompt):
     except Exception as e:
         completion = e
         raise gr.Error(e)
-    #finally:
-    #    wandb.finish()
+    finally:
+        wandb.finish()
     return completion
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application using a <strong>large language model (LLM)</strong> with