Update app.py

app.py

@@ -84,115 +84,112 @@ def export_conversation_to_pdf(conversation_history):
     return pdf_path
 
 def main():
-    # Initialize session state variables
+    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
-    if "chat_history" not in st.session_state:
-        st.session_state.chat_history = []  # Initialize chat history
-
-    st.set_page_config(page_title="Gemini Chat QA", layout="wide")
-    st.title("📄 Chat with Your Document using Gemini (Sayandip+Dhiraj+Fazal)")
 
-    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
+    uploaded_file = st.file_uploader(
+        "Upload a file (.txt, .docx, .csv, .xlsx, .pptx, .pdf, .html, .htm, .tex, .jpg, .jpeg, .png):",
+        type=["txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex", "jpg", "jpeg", "png"]
     )
 
-    if uploaded_files:
-        combined_context = ""
-        image_contents = []
-        
-        for uploaded_file in uploaded_files:
-            file_ext = uploaded_file.name.split('.')[-1].lower()
-
-            try:
-                if file_ext == "txt":
-                    combined_context += uploaded_file.read().decode("utf-8")
-                elif file_ext == "docx":
-                    combined_context += extract_text_from_docx(uploaded_file)
-                elif file_ext == "csv":
-                    combined_context += extract_text_from_csv(uploaded_file)
-                elif file_ext == "xlsx":
-                    combined_context += extract_text_from_xlsx(uploaded_file)
-                elif file_ext == "pptx":
-                    combined_context += extract_text_from_pptx(uploaded_file)
-                elif file_ext == "pdf":
-                    combined_context += extract_text_from_pdf(uploaded_file)
-                elif file_ext in ["html", "htm"]:
-                    combined_context += extract_text_from_html(uploaded_file)
-                elif file_ext == "tex":
-                    combined_context += extract_text_from_tex(uploaded_file)
-                elif file_ext in ["jpg", "jpeg", "png"]:
-                    image_data = process_image(uploaded_file)
-                    image_contents.append(image_data)
-                    st.image(uploaded_file, caption=uploaded_file.name, use_container_width=True)
-
-            except Exception as e:
-                st.error(f"Error processing {uploaded_file.name}: {e}")
-
-        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.conversation:
-                        history_text = "\n\n".join(
-                            f"Q: {entry[0]}\nA: {entry[1]}"
-                            for entry in st.session_state.conversation
-                        )
-                        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
-                    user_input = st.text_input("Ask a question about the documents (type 'exit' to stop):")
+    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)
+            elif filetype in ["html", "htm"]:
+                st.session_state.document_text = extract_text_from_html(uploaded_file)
+            elif filetype == "tex":
+                st.session_state.document_text = extract_text_from_tex(uploaded_file)
+            elif filetype in ["jpg", "jpeg", "png"]:
+                image_data = process_image(uploaded_file)
+                st.session_state.document_text = "Image uploaded and processed."
+                st.image(uploaded_file, caption=uploaded_file.name, use_container_width=True)
+                st.session_state.image_data = image_data  # Store the image data for later use
+            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:
-                        content_blocks.append({"text": f"Question: {user_input}"})
-
-                        response = client.models.generate_content(
-                            model="gemini-2.0-flash",
-                            contents=content_blocks
-                        )
-
-                        st.session_state.conversation.append((user_input, response.text))
-                        st.success("💡 Answer:")
-                        st.write(response.text)
-
-                        # Save to session history
-                        st.session_state.chat_history.append({
-                            "question": user_input,
-                            "answer": response.text
-                        })
-
-                except Exception as e:
-                    st.error(f"Error generating answer: {e}")
-
-        else:
-            st.warning("No valid documents or images processed.")
-
-    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"
-            )
+                        with st.spinner("Gemini is thinking..."):
+                            content_blocks = []
+
+                            # Add prior Q&A as part of the context
+                            if st.session_state.conversation:
+                                history_text = "\n\n".join(
+                                    f"Q: {entry[0]}\nA: {entry[1]}"
+                                    for entry in st.session_state.conversation
+                                )
+                                content_blocks.append({"text": f"Previous conversation:\n{history_text}"})
+
+                            # Add file text content
+                            if st.session_state.document_text.strip():
+                                content_blocks.append({"text": f"Context:\n{st.session_state.document_text}"})
+
+                            # Add image blocks (if image data exists)
+                            if "image_data" in st.session_state:
+                                content_blocks.append(st.session_state.image_data)
+
+                            # Add current question
+                            content_blocks.append({"text": f"Question: {user_input}"})
+
+                            response = client.models.generate_content(
+                                model="gemini-2.0-flash",
+                                contents=content_blocks
+                            )
+
+                            st.session_state.conversation.append((user_input, response.text))
+                            st.success("💡 Answer:")
+                            st.write(response.text)
+
+                            # Save to session history
+                            st.session_state.chat_history.append({
+                                "question": user_input,
+                                "answer": response.text
+                            })
+
+        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()
+    main()