Update app.py

app.py

@@ -189,11 +189,23 @@ def generate_download_button(file, label, mime_type):
 
 # --- Streamlit UI ---
 st.set_page_config("Deep Research Bot", layout="centered")
+
+# Inject 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; }
-    .css-1d391kg, .css-1kyxreq, .css-q8sbsg { background-color: #1e293b; color: white; border-radius: 10px; padding: 10px; }
+    .css-1d391kg, .css-1kyxreq, .css-q8sbsg {
+        background-color: #1e293b; color: white;
+        border-radius: 10px; padding: 10px;
+    }
     </style>
 """, unsafe_allow_html=True)
 
@@ -235,23 +247,25 @@ if research_button and topic:
 
             merged = merge_duplicates(all_sources)
 
-            st.markdown("---")
-            st.subheader("🖼 Source Previews")
-            image_shown = False
-            cols = st.columns(2)
-            for i, m in enumerate(merged):
-                if m.get("image_url"):
-                    with cols[i % 2]:
-                        st.image(m["image_url"], caption=m["title"], use_container_width=True)
-                        image_shown = True
-            if not image_shown:
-                st.info("ℹ️ No image previews available from sources.")
-
-            citations = [generate_apa_citation(m['title'], m['url'], m['source']) for m in merged]
-            combined_text = "\n\n".join(
-                [f"- [{m['title']}]({m['url']})\n> {m.get('snippet', m.get('summary', ''))[:300]}..." for m in merged]
-            )
-            status.update(label="🧠 Synthesizing final report...")
+        # Image previews
+        st.markdown("---")
+        st.subheader("🖼 Source Previews")
+        image_shown = False
+        cols = st.columns(2)
+        for i, m in enumerate(merged):
+            if m.get("image_url"):
+                with cols[i % 2]:
+                    st.image(m["image_url"], caption=m["title"], use_container_width=True)
+                    image_shown = True
+        if not image_shown:
+            st.info("ℹ️ No image previews available from sources.")
+
+        citations = [generate_apa_citation(m['title'], m['url'], m['source']) for m in merged]
+        combined_text = "\n\n".join(
+            [f"- [{m['title']}]({m['url']})\n> {m.get('snippet', m.get('summary', ''))[:300]}..." for m in merged]
+        )
+
+        st.status("🧠 Synthesizing final report...")
 
         prompt = f"""
 You are an expert research assistant.
@@ -260,13 +274,13 @@ You are an expert research assistant.
 3. Write a full academic-style report on the newly proposed topic using a {tone.lower()} tone.
 
 Report Format:
-- **Introduction**: Context and importance of the new topic.
-- **Identified Research Gap**: Summary of missing or underexplored areas.
-- **Proposed Novel Topic**: Explanation of the new idea and its significance.
-- **Methodology**: Outline methods or frameworks to explore the idea.
-- **Comparative Insight**: How the new idea differs/improves over previous works.
-- **Future Work**: Possible expansions or questions.
-- **Citations**: Include APA style citations at the end.
+- Introduction
+- Identified Research Gap
+- Proposed Novel Topic
+- Methodology
+- Comparative Insight
+- Future Work
+- Citations (APA)
 
 Sources:
 {combined_text}
@@ -275,6 +289,7 @@ APA Citations:
 {chr(10).join(citations)}
         """
 
+        # Generate report
         st.subheader(f"📝 {report_type} on '{topic}'")
         output_placeholder = st.empty()
         full_output = ""
@@ -282,49 +297,60 @@ APA Citations:
             full_output += chunk
             output_placeholder.markdown(full_output, unsafe_allow_html=True)
 
+        # Store report for mind map and follow-up
+        st.session_state["last_report"] = full_output
+
+        # Downloads
         st.markdown("---")
         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)
 
-        overlaps = check_plagiarism(full_output, topic)
+        # Plagiarism
         st.markdown("---")
         st.subheader("🔎 Plagiarism Check")
+        overlaps = check_plagiarism(full_output, topic)
         if overlaps:
-            st.warning("Potential overlaps detected:")
+            st.warning("⚠️ Overlaps found:")
             for hit in overlaps:
                 st.markdown(f"- [{hit['title']}]({hit['url']})")
         else:
-            st.success("✅ No major overlaps found.")
-        # Mind Map Generator (Visual, not code block)
+            st.success("✅ No significant overlaps found.")
+
+        # Mind Map Generator (Visual)
         st.markdown("---")
         st.subheader("🧭 Visual Mind Map")
-        
-        # Inject Mermaid.js script to enable 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)
-        
         if st.button("🗺 Generate Mind Map"):
             mindmap_prompt = [
-                {"role": "system", "content": "You are a researcher that summarizes a topic into a mind map using Mermaid.js."},
-                {"role": "user", "content": f"Generate a mermaid mind map for the following research paper:\n\n{st.session_state['last_report']}"}
+                {"role": "system", "content": "You are a researcher who converts reports into Mermaid.js mind maps."},
+                {"role": "user", "content": f"Create a mermaid mind map for this research:\n{full_output}"}
             ]
-        
             mindmap_code = ""
             for chunk in call_llm(mindmap_prompt):
                 mindmap_code += chunk
-        
-            # Clean any code block wrapping
+
             mindmap_code = mindmap_code.replace("```mermaid", "").replace("```", "").strip()
-        
-            st.markdown(f"""
-            <div class="mermaid">
-            {mindmap_code}
-            </div>
-            """, unsafe_allow_html=True)
+            st.markdown(f"<div class='mermaid'>{mindmap_code}</div>", unsafe_allow_html=True)
+
+        # Follow-up Q&A
+        st.markdown("---")
+        st.subheader("💬 Ask a Follow-Up Question")
+        follow_input = st.text_input("Ask a question related to the report above:")
+        if st.button("🔄 Submit Follow-Up"):
+            if "last_report" in st.session_state:
+                with st.spinner("Thinking..."):
+                    follow_prompt = [
+                        {"role": "system", "content": "You are a helpful research assistant."},
+                        {"role": "user", "content": f"Here is the report:\n{st.session_state['last_report']}"},
+                        {"role": "user", "content": follow_input}
+                    ]
+                    follow_output = ""
+                    follow_placeholder = st.empty()
+                    for chunk in call_llm(follow_prompt):
+                        follow_output += chunk
+                        follow_placeholder.markdown(follow_output, unsafe_allow_html=True)
+            else:
+                st.error("⚠️ Report not found. Please generate a report first.")
+
     except Exception as e:
         st.error(f"❌ Error occurred: {e}")