Update app.py

app.py

@@ -115,13 +115,13 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion
 
-def wandb_trace(rag_option, prompt, prompt_template, completion):
+def wandb_trace(rag_option, prompt, prompt_template, completion, chain_name):
     wandb.init(project = "openai-llm-rag")
     trace = Trace(
-        name="test",
-        kind="chain",
-        #status_code=status,
-        #status_message=status_message,
+        name = chain_name,
+        kind = "chain",
+        #status_code = status,
+        #status_message = status_message,
         metadata={
             "chunk_overlap": config["chunk_overlap"],
             "chunk_size": config["chunk_size"],
@@ -129,10 +129,10 @@ def wandb_trace(rag_option, prompt, prompt_template, completion):
             "model": config["model"],
             "temperature": config["temperature"],
         },
-        #start_time_ms=start_time_ms,
-        #end_time_ms=end_time_ms,
-        inputs={"rag_option": rag_option, "prompt": prompt, "prompt_template": prompt_template},
-        outputs={"completion": str(completion)},
+        #start_time_ms = start_time_ms,
+        #end_time_ms = end_time_ms,
+        inputs = {"rag_option": rag_option, "prompt": prompt, "prompt_template": prompt_template},
+        outputs = {"completion": str(completion)},
     )
     trace.log("test")
     wandb.finish()
@@ -155,6 +155,7 @@ def invoke(openai_api_key, rag_option, prompt):
             completion = rag_chain(llm, prompt, db)
             result = completion["result"]
             prompt_template = rag_template
+            chain_name = type(RetrievalQA).__name__
         elif (rag_option == "MongoDB"):
             #splits = document_loading_splitting()
             #document_storage_mongodb(splits)
@@ -162,15 +163,17 @@ def invoke(openai_api_key, rag_option, prompt):
             completion = rag_chain(llm, prompt, db)
             result = completion["result"]
             prompt_template = rag_template
+            chain_name = type(RetrievalQA).__name__
         else:
             result = llm_chain(llm, prompt)
             completion = result
             prompt_template = llm_template
+            chain_name = type(LLMChain).__name__
     except Exception as e:
         completion = e
         raise gr.Error(e)
     finally:
-        wandb_trace(rag_option, prompt, prompt_template, completion)
+        wandb_trace(rag_option, prompt, prompt_template, completion, chain_name)
     return result
 
 description = """<strong>Overview:</strong> Context-aware multimodal reasoning application using a <strong>large language model (LLM)</strong> with