Update app.py

app.py

@@ -190,13 +190,13 @@ def generate_download_button(file, label, mime_type):
 # --- Streamlit UI ---
 st.set_page_config("Deep Research Assistant", layout="centered")
 
-# Initialize session state
+# 🧠 Initialize session state
 if "last_report" not in st.session_state:
     st.session_state["last_report"] = ""
 if "follow_up_input" not in st.session_state:
     st.session_state["follow_up_input"] = ""
-if "methodology_recommended" not in st.session_state:
-    st.session_state["methodology_recommended"] = False
+if "methodology_notes" not in st.session_state:
+    st.session_state["methodology_notes"] = ""
 if "chat_history" not in st.session_state:
     st.session_state["chat_history"] = []
 
@@ -296,7 +296,32 @@ APA Citations:
     except Exception as e:
         st.error(f"❌ Error occurred: {e}")
 
-# --- Follow-up Q&A (always visible) ---
+# --- Methodology Recommender ---
+st.divider()
+st.subheader("🧪 Methodology Recommender")
+
+if st.button("🧠 Suggest Research Methodologies"):
+    if st.session_state["last_report"]:
+        try:
+            method_prompt = [
+                {"role": "system", "content": "You are a research advisor. Based on the following research report, suggest suitable research methodologies (quantitative, qualitative, ML/AI techniques, etc.). Be concise and give bullet-point suggestions."},
+                {"role": "user", "content": st.session_state["last_report"]}
+            ]
+            method_output = ""
+            method_box = st.empty()
+            for chunk in call_llm(method_prompt):
+                method_output += chunk
+                method_box.markdown(method_output, unsafe_allow_html=True)
+
+            # ✅ Store methodology context for follow-up
+            st.session_state["methodology_notes"] = method_output
+
+        except Exception as e:
+            st.error(f"❌ Methodology suggestion failed: {e}")
+    else:
+        st.warning("⚠️ Generate the research report first.")
+
+# --- Follow-up Q&A (Contextual to Report + Methodology) ---
 st.divider()
 st.subheader("💬 Follow-up Q&A")
 
@@ -305,40 +330,33 @@ followup = st.text_input("Ask a follow-up question:", key="follow_up_input")
 if st.button("Ask"):
     if followup:
         try:
-            chat = st.session_state.chat_history + [{"role": "user", "content": followup}]
+            context_intro = (
+                "Below is a research report followed by methodology suggestions.\n"
+                "Use both to answer the user's follow-up question."
+            )
+            combined_context = f"{context_intro}\n\n=== Report ===\n{st.session_state['last_report']}\n\n=== Methodology ===\n{st.session_state['methodology_notes']}"
+
+            chat = st.session_state.chat_history + [
+                {"role": "system", "content": "You are an academic research assistant."},
+                {"role": "user", "content": combined_context},
+                {"role": "user", "content": followup}
+            ]
+
             response = ""
             for chunk in call_llm(chat, max_tokens=1500):
                 response += chunk
+
             st.session_state.chat_history.append({"role": "user", "content": followup})
             st.session_state.chat_history.append({"role": "assistant", "content": response})
+
             st.markdown(response)
+
         except Exception as e:
             st.error(f"Follow-up error: {e}")
 
-# --- Chat history display (optional) ---
+# --- Optional: View Chat History ---
 with st.expander("📜 View Full Chat History", expanded=False):
     for msg in st.session_state.chat_history:
         role = msg["role"]
         prefix = "👤 You" if role == "user" else "🤖 Assistant"
-        st.markdown(f"**{prefix}:** {msg['content']}")
-
-# --- Methodology Recommender ---
-st.divider()
-st.subheader("🧪 Methodology Recommender")
-
-if st.button("🧠 Suggest Research Methodologies"):
-    if st.session_state["last_report"]:
-        try:
-            method_prompt = [
-                {"role": "system", "content": "You are a research advisor. Based on the following research report, suggest suitable research methodologies (quantitative, qualitative, ML/AI techniques, etc.). Be concise and give bullet-point suggestions."},
-                {"role": "user", "content": st.session_state["last_report"]}
-            ]
-            method_output = ""
-            method_box = st.empty()
-            for chunk in call_llm(method_prompt):
-                method_output += chunk
-                method_box.markdown(method_output, unsafe_allow_html=True)
-        except Exception as e:
-            st.error(f"❌ Methodology suggestion failed: {e}")
-    else:
-        st.warning("⚠️ Generate the research report first.")
+        st.markdown(f"**{prefix}:** {msg['content']}")