Update app.py

app.py

@@ -14,7 +14,8 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.vectorstores import MongoDBAtlasVectorSearch
 from pymongo import MongoClient
-from wandb.sdk.data_types.trace_tree import Trace
+from trace import wandb_trace
+#from wandb.sdk.data_types.trace_tree import Trace
 
 _ = load_dotenv(find_dotenv())
 
@@ -35,7 +36,7 @@ MONGODB_INDEX_NAME        = "default"
 LLM_CHAIN_PROMPT = PromptTemplate(input_variables = ["question"], template = os.environ["LLM_TEMPLATE"])
 RAG_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"], template = os.environ["RAG_TEMPLATE"])
 
-WANDB_API_KEY = os.environ["WANDB_API_KEY"]
+#WANDB_API_KEY = os.environ["WANDB_API_KEY"]
 
 RAG_OFF     = "Off"
 RAG_CHROMA  = "Chroma"
@@ -115,42 +116,42 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
-def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms):
-    wandb.init(project = "openai-llm-rag")
-    
-    trace = Trace(
-        kind = "chain",
-        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"],
-                   } if (str(err_msg) == "") else {},
-        inputs = {"rag_option": rag_option,
-                  "prompt": prompt,
-                  "chain_prompt": (str(chain.prompt) if (rag_option == 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"]]),
-                 } 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
-                                 str(chain.combine_documents_chain.llm_chain.llm.client)),
-                      "model_name": (str(chain.llm.model_name) if (rag_option == RAG_OFF) else
-                                     str(chain.combine_documents_chain.llm_chain.llm.model_name)),
-                      "temperature": (str(chain.llm.temperature) if (rag_option == RAG_OFF) else
-                                      str(chain.combine_documents_chain.llm_chain.llm.temperature)),
-                      "retriever": ("" if (rag_option == RAG_OFF) else str(chain.retriever)),
-                     } if (str(err_msg) == "") else {},
-        start_time_ms = start_time_ms,
-        end_time_ms = end_time_ms
-    )
-    
-    trace.log("evaluation")
-    wandb.finish()
+#def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms):
+#    wandb.init(project = "openai-llm-rag")
+#    
+#    trace = Trace(
+#        kind = "chain",
+#        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"],
+#                   } if (str(err_msg) == "") else {},
+#        inputs = {"rag_option": rag_option,
+#                  "prompt": prompt,
+#                  "chain_prompt": (str(chain.prompt) if (rag_option == 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"]]),
+#                 } 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
+#                                 str(chain.combine_documents_chain.llm_chain.llm.client)),
+#                      "model_name": (str(chain.llm.model_name) if (rag_option == RAG_OFF) else
+#                                     str(chain.combine_documents_chain.llm_chain.llm.model_name)),
+#                      "temperature": (str(chain.llm.temperature) if (rag_option == RAG_OFF) else
+#                                      str(chain.combine_documents_chain.llm_chain.llm.temperature)),
+#                      "retriever": ("" if (rag_option == RAG_OFF) else str(chain.retriever)),
+#                     } if (str(err_msg) == "") else {},
+#        start_time_ms = start_time_ms,
+#        end_time_ms = end_time_ms
+#    )
+#    
+#    trace.log("evaluation")
+#    wandb.finish()
 
 def invoke(openai_api_key, rag_option, prompt):
     if (openai_api_key == ""):