Update app.py

app.py

@@ -8,6 +8,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
 from langchain.embeddings import HuggingFaceEmbeddings
 import fitz  # PyMuPDF
+import base64
 
 load_dotenv()
 client = groq.Client(api_key=os.getenv("GROQ_FINANCE_API_KEY"))
@@ -30,31 +31,42 @@ custom_css = """
     --border-grey: #E5E7EB;
 }
 body { background-color: var(--light-grey); font-family: 'Inter', sans-serif; }
-.container { max-width: 900px !important; margin: 0 auto !important; }
+.container { max-width: 1200px !important; margin: 0 auto !important; padding: 10px; }
 .header { background-color: var(--white); border-bottom: 2px solid var(--border-grey); padding: 15px 0; margin-bottom: 20px; border-radius: 12px 12px 0 0; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
-.header-title { color: var(--dark-green); font-size: 1.5rem; font-weight: 700; text-align: center; }
-.header-subtitle { color: var(--medium-grey); font-size: 0.9rem; text-align: center; margin-top: 5px; }
-.chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-grey) !important; }
+.header-title { color: var(--dark-green); font-size: 1.8rem; font-weight: 700; text-align: center; }
+.header-subtitle { color: var(--medium-grey); font-size: 1rem; text-align: center; margin-top: 5px; }
+.chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-grey) !important; min-height: 500px; }
 .message-user { background-color: var(--primary-green) !important; color: var(--white) !important; border-radius: 18px 18px 4px 18px !important; padding: 12px 16px !important; margin-left: auto !important; max-width: 80% !important; }
 .message-bot { background-color: var(--light-grey) !important; color: var(--medium-grey) !important; border-radius: 18px 18px 18px 4px !important; padding: 12px 16px !important; margin-right: auto !important; max-width: 80% !important; }
 .input-area { background-color: var(--white) !important; border-top: 1px solid var(--border-grey) !important; padding: 12px !important; border-radius: 0 0 12px 12px !important; }
 .input-box { border: 1px solid var(--border-grey) !important; border-radius: 24px !important; padding: 12px 16px !important; box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; }
 .send-btn { background-color: var(--primary-green) !important; border-radius: 24px !important; color: var(--white) !important; padding: 10px 20px !important; font-weight: 500 !important; }
 .clear-btn { background-color: var(--light-grey) !important; border: 1px solid var(--border-grey) !important; border-radius: 24px !important; color: var(--medium-grey) !important; padding: 8px 16px !important; font-weight: 500 !important; }
+.pdf-viewer-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-grey) !important; padding: 20px; }
+.pdf-viewer-image { max-width: 100%; height: auto; border: 1px solid var(--border-grey); border-radius: 12px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
+.stats-box { background-color: var(--light-green); padding: 10px; border-radius: 8px; margin-top: 10px; }
 """
 
 def process_pdf(pdf_file):
     if pdf_file is None:
-        return None, "No file uploaded"
+        return None, "No file uploaded", {"page_images": [], "total_pages": 0, "total_words": 0}
     try:
         session_id = str(uuid.uuid4())
         with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
             temp_file.write(pdf_file)
             pdf_path = temp_file.name
         
-        # Use fitz to extract text from PDF
+        # Use fitz to extract text and images from PDF
         doc = fitz.open(pdf_path)
         texts = [page.get_text() for page in doc]
+        page_images = []
+        for page in doc:
+            pix = page.get_pixmap()
+            img_bytes = pix.tobytes("png")
+            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
+            page_images.append(img_base64)
+        total_pages = len(doc)
+        total_words = sum(len(text.split()) for text in texts)
         doc.close()
 
         # Split the extracted text into chunks
@@ -68,11 +80,12 @@ def process_pdf(pdf_file):
         user_vectorstores[session_id] = vectorstore
 
         os.unlink(pdf_path)
-        return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF"
+        pdf_state = {"page_images": page_images, "total_pages": total_pages, "total_words": total_words}
+        return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF", pdf_state
     except Exception as e:
         if "pdf_path" in locals() and os.path.exists(pdf_path):
             os.unlink(pdf_path)
-        return None, f"Error processing PDF: {str(e)}"
+        return None, f"Error processing PDF: {str(e)}", {"page_images": [], "total_pages": 0, "total_words": 0}
 
 def generate_response(message, session_id, model_name, history):
     if not message:
@@ -103,16 +116,27 @@ def generate_response(message, session_id, model_name, history):
         history.append((message, f"Error generating response: {str(e)}"))
         return history
 
+def update_pdf_viewer(pdf_state):
+    if not pdf_state["total_pages"]:
+        return 0, None, "No PDF uploaded yet"
+    return pdf_state["total_pages"], f"data:image/png;base64,{pdf_state['page_images'][0]}", f"**Total Pages:** {pdf_state['total_pages']}\n**Total Words:** {pdf_state['total_words']}"
+
+def update_image(page_num, pdf_state):
+    if not pdf_state["total_pages"] or page_num < 1 or page_num > pdf_state["total_pages"]:
+        return None
+    return f"data:image/png;base64,{pdf_state['page_images'][page_num - 1]}"
+
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     current_session_id = gr.State(None)
+    pdf_state = gr.State({"page_images": [], "total_pages": 0, "total_words": 0})
     gr.HTML("""
     <div class="header">
         <div class="header-title">Financial Report Summarizer</div>
         <div class="header-subtitle">Summarize financial data with Groq's LLM API</div>
     </div>
     """)
-    with gr.Row():
-        with gr.Column(scale=1):
+    with gr.Row(elem_classes="container"):
+        with gr.Column(scale=1, min_width=300):
             pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary")
             upload_button = gr.Button("Process PDF", variant="primary")
             pdf_status = gr.Markdown("No PDF uploaded yet")
@@ -121,17 +145,54 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                 value="llama3-70b-8192",
                 label="Select Groq Model"
             )
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container")
-            with gr.Row():
-                msg = gr.Textbox(show_label=False, placeholder="Ask about your financial report...", scale=5)
-                send_btn = gr.Button("Send", scale=1)
-            clear_btn = gr.Button("Clear Conversation")
+        with gr.Column(scale=2, min_width=600):
+            with gr.Tabs():
+                with gr.TabItem("Chatbot"):
+                    chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container")
+                    with gr.Row():
+                        msg = gr.Textbox(show_label=False, placeholder="Ask about your financial report...", scale=5)
+                        send_btn = gr.Button("Send", scale=1)
+                    clear_btn = gr.Button("Clear Conversation")
+                with gr.TabItem("PDF Viewer"):
+                    with gr.Column(elem_classes="pdf-viewer-container"):
+                        page_slider = gr.Slider(minimum=1, maximum=0, step=1, label="Page Number", value=1)
+                        pdf_image = gr.Image(label="PDF Page", type="pil", elem_classes="pdf-viewer-image")
+                        stats_display = gr.Markdown("No PDF uploaded yet", elem_classes="stats-box")
+    
+    # Event Handlers
+    upload_button.click(
+        process_pdf,
+        inputs=[pdf_file],
+        outputs=[current_session_id, pdf_status, pdf_state]
+    ).then(
+        update_pdf_viewer,
+        inputs=[pdf_state],
+        outputs=[page_slider, pdf_image, stats_display]
+    )
+    
+    msg.submit(
+        generate_response,
+        inputs=[msg, current_session_id, model_dropdown, chatbot],
+        outputs=[chatbot]
+    ).then(lambda: "", None, [msg])
+    
+    send_btn.click(
+        generate_response,
+        inputs=[msg, current_session_id, model_dropdown, chatbot],
+        outputs=[chatbot]
+    ).then(lambda: "", None, [msg])
+    
+    clear_btn.click(
+        lambda: ([], None, "No PDF uploaded yet", {"page_images": [], "total_pages": 0, "total_words": 0}, 0, None, "No PDF uploaded yet"),
+        None,
+        [chatbot, current_session_id, pdf_status, pdf_state, page_slider, pdf_image, stats_display]
+    )
     
-    upload_button.click(process_pdf, inputs=[pdf_file], outputs=[current_session_id, pdf_status])
-    msg.submit(generate_response, inputs=[msg, current_session_id, model_dropdown, chatbot], outputs=[chatbot]).then(lambda: "", None, [msg])
-    send_btn.click(generate_response, inputs=[msg, current_session_id, model_dropdown, chatbot], outputs=[chatbot]).then(lambda: "", None, [msg])
-    clear_btn.click(lambda: ([], None, "No PDF uploaded yet"), None, [chatbot, current_session_id, pdf_status])
+    page_slider.change(
+        update_image,
+        inputs=[page_slider, pdf_state],
+        outputs=[pdf_image]
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()