Update app.py

app.py

@@ -120,9 +120,11 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
-def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms, end_time_ms):
-    if (rag_option == RAG_OFF or str(status_msg) != ""):
-        result = completion
+def wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, end_time_ms):
+    if (str(err_msg) != ""):
+        result = ""
+    else if (rag_option == RAG_OFF):
+        result = completion.text
     else:
         result = completion["result"]
         docs_meta = str([doc.metadata for doc in completion["source_documents"]])
@@ -130,8 +132,8 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
     trace = Trace(
         kind = "chain",
         name = type(chain).__name__ if (chain != None) else "",
-        status_code = "success" if (str(status_msg) == "") else "error",
-        status_message = str(status_msg),
+        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"],
@@ -139,11 +141,12 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
             "model": config["model"],
             "temperature": config["temperature"],
         },
-        inputs = {"rag_option": rag_option if (str(status_msg) == "") else "",
-                  "prompt": str(prompt if (str(status_msg) == "") else ""),
-                  "prompt_template": str((llm_template if (rag_option == RAG_OFF) else rag_template) if (str(status_msg) == "") else ""),
-                  "docs_meta": "" if (rag_option == RAG_OFF or str(status_msg) != "") else docs_meta},
-        outputs = {"result": result},
+        inputs = {"rag_option": rag_option if (str(err_msg) == "") else "",
+                  "prompt": str(prompt if (str(err_msg) == "") else ""),
+                  "prompt_template": str((llm_template if (rag_option == RAG_OFF) else rag_template) if (str(err_msg) == "") else ""),
+                  "docs_meta": "" if (str(err_msg) != "" or rag_option == RAG_OFF) else docs_meta},
+        outputs = {"result": result,
+                   "completion": completion},
         start_time_ms = start_time_ms,
         end_time_ms = end_time_ms
     )
@@ -160,7 +163,7 @@ def invoke(openai_api_key, rag_option, prompt):
     completion = ""
     result = ""
     chain = None
-    status_msg = ""
+    err_msg = ""
     try:
         start_time_ms = round(time.time() * 1000)
         llm = ChatOpenAI(model_name = config["model"], 
@@ -179,15 +182,14 @@ def invoke(openai_api_key, rag_option, prompt):
             completion, chain = rag_chain(llm, prompt, db)
             result = completion["result"]
         else:
-            result, chain = llm_chain(llm, prompt)
-            print(result)
-            completion = result
+            completion, chain = llm_chain(llm, prompt)
+            result = completion.text
     except Exception as e:
-        status_msg = e
+        err_msg = e
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
-        wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, chain, err_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()