Update app.py

app.py

@@ -299,29 +299,23 @@ Citations:
 st.divider()
 st.subheader("📂 Your Research Threads")
 
-user_avatar = "https://cdn-icons-png.flaticon.com/512/9131/9131529.png"
-assistant_avatar = "https://cdn-icons-png.flaticon.com/512/4712/4712107.png"
-
 for tid, chats in st.session_state.chat_threads.items():
-    with st.container():
-        with st.expander(f"🧵 Thread {tid[:8]}", expanded=False):
-            for msg in chats:
-                avatar = user_avatar if msg['role'] == 'user' else assistant_avatar
-                bubble_color = "#DCF8C6" if msg['role'] == 'user' else "#F0F0F0"
-                align = "flex-end" if msg['role'] == 'user' else "flex-start"
-                st.markdown({msg['content']})
-            followup = st.text_input(f"💬 Continue Thread {tid[:8]}:", key=f"followup_{tid}")
-            if st.button(f"Ask Follow-up {tid}", key=f"button_{tid}"):
-                if followup:
-                    with st.spinner("🤖 Assistant is typing..."):
-                        response = ""
-                        for chunk in call_llm(st.session_state.chat_threads[tid] + [{"role": "user", "content": followup}], max_tokens=2000):
-                            response += chunk
-                            st.markdown(response)
-                    st.session_state.chat_threads[tid].append({"role": "user", "content": followup})
-                    st.session_state.chat_threads[tid].append({"role": "assistant", "content": response})
-                    save_session_data()
-                    st.rerun()
+    with st.expander(f"🧵 Thread {tid[:8]}", expanded=False):
+        for msg in chats:
+            with st.chat_message(msg["role"] if msg["role"] in ["user", "assistant"] else "assistant"):
+                st.markdown(msg["content"])
+
+        followup = st.text_input(f"💬 Continue Thread {tid[:8]}:", key=f"followup_{tid}")
+        if st.button(f"Ask Follow-up {tid}", key=f"button_{tid}"):
+            if followup:
+                with st.spinner("🤖 Assistant is typing..."):
+                    response = ""
+                    for chunk in call_llm(st.session_state.chat_threads[tid] + [{"role": "user", "content": followup}], max_tokens=2000):
+                        response += chunk
+                st.session_state.chat_threads[tid].append({"role": "user", "content": followup})
+                st.session_state.chat_threads[tid].append({"role": "assistant", "content": response})
+                save_session_data()
+                st.rerun()
 
 # --- Download All Threads Section ---
 if st.session_state.chat_threads:
@@ -335,7 +329,7 @@ st.divider()
 st.subheader("🧪 Methodology Recommender")
 
 if st.button("🧠 Suggest Research Methodologies"):
-    if st.session_state["last_report"]:
+    if st.session_state.get("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."},
@@ -352,7 +346,7 @@ if st.button("🧠 Suggest Research Methodologies"):
         except Exception as e:
             st.error(f"❌ Methodology suggestion failed: {e}")
     else:
-        st.warning("⚠️ Generate the research report first.")
+        st.warning("⚠️ No research report available. Please generate research first.")
 
 # --- Follow-up Q&A (Contextual to Report + Methodology) ---
 st.divider()
@@ -360,7 +354,7 @@ st.subheader("💬 Follow-up Q&A")
 
 followup = st.text_input("Ask a follow-up question:", key="follow_up_input")
 
-if st.button("Ask"):
+if st.button("Ask Follow-up"):
     if followup:
         try:
             context_intro = (
@@ -382,7 +376,8 @@ if st.button("Ask"):
             st.session_state.chat_history.append({"role": "user", "content": followup})
             st.session_state.chat_history.append({"role": "assistant", "content": response})
 
-            st.markdown(response)
+            with st.chat_message("assistant"):
+                st.markdown(response)
 
         except Exception as e:
             st.error(f"Follow-up error: {e}")
@@ -393,7 +388,7 @@ st.subheader("📤 Upload Your Paper for Feedback")
 
 uploaded_file = st.file_uploader("Upload your research paper (.pdf or .txt)", type=["pdf", "txt"])
 
-if uploaded_file and st.button("🧠 Analyze and Suggest Improvements"):
+if uploaded_file and st.button("🧠 Analyze Paper for Improvements"):
     try:
         def extract_text_from_file(file):
             if file.name.endswith(".pdf"):
@@ -421,21 +416,21 @@ Be honest and constructive. Here's the full text:
 \"\"\"{paper_text}\"\"\""""}
             ]
 
-            st.status("🔎 Analyzing your paper...")
-            improvement_output = ""
-            feedback_box = st.empty()
-            for chunk in call_llm(feedback_prompt, max_tokens=2500):
-                improvement_output += chunk
-                feedback_box.markdown(improvement_output, unsafe_allow_html=True)
+            with st.status("🔎 Analyzing your paper..."):
+                improvement_output = ""
+                feedback_box = st.empty()
+                for chunk in call_llm(feedback_prompt, max_tokens=2500):
+                    improvement_output += chunk
+                    feedback_box.markdown(improvement_output, unsafe_allow_html=True)
 
     except Exception as e:
         st.error(f"❌ Error while analyzing paper: {e}")
 
-
 # --- Full Chat History Viewer ---
 st.divider()
-with st.expander("📜 View Full Chat History", expanded=False):
+st.subheader("📜 Full Chat History")
+
+with st.expander("View 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']}")
+        with st.chat_message(msg["role"] if msg["role"] in ["user", "assistant"] else "assistant"):
+            st.markdown(msg["content"])