Update app.py

app.py

@@ -191,11 +191,14 @@ def generate_download_button(file, label, mime_type):
 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"] = ""
+    st.session_state.last_report = ""
+
+if "mindmap_triggered" not in st.session_state:
+    st.session_state.mindmap_triggered = False
+
+if "followup_question" not in st.session_state:
+    st.session_state.followup_question = ""
+
 
 # Mermaid.js for mind map rendering
 st.markdown("""
@@ -288,7 +291,8 @@ APA Citations:
             output_placeholder.markdown(full_output, unsafe_allow_html=True)
 
         # Save output for follow-up
-        st.session_state["last_report"] = full_output
+        st.session_state.last_report = full_output
+
 
         # 🔹 Downloads
         st.subheader("📄 Downloads")
@@ -305,46 +309,50 @@ APA Citations:
         else:
             st.success("✅ No major overlaps found.")
 
-        # 🔹 Mind Map Generator
-        st.subheader("🧭 Visual Mind Map")
-        if st.button("🗺 Generate Mind Map"):
-            st.session_state["generate_mindmap"] = True
+        st.subheader("🧭 Visual Mind Map")        
+        mindmap_button = st.button("🗺 Generate Mind Map")
+        if mindmap_button:
+            st.session_state.mindmap_triggered = 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"]}
-                ]
-                mindmap_code = ""
-                for chunk in call_llm(mindmap_prompt):
-                    mindmap_code += chunk
-                mindmap_code = mindmap_code.replace("```mermaid", "").replace("```", "").strip()
-                st.markdown(f"<div class='mermaid'>{mindmap_code}</div>", unsafe_allow_html=True)
-            except Exception as e:
-                st.error(f"Mind map generation failed: {e}")
-            finally:
-                st.session_state["generate_mindmap"] = False
+        if st.session_state.mindmap_triggered and st.session_state.last_report:
+            mindmap_prompt = [
+                {"role": "system", "content": "Convert this research report into a mermaid.js mind map."},
+                {"role": "user", "content": st.session_state.last_report}
+            ]
+            mindmap_code = ""
+            for chunk in call_llm(mindmap_prompt):
+                mindmap_code += chunk
+        
+            mindmap_code = mindmap_code.replace("```mermaid", "").replace("```", "").strip()
+            st.markdown(f"<div class='mermaid'>{mindmap_code}</div>", unsafe_allow_html=True)
+        
+            # Reset trigger after rendering
+            st.session_state.mindmap_triggered = False
+
 
         # 🔹 Follow-up Interaction
-        if st.button("🔄 Submit Follow-Up"):
-            st.session_state["follow_up"] = follow_input
+        st.subheader("💬 Ask a Follow-Up")
         
-        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": 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)
-            except Exception as e:
-                st.error(f"Follow-up failed: {e}")
-            finally:
-                st.session_state["follow_up"] = ""
+        followup_input = st.text_input("Ask a question about the report:")
+        
+        followup_submit = st.button("🔄 Submit Follow-Up")
+        if followup_submit and followup_input:
+            st.session_state.followup_question = followup_input
+        
+        if st.session_state.followup_question and st.session_state.last_report:
+            follow_prompt = [
+                {"role": "system", "content": "You are a helpful academic assistant."},
+                {"role": "user", "content": st.session_state.last_report},
+                {"role": "user", "content": st.session_state.followup_question}
+            ]
+            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)
+        
+            # Reset after showing
+            st.session_state.followup_question = ""
+
     except Exception as e:
         st.error(f"❌ Error occurred: {e}")