Update app.py

app.py

@@ -13,17 +13,31 @@ from fpdf import FPDF
 from duckduckgo_search import DDGS
 from io import BytesIO
 
-# --- Load API Keys ---
+# Load environment variables
 load_dotenv()
 OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 tavily = TavilyClient(api_key=TAVILY_API_KEY)
 
-# --- Utility Functions ---
+# Initialize session state
+if "memory_bank" not in st.session_state:
+    st.session_state.memory_bank = []
+if "chat_threads" not in st.session_state:
+    st.session_state.chat_threads = {}
+if "current_thread_id" not in st.session_state:
+    st.session_state.current_thread_id = None
+if "last_report" not in st.session_state:
+    st.session_state.last_report = ""
+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 = []
+
+# Session data functions
 def save_session_data():
     data = {
-        "memory_bank": st.session_state.get("memory_bank", []),
-        "chat_threads": st.session_state.get("chat_threads", {})
+        "memory_bank": st.session_state.memory_bank,
+        "chat_threads": st.session_state.chat_threads
     }
     with open("session_memory.json", "w", encoding="utf-8") as f:
         json.dump(data, f, ensure_ascii=False, indent=4)
@@ -35,12 +49,14 @@ def load_session_data():
             st.session_state.memory_bank = data.get("memory_bank", [])
             st.session_state.chat_threads = data.get("chat_threads", {})
 
+load_session_data()
+
+# LLM call
 def call_llm(messages, model="deepseek/deepseek-chat-v3-0324:free", max_tokens=4000, temperature=0.7):
     url = "https://openrouter.ai/api/v1/chat/completions"
     headers = {
         "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "Content-Type": "application/json",
-        "X-Title": "Deep Research Assistant"
+        "Content-Type": "application/json"
     }
     data = {
         "model": model,
@@ -64,8 +80,8 @@ def call_llm(messages, model="deepseek/deepseek-chat-v3-0324:free", max_tokens=4
                                 yield token
                         except json.JSONDecodeError:
                             pass
-
-def get_image_urls(query, max_images=5):
+# --- Source Gathering Functions ---
+def get_image_urls(query, max_images=6):
     with DDGS() as ddgs:
         return [img["image"] for img in ddgs.images(query, max_results=max_images)]
 
@@ -175,21 +191,11 @@ def download_threads_as_pdf(chat_threads):
     pdf_output.seek(0)
     return pdf_output
 
-# --- Streamlit UI Setup ---
+# --- Streamlit UI Start ---
 st.set_page_config(page_title="🧠 Deep Research Assistant 4.0", layout="centered")
 
-if "memory_bank" not in st.session_state:
-    st.session_state.memory_bank = []
-if "chat_threads" not in st.session_state:
-    st.session_state.chat_threads = {}
-if "current_thread_id" not in st.session_state:
-    st.session_state.current_thread_id = None
-
-load_session_data()
-
 # --- Sidebar ---
 with st.sidebar:
-    st.image("https://raw.githubusercontent.com/mk-gurucharan/streamlit-deep-research/main/deep_research_logo.png", use_container_width=True)
     st.markdown("## 🔍 Start New Research")
     topic = st.text_input("🧠 Topic")
     report_type = st.selectbox("📄 Report Type", ["Summary", "Detailed Report", "Thorough Academic Research"])
@@ -197,17 +203,12 @@ with st.sidebar:
     source_type = st.selectbox("📚 Sources", ["Web Only", "Academic Only", "Hybrid"])
     custom_domains = st.text_input("🌐 Optional Domains", placeholder="forbes.com, mit.edu")
     research_button = st.button("🚀 Run Deep Research", use_container_width=True)
-    st.markdown("---")
-    st.markdown("Made with ❤️ by Cutie AI ✨")
 
-# --- Main Title ---
 st.title("🌙 Deep Research Assistant 4.0")
 st.markdown("Where serious research meets serious style. 🧠💖")
 st.divider()
 
-# Continue from here for Main Research Section...
-
-# --- Show Web Images from Topic ---
+# --- Web Images Section ---
 if topic and research_button:
     st.subheader("🖼 Related Images from the Web")
     try:
@@ -294,13 +295,12 @@ Citations:
 
     except Exception as e:
         st.error(f"❌ Error: {e}")
-
 # --- Chat Threads Section ---
 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"  
+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():
@@ -339,18 +339,6 @@ if st.session_state.chat_threads:
     pdf_file = download_threads_as_pdf(st.session_state.chat_threads)
     st.download_button("📥 Download All Threads as PDF", data=pdf_file, file_name="Research_Threads.pdf", mime="application/pdf", use_container_width=True)
 
-
-
-# 🧠 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_notes" not in st.session_state:
-    st.session_state["methodology_notes"] = ""
-if "chat_history" not in st.session_state:
-    st.session_state["chat_history"] = []
-
 # --- Methodology Recommender ---
 st.divider()
 st.subheader("🧪 Methodology Recommender")
@@ -368,7 +356,6 @@ if st.button("🧠 Suggest Research Methodologies"):
                 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:
@@ -433,15 +420,8 @@ if uploaded_file and st.button("🧠 Analyze and Suggest Improvements"):
         else:
             feedback_prompt = [
                 {"role": "system", "content": "You are an expert academic advisor."},
-                {"role": "user", "content": f"""I have written the following research paper. Please analyze it and provide detailed suggestions on:
-                                                - Areas where the paper is weak or unclear
-                                                - How to improve the novelty or originality
-                                                - Structural improvements or better ways to present arguments
-
-                                            Be honest and constructive. Here's the full text:
-
-                                            \"\"\"{paper_text}\"\"\""""}
-                             ]
+                {"role": "user", "content": f\"\"\"I have written the following research paper. Please analyze it and provide detailed suggestions on:\n- Areas where the paper is weak or unclear\n- How to improve the novelty or originality\n- Structural improvements or better ways to present arguments\n\nBe honest and constructive. Here's the full text:\n\n\"\"\"{paper_text}\"\"\"\"}
+            ]
 
             with st.status("🔎 Analyzing your paper..."):
                 improvement_output = ""
@@ -453,8 +433,8 @@ if uploaded_file and st.button("🧠 Analyze and Suggest Improvements"):
     except Exception as e:
         st.error(f"❌ Error while analyzing paper: {e}")
 
-
-# --- Optional: View Chat History ---
+# --- Full Chat History Viewer ---
+st.divider()
 with st.expander("📜 View Full Chat History", expanded=False):
     for msg in st.session_state.chat_history:
         role = msg["role"]