Update app.py

app.py

@@ -123,9 +123,28 @@ def upload_pdf(pdf_file):
 
 
 def respond(message, history):
-    bot_message = pdf_chatbot.chat(message)
-    history.append((message, bot_message))
-    return "", history
+    if not pdf_chatbot.qa_chain:
+        return "", history, "", "", "", "", "", ""
+
+    # Generate response using QA chain
+    response = pdf_chatbot.chat(message)
+    response_answer = response["answer"]
+    if response_answer.find("Helpful Answer:") != -1:
+        response_answer = response_answer.split("Helpful Answer:")[-1]
+    response_sources = response["source_documents"]
+
+    # Extract source documents and page numbers
+    response_source1 = response_sources[0].page_content.strip() if len(response_sources) > 0 else ""
+    response_source2 = response_sources[1].page_content.strip() if len(response_sources) > 1 else ""
+    response_source3 = response_sources[2].page_content.strip() if len(response_sources) > 2 else ""
+    response_source1_page = response_sources[0].metadata["page"] + 1 if len(response_sources) > 0 else ""
+    response_source2_page = response_sources[1].metadata["page"] + 1 if len(response_sources) > 1 else ""
+    response_source3_page = response_sources[2].metadata["page"] + 1 if len(response_sources) > 2 else ""
+
+    # Append user message and response to chat history
+    history.append((message, response_answer))
+
+    return "", history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
 
 def clear_chatbot():
     pdf_chatbot.memory.clear()
@@ -153,7 +172,18 @@ with gr.Blocks() as demo:
     msg = gr.Textbox()
     clear = gr.Button("Clear")
 
-    msg.submit(respond, inputs=[msg, chatbot_interface], outputs=[msg, chatbot_interface])
+    with gr.Accordion("Advanced - Document references", open=False):
+        with gr.Row():
+            doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
+            source1_page = gr.Number(label="Page", scale=1)
+        with gr.Row():
+            doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
+            source2_page = gr.Number(label="Page", scale=1)
+        with gr.Row():
+            doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
+            source3_page = gr.Number(label="Page", scale=1)
+
+    msg.submit(respond, inputs=[msg, chatbot_interface], outputs=[msg, chatbot_interface, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page])
     clear.click(clear_chatbot, outputs=[chatbot_interface])
     path_button.click(get_pdf_path, outputs=[pdf_path_display])