Update app.py

app.py

@@ -2,15 +2,9 @@
 # MedGenesis AI · CPU-only Streamlit app (OpenAI / Gemini)
 
 import os, pathlib
-
-# ── Streamlit telemetry dir fix ───────────────────────────────────────
-os.environ["STREAMLIT_DATA_DIR"]               = "/tmp/.streamlit"
-os.environ["XDG_STATE_HOME"]                  = "/tmp"
-os.environ["STREAMLIT_BROWSER_GATHERUSAGESTATS"] = "false"
-pathlib.Path("/tmp/.streamlit").mkdir(parents=True, exist_ok=True)
-
 import asyncio, re
 from pathlib import Path
+
 import streamlit as st
 import pandas as pd
 import plotly.express as px
@@ -18,10 +12,16 @@ from fpdf import FPDF
 from streamlit_agraph import agraph
 
 from mcp.orchestrator import orchestrate_search, answer_ai_question
-from mcp.workspace      import get_workspace, save_query
+from mcp.workspace import get_workspace, save_query
 from mcp.knowledge_graph import build_agraph
 from mcp.graph_metrics import build_nx, get_top_hubs, get_density
-from mcp.alerts         import check_alerts
+from mcp.alerts import check_alerts
+
+# ── Streamlit telemetry dir fix ───────────────────────────────────────
+os.environ["STREAMLIT_DATA_DIR"] = "/tmp/.streamlit"
+os.environ["XDG_STATE_HOME"] = "/tmp"
+os.environ["STREAMLIT_BROWSER_GATHERUSAGESTATS"] = "false"
+pathlib.Path("/tmp/.streamlit").mkdir(parents=True, exist_ok=True)
 
 ROOT = Path(__file__).parent
 LOGO = ROOT / "assets" / "logo.png"
@@ -59,9 +59,17 @@ def _workspace_sidebar():
 def render_ui():
     st.set_page_config("MedGenesis AI", layout="wide")
 
-    # Persist follow-up input
+    # Initialize session state keys if missing
+    if "query_result" not in st.session_state:
+        st.session_state.query_result = None
     if "followup_input" not in st.session_state:
         st.session_state.followup_input = ""
+    if "followup_response" not in st.session_state:
+        st.session_state.followup_response = None
+    if "last_query" not in st.session_state:
+        st.session_state.last_query = ""
+    if "last_llm" not in st.session_state:
+        st.session_state.last_llm = ""
 
     _workspace_sidebar()
 
@@ -76,6 +84,7 @@ def render_ui():
     llm = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
     query = st.text_input("Enter biomedical question", placeholder="e.g. CRISPR glioblastoma therapy")
 
+    # Alerts
     if get_workspace():
         try:
             news = asyncio.run(check_alerts([w["query"] for w in get_workspace()]))
@@ -87,20 +96,23 @@ def render_ui():
         except Exception:
             pass
 
-    # Trigger search
+    # Run search button
     if st.button("Run Search 🚀") and query:
         with st.spinner("Collecting literature & biomedical data …"):
             res = asyncio.run(orchestrate_search(query, llm=llm))
         st.success(f"Completed with **{res['llm_used'].title()}**")
         st.session_state.query_result = res
+        st.session_state.last_query = query
+        st.session_state.last_llm = llm
         st.session_state.followup_input = ""
-    else:
-        res = st.session_state.get("query_result", None)
+        st.session_state.followup_response = None
+
+    res = st.session_state.query_result
 
     if res:
         tabs = st.tabs(["Results", "Genes", "Trials", "Graph", "Metrics", "Visuals"])
 
-        with tabs[0]:  # Results
+        with tabs[0]:
             for i, p in enumerate(res["papers"], 1):
                 st.markdown(f"**{i}. [{p['title']}]({p['link']})**  *{p['authors']}*")
                 st.write(p["summary"])
@@ -111,22 +123,25 @@ def render_ui():
             with col2:
                 st.download_button("PDF", _pdf(res["papers"]), "papers.pdf", "application/pdf")
             if st.button("💾 Save"):
-                save_query(query, res)
+                save_query(st.session_state.last_query, res)
                 st.success("Saved to workspace")
+
             st.subheader("UMLS concepts")
             for c in res["umls"]:
                 if c.get("cui"):
                     st.write(f"- **{c['name']}** ({c['cui']})")
+
             st.subheader("OpenFDA safety")
             for d in res["drug_safety"]:
                 st.json(d)
+
             st.subheader("AI summary")
             st.info(res["ai_summary"])
 
-        with tabs[1]:  # Genes
+        with tabs[1]:
             st.header("Gene / Variant signals")
             for g in res["genes"]:
-                st.write(f"- **{g.get('name', g.get('geneid'))}** {g.get('description', '')}")
+                st.write(f"- **{g.get('name', g.get('geneid'))}** {g.get('description','')}")
             if res["gene_disease"]:
                 st.markdown("### DisGeNET links")
                 st.json(res["gene_disease"][:15])
@@ -136,7 +151,7 @@ def render_ui():
                     if d:
                         st.write("-", d)
 
-        with tabs[2]:  # Trials
+        with tabs[2]:
             st.header("Clinical trials")
             if not res["clinical_trials"]:
                 st.info("No trials (rate-limited or none found).")
@@ -144,7 +159,7 @@ def render_ui():
                 st.markdown(f"**{t['NCTId'][0]}** – {t['BriefTitle'][0]}")
                 st.write(f"Phase {t.get('Phase',[''])[0]} | Status {t['OverallStatus'][0]}")
 
-        with tabs[3]:  # Graph
+        with tabs[3]:
             nodes, edges, cfg = build_agraph(res["papers"], res["umls"], res["drug_safety"])
             hl = st.text_input("Highlight node:", key="hl")
             if hl:
@@ -153,7 +168,7 @@ def render_ui():
                     n.color = "#f1c40f" if pat.search(n.label) else "#d3d3d3"
             agraph(nodes, edges, cfg)
 
-        with tabs[4]:  # Metrics
+        with tabs[4]:
             nodes, edges, _ = build_agraph(res["papers"], res["umls"], res["drug_safety"])
             G = build_nx([n.__dict__ for n in nodes], [e.__dict__ for e in edges])
             st.metric("Density", f"{get_density(G):.3f}")
@@ -162,26 +177,30 @@ def render_ui():
                 lab = next((n.label for n in nodes if n.id == nid), nid)
                 st.write(f"- {lab}  {sc:.3f}")
 
-        with tabs[5]:  # Visuals
+        with tabs[5]:
             years = [p["published"] for p in res["papers"] if p.get("published")]
             if years:
                 st.plotly_chart(px.histogram(years, nbins=12, title="Publication Year"))
 
+        # Follow-up Q&A block with callback
         st.markdown("---")
-        follow = st.text_input(
-            "Ask follow‑up question:",
-            value=st.session_state.followup_input,
-            key="followup_input"
-        )
-        if st.button("Ask AI"):
-            st.session_state.followup_input = follow
+        st.text_input("Ask follow‑up question:", key="followup_input")
+        def handle_followup():
+            follow = st.session_state.followup_input
             if follow.strip():
-                with st.spinner("Generating AI response..."):
-                    ans = asyncio.run(answer_ai_question(
-                        follow, context=query, llm=llm))
-                st.write(ans["answer"])
+                ans = asyncio.run(answer_ai_question(
+                    follow,
+                    context=st.session_state.last_query,
+                    llm=st.session_state.last_llm))
+                st.session_state.followup_response = ans["answer"]
             else:
-                st.warning("Please type a follow-up question before submitting.")
+                st.session_state.followup_response = None
+
+        st.button("Ask AI", on_click=handle_followup)
+
+        if st.session_state.followup_response:
+            st.write(st.session_state.followup_response)
+
     else:
         st.info("Enter a question and press **Run Search 🚀**")