Update app.py

app.py

@@ -136,13 +136,15 @@ 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 = {"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 
+        model_dict = {"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 "",
+                      "llm_model_name": (str(chain.llm.model_name) if (rag_option == RAG_OFF) else
+                                         str(chain.combine_documents_chain.llm_chain.llm.model_name)) if (str(err_msg) == "") else "",
+                      "llm_temperature": (str(chain.llm.temperature) if (rag_option == RAG_OFF) else
+                                          str(chain.combine_documents_chain.llm_chain.llm.temperature)) 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 "",
+                      "chain_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,7 +189,6 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
-        print(chain)
         wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_time_ms, end_time_ms)
     return result