Update app.py

app.py

@@ -1,174 +1,150 @@
-import os
-import base64
 import streamlit as st
+import os
+import pandas as pd
+import fitz  # PyMuPDF
+from google import genai
+from google.genai import types
 from docx import Document
-import openpyxl
-from pdfminer.high_level import extract_text
-import csv
 from pptx import Presentation
-from io import StringIO
-from bs4 import BeautifulSoup
-import re
-import google.generativeai as genai
-
-# Setup
-st.set_page_config(page_title="Gemini Q&A App", layout="centered")
-st.title("🤖 Q&A - Document & Image")
-
-# Initialize chat history
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
-
-# Optional: Clear history button
-if st.button("🧹 Clear Chat History"):
-    st.session_state.chat_history = []
-
-# Load API key
-GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY") or st.secrets.get("GOOGLE_API_KEY")
-if not GOOGLE_API_KEY:
-    st.error("GOOGLE_API_KEY not found in environment or Streamlit secrets.")
-    st.stop()
-
-genai.configure(api_key=GOOGLE_API_KEY)
-model = genai.GenerativeModel(model_name="gemini-2.0-flash")
-
-# Helper: Extract text from various documents
-def load_and_extract_text(file_path):
-    try:
-        ext = file_path.split('.')[-1].lower()
-
-        if ext == 'docx':
-            doc = Document(file_path)
-            return '\n'.join(p.text for p in doc.paragraphs)
-
-        elif ext == 'xlsx':
-            workbook = openpyxl.load_workbook(file_path)
-            return '\n'.join(
-                ' '.join(str(cell.value) for cell in row if cell.value is not None)
-                for sheet in workbook for row in sheet.iter_rows()
-            )
-
-        elif ext == 'pdf':
-            return extract_text(file_path)
-
-        elif ext == 'csv':
-            with open(file_path, 'r', encoding='utf-8') as f:
-                reader = csv.reader(f)
-                return '\n'.join(' '.join(row) for row in reader)
-
-        elif ext == 'txt':
-            with open(file_path, 'r', encoding='utf-8') as f:
-                return f.read()
-
-        elif ext == 'pptx':
-            prs = Presentation(file_path)
-            return '\n'.join(
-                shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, 'text')
-            )
-
-        elif ext == 'tex':
-            with open(file_path, 'r', encoding='utf-8') as f:
-                content = f.read()
-                content = re.sub(r'\\[a-zA-Z]+\*?(?:\[[^\]]*\])?(?:\{[^}]*\})*', '', content)
-                content = re.sub(r'%.*', '', content)
-                return content.strip()
-
-        elif ext in ['html', 'htm']:
-            with open(file_path, 'r', encoding='utf-8') as f:
-                soup = BeautifulSoup(f, 'html.parser')
-                for script_or_style in soup(['script', 'style']):
-                    script_or_style.decompose()
-                return soup.get_text(separator='\n', strip=True)
-
-        else:
-            return "Unsupported file format."
-
-    except Exception as e:
-        return f"Error processing file: {e}"
-
-# Multi-file uploader
-uploaded_files = st.file_uploader(
-    "Upload one or more documents or images",
-    type=["docx", "xlsx", "pdf", "csv", "txt", "pptx", "tex", "html", "htm", "jpg", "jpeg", "png"],
-    accept_multiple_files=True
-)
-
-# Question input
-question = st.text_input("💬 Enter your question:")
-
-# Processing
-if uploaded_files and question:
-    combined_context = ""
-    image_contents = []
-
-    for uploaded_file in uploaded_files:
-        file_ext = uploaded_file.name.split('.')[-1].lower()
-
-        if file_ext in ['jpg', 'jpeg', 'png']:
-            image_bytes = uploaded_file.read()
-            encoded_image = base64.b64encode(image_bytes).decode("utf-8")
-            image_contents.append({
-                "inline_data": {
-                    "mime_type": uploaded_file.type,
-                    "data": encoded_image
-                }
-            })
-            st.image(image_bytes, caption=uploaded_file.name, use_container_width=True)
-            continue
-
-        temp_file_path = "temp." + file_ext
-        with open(temp_file_path, "wb") as temp_file:
-            temp_file.write(uploaded_file.read())
-
-        extracted_text = load_and_extract_text(temp_file_path)
-        os.remove(temp_file_path)
-
-        if "Error" in extracted_text or "Unsupported" in extracted_text:
-            st.error(f"{uploaded_file.name}: {extracted_text}")
-        else:
-            combined_context += f"\n\n---\nDocument: {uploaded_file.name}\n\n{extracted_text}"
-
-    if combined_context.strip() or image_contents:
-        with st.spinner("Generating answer..."):
-            try:
-                content_blocks = []
-
-                # Add prior Q&A as part of the context
-                if st.session_state.chat_history:
-                    history_text = "\n\n".join(
-                        f"Q: {entry['question']}\nA: {entry['answer']}"
-                        for entry in st.session_state.chat_history
-                    )
-                    content_blocks.append({"text": f"Previous conversation:\n{history_text}"})
-
-                # Add file text content
-                if combined_context.strip():
-                    content_blocks.append({"text": f"Context:\n{combined_context}"})
-
-                # Add image blocks
-                content_blocks.extend(image_contents)
-
-                # Add current question
-                content_blocks.append({"text": f"Question: {question}"})
-
-                response = model.generate_content(contents=content_blocks)
-                st.success("💡 Answer:")
-                st.write(response.text)
-
-                # Save to session history
-                st.session_state.chat_history.append({
-                    "question": question,
-                    "answer": response.text
-                })
-
-            except Exception as e:
-                st.error(f"Error generating answer: {e}")
-    else:
-        st.warning("No valid documents or images processed.")
-
-# Optional: Show full conversation history
-if st.session_state.chat_history:
-    st.markdown("### 🗂️ Chat History")
-    for i, entry in enumerate(st.session_state.chat_history, 1):
-        st.markdown(f"**Q{i}:** {entry['question']}")
-        st.markdown(f"**A{i}:** {entry['answer']}")
+from reportlab.lib.pagesizes import A4
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
+from reportlab.lib.styles import getSampleStyleSheet
+
+# Set Gemini API Key
+os.environ["GEMINI_API_KEY"] = "AIzaSyAQLUStrjXv2_zrAjssapa_sYXQWOqVeiE"  # Ensure this is set securely in Hugging Face environment
+client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
+
+# Function to extract text from different file formats
+def extract_text_from_docx(docx_file):
+    document = Document(docx_file)
+    return "\n".join([para.text for para in document.paragraphs])
+
+def extract_text_from_csv(csv_file):
+    df = pd.read_csv(csv_file)
+    return df.to_string(index=False)
+
+def extract_text_from_xlsx(xlsx_file):
+    df = pd.read_excel(xlsx_file)
+    return df.to_string(index=False)
+
+def extract_text_from_pptx(pptx_file):
+    prs = Presentation(pptx_file)
+    text_runs = []
+    for slide in prs.slides:
+        for shape in slide.shapes:
+            if hasattr(shape, "text"):
+                text_runs.append(shape.text)
+    return "\n".join(text_runs)
+
+def extract_text_from_pdf(pdf_file):
+    doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
+    text = ""
+    for page in doc:
+        text += page.get_text()
+    return text
+
+def export_conversation_to_pdf(conversation_history):
+    pdf_path = "conversation_clean.pdf"
+    doc = SimpleDocTemplate(pdf_path, pagesize=A4)
+    styles = getSampleStyleSheet()
+    elements = []
+
+    for i, (user_q, gemini_a) in enumerate(conversation_history):
+        elements.append(Paragraph(f"<b>Q{i+1}:</b> {user_q}", styles["Normal"]))
+        elements.append(Spacer(1, 8))
+        elements.append(Paragraph(f"<b>A{i+1}:</b> {gemini_a.replace(chr(10), '<br/>')}", styles["Normal"]))
+        elements.append(Spacer(1, 16))
+
+    doc.build(elements)
+    return pdf_path
+
+def main():
+    st.set_page_config(page_title="Gemini Chat QA", layout="wide")
+    st.title("📄 Chat with Your Document using Gemini (Sayandip+Dhiraj+Fazal)")
+
+    if "conversation" not in st.session_state:
+        st.session_state.conversation = []
+    if "document_text" not in st.session_state:
+        st.session_state.document_text = ""
+    if "chat_active" not in st.session_state:
+        st.session_state.chat_active = True
+
+    uploaded_file = st.file_uploader("Upload a file (.txt, .docx, .csv, .xlsx, .pptx, .pdf):",
+                                     type=["txt", "docx", "csv", "xlsx", "pptx", "pdf"])
+
+    if uploaded_file and not st.session_state.document_text:
+        filetype = uploaded_file.name.split(".")[-1].lower()
+        try:
+            if filetype == "txt":
+                st.session_state.document_text = uploaded_file.read().decode("utf-8")
+            elif filetype == "docx":
+                st.session_state.document_text = extract_text_from_docx(uploaded_file)
+            elif filetype == "csv":
+                st.session_state.document_text = extract_text_from_csv(uploaded_file)
+            elif filetype == "xlsx":
+                st.session_state.document_text = extract_text_from_xlsx(uploaded_file)
+            elif filetype == "pptx":
+                st.session_state.document_text = extract_text_from_pptx(uploaded_file)
+            elif filetype == "pdf":
+                st.session_state.document_text = extract_text_from_pdf(uploaded_file)
+            else:
+                st.error("Unsupported file format.")
+        except Exception as e:
+            st.error(f"Failed to extract text: {e}")
+
+    if st.session_state.document_text:
+        st.markdown("### 💬 Conversation")
+        for user_q, gemini_a in st.session_state.conversation:
+            st.markdown(f"**You:** {user_q}")
+            st.markdown(f"**Gemini:**\\n\\n{gemini_a}")
+
+        if st.session_state.chat_active:
+            with st.form(key="chat_form", clear_on_submit=True):
+                user_input = st.text_input("Ask a question about the document (type 'exit' to stop):")
+                submit = st.form_submit_button("Send")
+
+                if submit and user_input:
+                    if user_input.strip().lower() == "exit":
+                        st.session_state.chat_active = False
+                        st.success("Chat ended. Reload to start again.")
+                    else:
+                        with st.spinner("Gemini is thinking..."):
+                            contents = [
+                                types.Content(
+                                    role="user",
+                                    parts=[
+                                        types.Part.from_text(text=st.session_state.document_text),
+                                        types.Part.from_text(text=user_input),
+                                        types.Part.from_text(text="If needed, search the web for more context.")
+                                    ]
+                                )
+                            ]
+                            tools = [types.Tool(google_search=types.GoogleSearch())]
+                            generate_config = types.GenerateContentConfig(
+                                tools=tools,
+                                response_mime_type="text/plain",
+                            )
+
+                            response_text = ""
+                            for chunk in client.models.generate_content_stream(
+                                model="gemini-2.0-flash",
+                                contents=contents,
+                                config=generate_config,
+                            ):
+                                response_text += chunk.text
+
+                            st.session_state.conversation.append((user_input, response_text))
+                            st.rerun()
+
+        if st.session_state.conversation:
+            pdf_path = export_conversation_to_pdf(st.session_state.conversation)
+            with open(pdf_path, "rb") as f:
+                st.download_button(
+                    label="📥 Download Cleaned Conversation as PDF",
+                    data=f,
+                    file_name="clean_conversation.pdf",
+                    mime="application/pdf"
+                )
+
+if __name__ == "__main__":
+    main()