Update app.py

app.py

@@ -124,11 +124,11 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
             print("1=" + str(chain.llm.temperature))
             print("2=" + str(chain.prompt.input_variables))
             print("2=" + str(chain.prompt.template))
-        if (type(chain).__name__ == "RetrievalQA"):
+        #if (type(chain).__name__ == "RetrievalQA"):
             #print(chain.combine_documents_chain.llm_chain.llm)
             #print(chain.combine_documents_chain.llm_chain.async_client)
-            print(chain.combine_documents_chain.retriever)
-            print(chain.combine_documents_chain.vectorstore)
+            #print(chain.combine_documents_chain.retriever)
+            #print(chain.combine_documents_chain.vectorstore)
             #print("3=" + str(chain.llm))
             #print("4=" + str(chain.chain_type_kwargs))
             #print("5=" + str(chain.retriever))
@@ -137,6 +137,7 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
         result = completion
     else:
         result = completion["result"]
+        documents = completion["source_documents"]
         document_0 = completion["source_documents"][0]
         document_1 = completion["source_documents"][1]
         document_2 = completion["source_documents"][2]
@@ -155,6 +156,7 @@ def wandb_trace(rag_option, prompt, completion, chain, status_msg, start_time_ms
         inputs = {"rag_option": rag_option if (str(status_msg) == "") else "",
                   "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 ""),
+                  "documents": "" if (rag_option == "Off" or str(status_msg) != "") else str(documents),
                   "document_0": "" if (rag_option == "Off" or str(status_msg) != "") else str(document_0),
                   "document_1": "" if (rag_option == "Off" or str(status_msg) != "") else str(document_1),
                   "document_2": "" if (rag_option == "Off" or str(status_msg) != "") else str(document_2)},