Update app.py

app.py

@@ -115,6 +115,29 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion
 
+def wandb_trace(prompt, completion, rag_option)
+    wandb.init(project = "openai-llm-rag")
+    trace = Trace(
+        name="test",
+        kind="chain",
+        #status_code=status,
+        #status_message=status_message,
+        metadata={
+            "rag_option": rag_option 
+            "chunk_overlap": config["chunk_overlap"],
+            "chunk_size": config["chunk_size"],
+            "k": config["k"],
+            "model": config["model"],
+            "temperature": config["temperature"],
+        },
+        #start_time_ms=start_time_ms,
+        #end_time_ms=end_time_ms,
+        inputs={"prompt": prompt},
+        outputs={"completion": completion},
+    )
+    trace.log("test")
+    wandb.finish()
+
 def invoke(openai_api_key, rag_option, prompt):
     if (openai_api_key == ""):
         raise gr.Error("OpenAI API Key is required.")
@@ -122,7 +145,6 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error("Retrieval Augmented Generation is required.")
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
-    wandb.init(project = "openai-llm-rag")
     completion = ""
     try:
         llm = ChatOpenAI(model_name = config["model"], 
@@ -146,23 +168,7 @@ 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()
+        wandb_trace(prompt, completion, rag_option)
     return completion
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application using a <strong>large language model (LLM)</strong> with