Update app.py

app.py

@@ -122,12 +122,9 @@ def wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_t
         result = completion
     else:
         result = completion["result"]
-        document_0_page_content = completion["source_documents"][0]["page_content"]
-        document_0_metadata = completion["source_documents"][0]["metadata"]
-        document_1_page_content = completion["source_documents"][1]["page_content"]
-        document_1_metadata = completion["source_documents"][1]["metadata"]
-        document_2_page_content = completion["source_documents"][2]["page_content"]
-        document_2_metadata = completion["source_documents"][2]["metadata"]
+        document_0 = completion["source_documents"][0]
+        document_1 = completion["source_documents"][1]
+        document_2 = completion["source_documents"][2]
     trace = Trace(
         kind = "chain",
         name = "LLMChain" if (rag_option == "Off") else "RetrievalQA",
@@ -144,12 +141,9 @@ def wandb_trace(rag_option, prompt, completion, status_msg, start_time_ms, end_t
                   "prompt": str(prompt if (str(status_msg) == "") else ""), 
                   "prompt_template": str((llm_template if (rag_option == "Off") else rag_template) if (str(status_msg) == "") else "")},
         outputs = {"result": str(result),
-                   "document_0_page_content": "" if (rag_option == "Off") else str(document_0_page_content),
-                   "document_0_metadata": "" if (rag_option == "Off") else str(document_0_metadata),
-                   "document_1_page_content": "" if (rag_option == "Off") else str(document_1_page_content),
-                   "document_1_metadata": "" if (rag_option == "Off") else str(document_1_metadata),
-                   "document_2_page_content": "" if (rag_option == "Off") else str(document_2_page_content),
-                   "document_2_metadata": "" if (rag_option == "Off") else str(document_2_metadata)},
+                   "document_0": "" if (rag_option == "Off") else str(document_0),
+                   "document_1": "" if (rag_option == "Off") else str(document_1),
+                   "document_2": "" if (rag_option == "Off") else str(document_2)},
         start_time_ms = start_time_ms,
         end_time_ms = end_time_ms
     )