Update app.py

app.py

@@ -119,7 +119,7 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
-def wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_time_ms, end_time_ms):
+def wandb_trace(rag_option, prompt, completion, result, llm_output, chain, err_msg, start_time_ms, end_time_ms):
     wandb.init(project = "openai-llm-rag")
     trace = Trace(
         kind = "chain",
@@ -134,14 +134,15 @@ def wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_ti
                   "prompt": prompt if (str(err_msg) == "") else "",
         },
         outputs = {"result": result if (str(err_msg) == "") else "",
+                   "llm_output": llm_output 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)) 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": (chain.llm.temperature if (rag_option == RAG_OFF) else
-                                          chain.combine_documents_chain.llm_chain.llm.temperature) 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 "",
                       "chain_retriever": ("" if (rag_option == RAG_OFF) else str(chain.retriever)) if (str(err_msg) == "") else "",
@@ -162,6 +163,7 @@ def invoke(openai_api_key, rag_option, prompt):
     chain = None
     completion = ""
     result = ""
+    llm_output = ""
     err_msg = ""
     try:
         start_time_ms = round(time.time() * 1000)
@@ -182,14 +184,15 @@ def invoke(openai_api_key, rag_option, prompt):
             result = completion["result"]
         else:
             completion, chain = llm_chain(llm, prompt)
-            result = completion.generations[0][0].text if (completion.generations[0] != None and 
-                                                           completion.generations[0][0] != None) else ""
+            if (completion.generations[0] != None and completion.generations[0][0] != None):
+                result = completion.generations[0][0].text
+                llm_output = completion.generations[0][0].llm_output
     except Exception as e:
         err_msg = e
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
-        wandb_trace(rag_option, prompt, completion, result, chain, err_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, result, llm_output, chain, err_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()