Update app.py

app.py

@@ -47,10 +47,8 @@ template = """If you don't know the answer, just say that you don't know, don't
 llm_template = "Answer the question at the end. " + template + "Question: {question} Helpful Answer: "
 rag_template = "Use the following pieces of context to answer the question at the end. " + template + "{context} Question: {question} Helpful Answer: "
 
-LLM_CHAIN_PROMPT = PromptTemplate(input_variables = ["question"], 
-                                  template = llm_template)
-RAG_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"], 
-                                  template = rag_template)
+LLM_CHAIN_PROMPT = PromptTemplate(input_variables = ["question"], template = llm_template)
+RAG_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"], template = rag_template)
 
 CHROMA_DIR  = "/data/chroma"
 YOUTUBE_DIR = "/data/youtube"
@@ -118,10 +116,10 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion
 
-def wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_name, status_msg, start_time_ms, end_time_ms):
+def wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_time_ms):
     wandb.init(project = "openai-llm-rag")
     trace = Trace(
-        name = chain_name,
+        name = "LLMChain" if (rag_option == "Off") else "RetrievalQA",
         kind = "chain",
         status_code = "SUCCESS" if (str(status_msg) == "") else "ERROR",
         status_message = str(status_msg),
@@ -134,8 +132,8 @@ def wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_n
         },
         start_time_ms = start_time_ms,
         end_time_ms = end_time_ms,
-        inputs = {"rag_option": rag_option, "prompt": str(prompt), "prompt_template": str(prompt_template)},
-        outputs = {"result": str(result), "completion": str(completion)},
+        inputs = {"rag_option": rag_option, "prompt": str(prompt), "prompt_template": str(llm_template if (rag_option == "Off") else rag_template)},
+        outputs = {"completion": str(completion)},
         #model_dict = {"llm": str(llm)}
     )
     trace.log("test")
@@ -150,8 +148,6 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error("Prompt is required.")
     completion = ""
     result = ""
-    prompt_template = ""
-    chain_name = ""
     status_msg = ""
     try:
         start_time_ms = round(time.time() * 1000)
@@ -164,26 +160,20 @@ def invoke(openai_api_key, rag_option, prompt):
             db = document_retrieval_chroma(llm, prompt)
             completion = rag_chain(llm, prompt, db)
             result = completion["result"]
-            prompt_template = rag_template
-            chain_name = "RetrievalQA"
         elif (rag_option == "MongoDB"):
             #splits = document_loading_splitting()
             #document_storage_mongodb(splits)
             db = document_retrieval_mongodb(llm, prompt)
             completion = rag_chain(llm, prompt, db)
             result = completion["result"]
-            prompt_template = rag_template
-            chain_name = "RetrievalQA"
         else:
             result = llm_chain(llm, prompt)
-            prompt_template = llm_template
-            chain_name = "LLMChain"
     except Exception as e:
         status_msg = e
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
-        wandb_trace(rag_option, prompt, prompt_template, result, completion, chain_name, status_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()