Update trace.py

trace.py

@@ -4,9 +4,7 @@ from wandb.sdk.data_types.trace_tree import Trace
 
 WANDB_API_KEY = os.environ["WANDB_API_KEY"]
 
-RAG_OFF = "Off"
-
-def wandb_trace(rag_option, 
+def wandb_trace(is_rag_off, 
                 prompt, 
                 completion, 
                 result, 
@@ -23,27 +21,27 @@ def wandb_trace(rag_option,
         name = "" if (chain == None) else type(chain).__name__,
         status_code = "success" if (str(err_msg) == "") else "error",
         status_message = str(err_msg),
-        metadata = {"chunk_overlap": "" if (rag_option == RAG_OFF) else config["chunk_overlap"],
-                    "chunk_size": "" if (rag_option == RAG_OFF) else config["chunk_size"],
+        metadata = {"chunk_overlap": "" if (is_rag_off) else config["chunk_overlap"],
+                    "chunk_size": "" if (is_rag_off) else config["chunk_size"],
                    } if (str(err_msg) == "") else {},
-        inputs = {"rag_option": rag_option,
+        inputs = {"is_rag": not is_rag_off,
                   "prompt": prompt,
-                  "chain_prompt": (str(chain.prompt) if (rag_option == RAG_OFF) else 
+                  "chain_prompt": (str(chain.prompt) if (is_rag_off) else 
                                    str(chain.combine_documents_chain.llm_chain.prompt)),
-                  "source_documents": "" if (rag_option == RAG_OFF) else str([doc.metadata["source"] for doc in completion["source_documents"]]),
+                  "source_documents": "" if (is_rag_off) else str([doc.metadata["source"] for doc in completion["source_documents"]]),
                  } if (str(err_msg) == "") else {},
         outputs = {"result": result,
                    "generation_info": str(generation_info),
                    "llm_output": str(llm_output),
                    "completion": str(completion),
                   } if (str(err_msg) == "") else {},
-        model_dict = {"client": (str(chain.llm.client) if (rag_option == RAG_OFF) else
+        model_dict = {"client": (str(chain.llm.client) if (is_rag_off) else
                                  str(chain.combine_documents_chain.llm_chain.llm.client)),
-                      "model_name": (str(chain.llm.model_name) if (rag_option == RAG_OFF) else
+                      "model_name": (str(chain.llm.model_name) if (is_rag_off) else
                                      str(chain.combine_documents_chain.llm_chain.llm.model_name)),
-                      "temperature": (str(chain.llm.temperature) if (rag_option == RAG_OFF) else
+                      "temperature": (str(chain.llm.temperature) if (is_rag_off) else
                                       str(chain.combine_documents_chain.llm_chain.llm.temperature)),
-                      "retriever": ("" if (rag_option == RAG_OFF) else str(chain.retriever)),
+                      "retriever": ("" if (is_rag_off) else str(chain.retriever)),
                      } if (str(err_msg) == "") else {},
         start_time_ms = start_time_ms,
         end_time_ms = end_time_ms