Update app.py

app.py

@@ -136,13 +136,13 @@ def wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_ti
         outputs = {"result": result if (str(err_msg) == "") else "",
                    "completion": str(completion) if (str(err_msg) == "") else "",
         },
-        model_dict = {"llm_client": str(chain.llm.client) if (rag_option == RAG_OFF) else 
-                                    str(chain.combine_documents_chain.llm_chain.llm.client),
-                      "model_name": config["model_name"],
-                      "temperature": config["temperature"],
-                      "chain_prompt": str(chain.prompt) if (rag_option == RAG_OFF) else 
-                                      str(chain.combine_documents_chain.llm_chain.prompt),
-                      "retriever": "" if (rag_option == RAG_OFF) else str(chain.retriever),
+        model_dict = {"model_name": config["model_name"] if (str(err_msg) == "") else "",
+                      "temperature": config["temperature"] if (str(err_msg) == "") else "",
+                      "llm_client": (str(chain.llm.client) if (rag_option == RAG_OFF) else 
+                                     str(chain.combine_documents_chain.llm_chain.llm.client) if (str(err_msg) == "") else "",
+                      "chain_prompt": (str(chain.prompt) if (rag_option == RAG_OFF) else 
+                                       str(chain.combine_documents_chain.llm_chain.prompt)) if (str(err_msg) == "") else "",
+                      "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
@@ -187,6 +187,7 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
+        print(completion)
         wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_time_ms, end_time_ms)
     return result