Create trace.py

trace.py

@@ -0,0 +1,48 @@
+import os, wandb
+
+from wandb.sdk.data_types.trace_tree import Trace
+
+WANDB_API_KEY = os.environ["WANDB_API_KEY"]
+
+AGENT_LANGCHAIN  = "LangChain"
+AGENT_LLAMAINDEX = "LlamaIndex"
+
+def trace_wandb(config,
+                agent_option,
+                prompt,
+                completion,
+                result,
+                callback,
+                err_msg,
+                start_time_ms,
+                end_time_ms):
+    wandb.init(project = "openai-llm-agent")
+
+    if (agent_option == AGENT_LANGCHAIN):
+        prompt_template = os.environ["LANGCHAIN_TEMPLATE"]
+    elif (agent_option == AGENT_LLAMAINDEX):
+        prompt_template = os.environ["LLAMAINDEX_TEMPLATE"]
+    else:
+        prompt_template = os.environ["TEMPLATE"]
+                    
+    trace = Trace(
+        kind = "LLM",
+        name = "Real-Time Reasoning Application",
+        status_code = "success" if (str(err_msg) == "") else "error",
+        status_message = str(err_msg),
+        inputs = {"prompt": prompt,
+                  "prompt_template": prompt_template,
+                  "agent_option": agent_option,
+                  "config": str(config)
+                 } if (str(err_msg) == "") else {},
+        outputs = {"result": str(result),
+                   "callback": str(callback),
+                   "completion": str(completion)
+                  } if (str(err_msg) == "") else {},
+        start_time_ms = start_time_ms,
+        end_time_ms = end_time_ms
+    )
+    
+    trace.log("evaluation")
+                    
+    wandb.finish()