Update trace.py

trace.py

@@ -14,14 +14,21 @@ def trace_wandb(config,
                 start_time_ms,
                 end_time_ms):
     wandb.init(project = "openai-llm-rag")
-    
+
+    if (rag_option = ""):
+        prompt_template = os.environ["LANGCHAIN_TEMPLATE"]
+    elif (rag_option = ""):
+        prompt_template = os.environ["LLAMAINDEX_TEMPLATE"]
+    else:
+        prompt_template = os.environ["TEMPLATE"]
+                    
     trace = Trace(
         kind = "LLM",
         name = "Context-Aware Reasoning Application",
         status_code = "success" if (str(err_msg) == "") else "error",
         status_message = str(err_msg),
         inputs = {"prompt": prompt,
-                  #"prompt_template": prompt_template,
+                  "prompt_template": prompt_template,
                   "rag_option": rag_option,
                   "config": str(config)
                  } if (str(err_msg) == "") else {},