Update app.py

app.py

@@ -189,28 +189,18 @@ 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 "chat_history" not in st.session_state:
+    st.session_state["chat_history"] = []
 
-
-# Mermaid.js for mind map rendering
-st.markdown("""
-<script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>
-<script>
-    mermaid.initialize({ startOnLoad: true });
-</script>
-""", unsafe_allow_html=True)
-
-st.markdown("""
-<style>
-    .stApp { background-color: #0f172a; color: white; }
-    h1, h2, h3 { color: #facc15; }
-</style>
-""", unsafe_allow_html=True)
-
+# --- Sidebar Inputs ---
 with st.sidebar:
     st.title("🧠 Deep Research Assistant")
     topic = st.text_input("🔍 Enter your research topic")
@@ -222,6 +212,7 @@ with st.sidebar:
 
 st.title("📘 Research Output")
 
+# --- Main Report Generation ---
 if research_button and topic:
     try:
         with st.status("🔍 Gathering and analyzing sources...") as status:
@@ -284,15 +275,15 @@ APA Citations:
             full_output += chunk
             output_placeholder.markdown(full_output, unsafe_allow_html=True)
 
-        # Save output for follow-up
+        # Save output for follow-up and methodology
         st.session_state["last_report"] = full_output
 
-        # 🔹 Downloads
+        # --- Downloads ---
         st.subheader("📄 Downloads")
         st.markdown(generate_download_button(generate_pdf(full_output), "Research_Report.pdf", "application/pdf"), unsafe_allow_html=True)
         st.markdown(generate_download_button(generate_latex(full_output), "Research_Report.tex", "application/x-latex"), unsafe_allow_html=True)
 
-        # 🔹 Plagiarism Check
+        # --- Plagiarism Check ---
         st.subheader("🔎 Plagiarism Check")
         overlaps = check_plagiarism(full_output, topic)
         if overlaps:
@@ -302,50 +293,52 @@ APA Citations:
         else:
             st.success("✅ No major overlaps found.")
 
-        st.subheader("💬 Ask a Follow-Up")
-
-        st.session_state["follow_up_input"] = st.text_input(
-            "Ask a question about the report:", value=st.session_state["follow_up_input"]
-        )
-        
-        if st.button("🔄 Submit Follow-Up") and st.session_state["follow_up_input"]:
-            if 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_input"]}
-                    ]
-                    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}")
-            else:
-                st.warning("⚠️ Generate the research report first.")
-
-        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.")
+    except Exception as e:
+        st.error(f"❌ Error occurred: {e}")
 
+# --- Follow-up Q&A (always visible) ---
+st.divider()
+st.subheader("💬 Follow-up Q&A")
 
+followup = st.text_input("Ask a follow-up question:", key="follow_up_input")
 
-    except Exception as e:
-        st.error(f"❌ Error occurred: {e}")
+if st.button("Ask"):
+    if followup:
+        try:
+            chat = st.session_state.chat_history + [{"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) ---
+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.")