Update app.py

app.py

@@ -295,6 +295,27 @@ Citations:
 
     except Exception as e:
         st.error(f"❌ Error: {e}")
+# --- Build Full Context (Research + Thread + Methodology) ---
+def build_full_context():
+    full_context = ""
+
+    # Add Research Report
+    if st.session_state.get("last_report"):
+        full_context += f"=== Research Report ===\n{st.session_state['last_report']}\n\n"
+
+    # Add Thread Messages
+    if st.session_state.get("current_thread_id"):
+        thread_msgs = st.session_state.chat_threads.get(st.session_state.current_thread_id, [])
+        for msg in thread_msgs:
+            who = "User" if msg["role"] == "user" else "Assistant"
+            full_context += f"{who}: {msg['content']}\n\n"
+
+    # Add Methodology if available
+    if st.session_state.get("methodology_notes"):
+        full_context += f"=== Methodology Suggestions ===\n{st.session_state['methodology_notes']}\n\n"
+
+    return full_context
+
 # --- Chat Threads Section ---
 st.divider()
 st.subheader("📂 Your Research Threads")
@@ -329,11 +350,14 @@ st.divider()
 st.subheader("🧪 Methodology Recommender")
 
 if st.button("🧠 Suggest Research Methodologies"):
-    if st.session_state.get("last_report"):
+    context = build_full_context()
+    if context:
         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"]}
+                {"role": "system", "content": "You are a research advisor."},
+                {"role": "user", "content": f"""Given the following conversation, research report, and context, suggest a very detailed and customized research methodology that matches the research objectives discussed.
+
+\"\"\"{context}\"\"\""""}
             ]
             method_output = ""
             method_box = st.empty()
@@ -346,41 +370,41 @@ if st.button("🧠 Suggest Research Methodologies"):
         except Exception as e:
             st.error(f"❌ Methodology suggestion failed: {e}")
     else:
-        st.warning("⚠️ No research report available. Please generate research first.")
+        st.warning("⚠️ No research context available. Please generate research first.")
 
-# --- Follow-up Q&A (Contextual to Report + Methodology) ---
+# --- Follow-up Q&A (Contextual to Full Thread) ---
 st.divider()
 st.subheader("💬 Follow-up Q&A")
 
 followup = st.text_input("Ask a follow-up question:", key="follow_up_input")
 
 if st.button("Ask Follow-up"):
-    if followup:
+    context = build_full_context()
+    if followup and context:
         try:
-            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}
+            combined_prompt = [
+                {"role": "system", "content": "You are an expert academic research assistant."},
+                {"role": "user", "content": f"""Use ONLY the following research report, conversation, and methodology suggestions to answer the follow-up question below. Stay fully topic-specific and context-aware.
+
+\"\"\"{context}\"\"\"
+
+Follow-up Question: {followup}
+"""}
             ]
 
             response = ""
-            for chunk in call_llm(chat, max_tokens=1500):
-                response += chunk
+            with st.chat_message("assistant"):
+                for chunk in call_llm(combined_prompt, max_tokens=2000):
+                    response += chunk
+                st.markdown(response)
 
             st.session_state.chat_history.append({"role": "user", "content": followup})
             st.session_state.chat_history.append({"role": "assistant", "content": response})
 
-            with st.chat_message("assistant"):
-                st.markdown(response)
-
         except Exception as e:
-            st.error(f"Follow-up error: {e}")
+            st.error(f"❌ Follow-up error: {e}")
+    else:
+        st.warning("⚠️ No sufficient context available. Please generate research first.")
 
 # --- Paper Upload for Review & Improvement ---
 st.divider()
@@ -433,4 +457,4 @@ st.subheader("📜 Full Chat History")
 with st.expander("View Chat History", expanded=False):
     for msg in st.session_state.chat_history:
         with st.chat_message(msg["role"] if msg["role"] in ["user", "assistant"] else "assistant"):
-            st.markdown(msg["content"])
+            st.markdown(msg["content"])