Update app.py

app.py

@@ -82,18 +82,12 @@ def initialize_llmchain(vector_db, temperature=0.5, max_tokens=4096, top_k=3):
 def format_citation(source_doc):
     content = source_doc.page_content.strip()
     page = source_doc.metadata["page"] + 1
-    return f"[Page {page}] {content}"
-
-def format_response_with_citations(answer, sources):
-    citations = [format_citation(source) for source in sources[:3]]
-    formatted_response = f"{answer}\n\nReferences:\n"
-    for idx, citation in enumerate(citations, 1):
-        formatted_response += f"^{idx}^ {citation}\n"
-    return formatted_response
+    return content, page
 
 def conversation(qa_chain, message, history):
     if not qa_chain:
-        return None, gr.update(value=""), history, "Please upload a document first.", None
+        return (None, gr.update(value=""), history, "", 0, "", 0, "", 0, 
+                "Please upload a document first.")
         
     formatted_history = []
     for user_msg, bot_msg in history:
@@ -108,14 +102,22 @@ def conversation(qa_chain, message, history):
     answer = response["answer"]
     if "Helpful Answer:" in answer:
         answer = answer.split("Helpful Answer:")[-1]
-        
-    formatted_response = format_response_with_citations(
-        answer,
-        response["source_documents"]
-    )
+
+    # Format answer with citation numbers
+    sources = response["source_documents"][:3]
+    modified_answer = answer
+    for i in range(len(sources)):
+        modified_answer = modified_answer + f" [{i+1}]"
+
+    # Get citation contents and page numbers
+    citations = [format_citation(source) for source in sources]
+    source1_content, page1 = citations[0] if len(citations) > 0 else ("", 0)
+    source2_content, page2 = citations[1] if len(citations) > 1 else ("", 0)
+    source3_content, page3 = citations[2] if len(citations) > 2 else ("", 0)
     
-    new_history = history + [(message, formatted_response)]
-    return qa_chain, gr.update(value=""), new_history
+    new_history = history + [(message, modified_answer)]
+    return (qa_chain, gr.update(value=""), new_history, 
+            source1_content, page1, source2_content, page2, source3_content, page3, "")
 
 def demo():
     with gr.Blocks(theme=gr.themes.Default(primary_hue="red", secondary_hue="pink", neutral_hue="sky")) as demo:
@@ -133,13 +135,23 @@ def demo():
                 document = gr.Files(
                     height=300,
                     file_count="multiple",
-                    file_types=[".pdf"],  # Changed from 'pdf' to '.pdf'
+                    file_types=[".pdf"],
                     label="Upload PDF documents"
                 )
                 upload_status = gr.Textbox(label="Upload Status", interactive=False)
                 
             with gr.Column(scale=2):
-                chatbot = gr.Chatbot(height=600)
+                chatbot = gr.Chatbot(height=500)
+                with gr.Accordion("Citations", open=False):
+                    with gr.Row():
+                        doc_source1 = gr.Textbox(label="[1]", lines=2, container=True, scale=20)
+                        source1_page = gr.Number(label="Page", scale=1)
+                    with gr.Row():
+                        doc_source2 = gr.Textbox(label="[2]", lines=2, container=True, scale=20)
+                        source2_page = gr.Number(label="Page", scale=1)
+                    with gr.Row():
+                        doc_source3 = gr.Textbox(label="[3]", lines=2, container=True, scale=20)
+                        source3_page = gr.Number(label="Page", scale=1)
                 with gr.Row():
                     msg = gr.Textbox(
                         placeholder="Ask a question about your documents...",
@@ -167,16 +179,34 @@ def demo():
             outputs=[vector_db, qa_chain, upload_status]
         )
         
+        # Clear citations when chat is cleared
+        def clear_all():
+            return ["", 0, "", 0, "", 0]
+            
         # Chatbot events
         submit_btn.click(
             conversation,
             inputs=[qa_chain, msg, chatbot],
-            outputs=[qa_chain, msg, chatbot]
+            outputs=[qa_chain, msg, chatbot, 
+                    doc_source1, source1_page,
+                    doc_source2, source2_page,
+                    doc_source3, source3_page,
+                    upload_status]
         )
         msg.submit(
             conversation,
             inputs=[qa_chain, msg, chatbot],
-            outputs=[qa_chain, msg, chatbot]
+            outputs=[qa_chain, msg, chatbot,
+                    doc_source1, source1_page,
+                    doc_source2, source2_page,
+                    doc_source3, source3_page,
+                    upload_status]
+        )
+        clear_btn.click(
+            clear_all,
+            outputs=[doc_source1, source1_page,
+                    doc_source2, source2_page,
+                    doc_source3, source3_page]
         )
         
     demo.queue().launch(debug=True)