| 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() |