Update trace.py

trace.py

@@ -12,6 +12,7 @@ def wandb_trace(config,
                 generation_info, 
                 llm_output, 
                 chain, 
+                cb, 
                 err_msg, 
                 start_time_ms, 
                 end_time_ms):
@@ -32,6 +33,7 @@ def wandb_trace(config,
                   "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,
+                   "cb": cb,
                    "generation_info": str(generation_info),
                    "llm_output": str(llm_output),
                    "completion": str(completion),