File size: 1,507 Bytes
e874531 fcd3a75 a857f12 9e0f123 6ac712e 301617d 5059289 86ffba3 a857f12 8792ffa 8a75a06 8792ffa 8a75a06 8792ffa a857f12 da04f66 5059289 5b7aea8 63dc30c 8792ffa 63dc30c 5b7aea8 a857f12 86ffba3 a857f12 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import os, wandb
from wandb.sdk.data_types.trace_tree import Trace
WANDB_API_KEY = os.environ["WANDB_API_KEY"]
RAG_LANGCHAIN = "LangChain"
RAG_LLAMAINDEX = "LlamaIndex"
def trace_wandb(config,
rag_option,
prompt,
completion,
result,
callback,
err_msg,
start_time_ms,
end_time_ms):
wandb.init(project = "openai-llm-rag")
if (rag_option == RAG_LANGCHAIN):
prompt_template = os.environ["LANGCHAIN_TEMPLATE"]
elif (rag_option == RAG_LLAMAINDEX):
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,
"rag_option": rag_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() |