Update app.py

app.py

@@ -16,6 +16,8 @@ from langchain.vectorstores import MongoDBAtlasVectorSearch
 
 from pymongo import MongoClient
 
+from wandb.sdk.data_types.trace_tree import Trace
+
 from dotenv import load_dotenv, find_dotenv
 _ = load_dotenv(find_dotenv())
 
@@ -120,8 +122,7 @@ 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["WANDB_PROJECT"] = "openai-llm-rag"
-    os.environ["LANGCHAIN_WANDB_TRACING"] = "true"
+    wandb.init(project = "openai-llm-rag")
     completion = ""
     try:
         llm = ChatOpenAI(model_name = config["model"], 
@@ -145,9 +146,23 @@ def invoke(openai_api_key, rag_option, prompt):
         completion = e
         raise gr.Error(e)
     finally:
+        trace = Trace(
+            name="test",
+            kind="chain",
+            #status_code=status,
+            #status_message=status_message,
+            metadata={
+                "temperature": config["temperature"],
+                #"token_usage": token_usage,
+                #"model_name": model_name,
+            },
+            #start_time_ms=start_time_ms,
+            #end_time_ms=end_time_ms,
+            #inputs={"system_prompt": system_message, "query": query},
+            #outputs={"response": response_text},
+        )
+        trace.log("test")
         wandb.finish()
-        del os.environ["LANGCHAIN_WANDB_TRACING"]
-        del os.environ["WANDB_PROJECT"]
     return completion
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application using a <strong>large language model (LLM)</strong> with