Update app.py

app.py

@@ -189,6 +189,13 @@ def generate_download_button(file, label, mime_type):
 
 # --- Streamlit UI ---
 st.set_page_config("Deep Research Assistant", layout="centered")
+# 🧠 Initialize session state
+if "last_report" not in st.session_state:
+    st.session_state["last_report"] = ""
+if "generate_mindmap" not in st.session_state:
+    st.session_state["generate_mindmap"] = False
+if "follow_up" not in st.session_state:
+    st.session_state["follow_up"] = ""
 
 # Mermaid.js for mind map rendering
 st.markdown("""
@@ -300,8 +307,11 @@ APA Citations:
 
         # 🔹 Mind Map Generator
         st.subheader("🧭 Visual Mind Map")
-        if st.button("🗺 Generate Mind Map", key="mindmap_btn"):
-            if "last_report" in st.session_state:
+        if st.button("🗺 Generate Mind Map"):
+            st.session_state["generate_mindmap"] = True
+        
+        if st.session_state.get("generate_mindmap") and st.session_state["last_report"]:
+            try:
                 mindmap_prompt = [
                     {"role": "system", "content": "Convert this research report into a mermaid.js mind map."},
                     {"role": "user", "content": st.session_state["last_report"]}
@@ -311,26 +321,30 @@ APA Citations:
                     mindmap_code += chunk
                 mindmap_code = mindmap_code.replace("```mermaid", "").replace("```", "").strip()
                 st.markdown(f"<div class='mermaid'>{mindmap_code}</div>", unsafe_allow_html=True)
-            else:
-                st.error("❌ No report found. Please run research first.")
+            except Exception as e:
+                st.error(f"Mind map generation failed: {e}")
+            finally:
+                st.session_state["generate_mindmap"] = False
 
         # 🔹 Follow-up Interaction
-        st.subheader("💬 Ask a Follow-Up")
-        follow_input = st.text_input("Ask a question about the report:")
         if st.button("🔄 Submit Follow-Up"):
-            if "last_report" in st.session_state:
+            st.session_state["follow_up"] = follow_input
+        
+        if st.session_state.get("follow_up") and st.session_state["last_report"]:
+            try:
                 follow_prompt = [
                     {"role": "system", "content": "You are a helpful academic assistant."},
                     {"role": "user", "content": st.session_state["last_report"]},
-                    {"role": "user", "content": follow_input}
+                    {"role": "user", "content": st.session_state["follow_up"]}
                 ]
                 follow_output = ""
                 follow_box = st.empty()
                 for chunk in call_llm(follow_prompt):
                     follow_output += chunk
                     follow_box.markdown(follow_output, unsafe_allow_html=True)
-            else:
-                st.warning("⚠️ Please generate the report first.")
-
+            except Exception as e:
+                st.error(f"Follow-up failed: {e}")
+            finally:
+                st.session_state["follow_up"] = ""
     except Exception as e:
-        st.error(f"❌ Error occurred: {e}")
+        st.error(f"❌ Error occurred: {e}")