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, llm_output, chain, err_msg, start_time_ms, end_time_ms):
+def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms):
     wandb.init(project = "openai-llm-rag")
     trace = Trace(
         kind = "chain",
@@ -134,6 +134,7 @@ def wandb_trace(rag_option, prompt, completion, result, llm_output, chain, err_m
                   "prompt": prompt if (str(err_msg) == "") else "",
         },
         outputs = {"result": result if (str(err_msg) == "") else "",
+                   "generation_info": generation_info if (str(err_msg) == "") else "",
                    "llm_output": llm_output if (str(err_msg) == "") else "",
                    "completion": str(completion) if (str(err_msg) == "") else "",
         },
@@ -186,7 +187,8 @@ def invoke(openai_api_key, rag_option, prompt):
             completion, chain = llm_chain(llm, prompt)
             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
+                generation_info = completion.generations[0][0].generation_info
+                llm_output = completion.generations.llm_output
                 print(completion)
                 print("###")
                 print(completion.generations[0])
@@ -195,7 +197,7 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
-        wandb_trace(rag_option, prompt, completion, result, llm_output, chain, err_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms)
     return result
 
 gr.close_all()