Update app.py

app.py

@@ -4,15 +4,12 @@ import requests
 import datetime
 import time
 import json
+import uuid
 from dotenv import load_dotenv
 from tavily import TavilyClient
 import feedparser
 from fuzzywuzzy import fuzz
-from PIL import Image
-from io import BytesIO
 from fpdf import FPDF
-import base64
-import uuid
 from duckduckgo_search import DDGS
 
 # --- Load API Keys ---
@@ -21,7 +18,21 @@ OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 tavily = TavilyClient(api_key=TAVILY_API_KEY)
 
-# --- Helper Functions ---
+# --- Utility Functions ---
+def save_session_data():
+    data = {
+        "memory_bank": st.session_state.get("memory_bank", []),
+        "chat_threads": st.session_state.get("chat_threads", {})
+    }
+    with open("session_memory.json", "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=4)
+
+def load_session_data():
+    if os.path.exists("session_memory.json"):
+        with open("session_memory.json", "r", encoding="utf-8") as f:
+            data = json.load(f)
+            st.session_state.memory_bank = data.get("memory_bank", [])
+            st.session_state.chat_threads = data.get("chat_threads", {})
 
 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"
@@ -53,7 +64,7 @@ def call_llm(messages, model="deepseek/deepseek-chat-v3-0324:free", max_tokens=4
                         except json.JSONDecodeError:
                             pass
 
-def get_image_urls(query, max_images=5):
+def get_image_urls(query, max_images=3):
     with DDGS() as ddgs:
         return [img["image"] for img in ddgs.images(query, max_results=max_images)]
 
@@ -141,34 +152,53 @@ def build_chronological_progression(sources):
         summary += f"**{year}**\n{entries}\n\n"
     return summary.strip()
 
-# --- Initialize Streamlit Session ---
-st.set_page_config(page_title="🧠 Deep Research Assistant 2.0", layout="centered")
-
+def download_threads_as_pdf(chat_threads):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.set_font("Arial", size=12)
+    for tid, chats in chat_threads.items():
+        pdf.cell(0, 10, f"Thread {tid[:8]}", ln=True)
+        for msg in chats:
+            role = "You" if msg["role"] == "user" else "Assistant"
+            text = f"{role}: {msg['content']}"
+            pdf.multi_cell(0, 10, text)
+        pdf.ln(5)
+    pdf_output = BytesIO()
+    pdf_bytes = pdf.output(dest='S').encode('latin-1')
+    pdf_output.write(pdf_bytes)
+    pdf_output.seek(0)
+    return pdf_output
+
+# --- Streamlit UI ---
+st.set_page_config(page_title="🧠 Deep Research Assistant 3.0", layout="centered")
+
+# --- Load Memory ---
 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
 
-# --- Sidebar Inputs ---
+load_session_data()
+
+# --- Sidebar ---
 with st.sidebar:
-    st.title("Deep Research Assistant 2.0 🚀")
-    topic = st.text_input("🔍 Enter your research topic")
+    st.title("Deep Research Assistant 3.0 🚀")
+    topic = st.text_input("🔍 Enter research topic")
     report_type = st.selectbox("📄 Report Type", ["Summary", "Detailed Report", "Thorough Academic Research"])
     tone = st.selectbox("🎯 Tone", ["Objective", "Persuasive", "Narrative"])
     source_type = st.selectbox("📚 Sources", ["Web Only", "Academic Only", "Hybrid"])
     custom_domains = st.text_input("🌐 Optional Domains", placeholder="example.com, nature.com")
     research_button = st.button("🚀 Start Research")
 
-# --- Main Logic ---
+# --- Main Area ---
 st.title("📘 Research Output")
 
 if research_button and topic:
     try:
-        with st.status("🔎 Gathering sources..."):
+        with st.status("🔍 Gathering sources..."):
             all_sources = []
             if source_type in ["Web Only", "Hybrid"]:
                 all_sources += get_sources(topic, custom_domains) if custom_domains.strip() else get_sources(topic)
@@ -183,8 +213,7 @@ if research_button and topic:
             merged = sort_sources_chronologically(merged)
             chronological_progress = build_chronological_progression(merged)
 
-        # Prepare previous learnings
-        previous_learnings = "\n\n".join(st.session_state.memory_bank[-5:])  # last 5 learnings
+        previous_learnings = "\n\n".join(st.session_state.memory_bank[-5:])
 
         citations = [f"- {s['title']} ({s['year']}) [{s['source']}]({s['url']})" for s in merged]
         sources_text = "\n".join([
@@ -198,35 +227,31 @@ if research_button and topic:
             "Thorough Academic Research": "Craft a full academic paper >1000 words."
         }[report_type]
 
-        # Create Thread ID
         thread_id = str(uuid.uuid4())
         st.session_state.current_thread_id = thread_id
         st.session_state.chat_threads[thread_id] = []
 
-        # --- LLM Prompt ---
         prompt = f"""
-You are an expert research assistant.
+Use past learnings:
 
-🔵 Past Knowledge:
 {previous_learnings}
 
-🔵 New Research Topic:
+New Topic:
 {topic}
 
-🔵 Writing Style:
+Writing:
 {tone} tone, {length_instruction}
 
-🔵 Research Timeline:
+Timeline:
 {chronological_progress}
 
-🔵 Sources:
+Sources:
 {sources_text}
 
-🔵 Citations:
+Citations:
 {chr(10).join(citations)}
 """
 
-        # --- Generate Report ---
         st.subheader(f"📝 {report_type} on '{topic}'")
         output_placeholder = st.empty()
         final_output = ""
@@ -237,10 +262,12 @@ You are an expert research assistant.
         st.session_state.memory_bank.append(final_output)
         st.session_state.chat_threads[thread_id].append({"role": "assistant", "content": final_output})
 
+        save_session_data()
+
     except Exception as e:
         st.error(f"❌ Error: {e}")
 
-# --- Show Chat Threads ---
+# --- Chat Threads and Follow-ups ---
 st.divider()
 st.subheader("📂 Your Research Threads")
 
@@ -250,6 +277,22 @@ for tid, chats in st.session_state.chat_threads.items():
             role = "🧑 You" if msg['role'] == 'user' else "🤖 Assistant"
             st.markdown(f"**{role}:** {msg['content']}")
 
+        followup = st.text_input(f"💬 Ask more in Thread {tid[:8]}:", key=f"followup_{tid}")
+        if st.button(f"Ask Follow-up {tid}", key=f"button_{tid}"):
+            if followup:
+                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.experimental_rerun()
+
+# --- Download Button ---
+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")
+
 
 
 # 🧠 Initialize session state