Update app.py

app.py

@@ -9,7 +9,6 @@ os.environ["XDG_STATE_HOME"] = "/tmp"
 os.environ["STREAMLIT_BROWSER_GATHERUSAGESTATS"] = "false"
 pathlib.Path("/tmp/.streamlit").mkdir(parents=True, exist_ok=True)
 
-# ── Std-lib / third-party imports ────────────────────────────────────
 import asyncio, re
 from pathlib import Path
 import streamlit as st
@@ -18,17 +17,15 @@ import plotly.express as px
 from fpdf import FPDF
 from streamlit_agraph import agraph
 
-# ── Internal helpers ────────────────────────────────────────────────
-from mcp.orchestrator    import orchestrate_search, answer_ai_question
-from mcp.workspace       import get_workspace, save_query
+from mcp.orchestrator import orchestrate_search, answer_ai_question
+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.graph_metrics import build_nx, get_top_hubs, get_density
+from mcp.alerts import check_alerts
 
 ROOT = Path(__file__).parent
 LOGO = ROOT / "assets" / "logo.png"
 
-# ── PDF export helper (UTF-8 → Latin-1 “safe”) ──────────────────────
 def _latin1_safe(txt: str) -> str:
     return txt.encode("latin-1", "replace").decode("latin-1")
 
@@ -37,24 +34,17 @@ def _pdf(papers):
     pdf.set_auto_page_break(auto=True, margin=15)
     pdf.add_page()
     pdf.set_font("Helvetica", size=11)
-    pdf.cell(200, 8, _latin1_safe("MedGenesis AI – Results"),  ln=True, align="C")
+    pdf.cell(200, 8, _latin1_safe("MedGenesis AI – Results"), ln=True, align="C")
     pdf.ln(3)
-
     for i, p in enumerate(papers, 1):
         pdf.set_font("Helvetica", "B", 11)
         pdf.multi_cell(0, 7, _latin1_safe(f"{i}. {p['title']}"))
         pdf.set_font("Helvetica", "", 9)
-        body = (
-            f"{p['authors']}\n"
-            f"{p['summary']}\n"
-            f"{p['link']}\n"
-        )
+        body = f"{p['authors']}\n{p['summary']}\n{p['link']}\n"
         pdf.multi_cell(0, 6, _latin1_safe(body))
         pdf.ln(1)
-
     return pdf.output(dest="S").encode("latin-1", "replace")
 
-# ── Sidebar workspace ───────────────────────────────────────────────
 def _workspace_sidebar():
     with st.sidebar:
         st.header("🗂️ Workspace")
@@ -66,9 +56,15 @@ def _workspace_sidebar():
             with st.expander(f"{i}. {item['query']}"):
                 st.write(item["result"]["ai_summary"])
 
-# ── Main UI ─────────────────────────────────────────────────────────
 def render_ui():
     st.set_page_config("MedGenesis AI", layout="wide")
+
+    # Initialize session-state
+    if "followup_input" not in st.session_state:
+        st.session_state.followup_input = ""
+    if "tab_index" not in st.session_state:
+        st.session_state.tab_index = 0
+
     _workspace_sidebar()
 
     c1, c2 = st.columns([0.15, 0.85])
@@ -79,7 +75,7 @@ def render_ui():
         st.markdown("## 🧬 **MedGenesis AI**")
         st.caption("Multi-source biomedical assistant · OpenAI / Gemini")
 
-    llm   = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
+    llm = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
     query = st.text_input("Enter biomedical question", placeholder="e.g. CRISPR glioblastoma therapy")
 
     # Alert check
@@ -94,98 +90,111 @@ def render_ui():
         except Exception:
             pass
 
-    # Run search
+    # Trigger search
     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()}**")
-
-        tabs = st.tabs(["Results", "Genes", "Trials", "Graph", "Metrics", "Visuals"])
-
-        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"])
-
-            col1, col2 = st.columns(2)
-            with col1:
-                st.download_button("CSV", pd.DataFrame(res["papers"]).to_csv(index=False), "papers.csv", "text/csv")
-            with col2:
-                st.download_button("PDF", _pdf(res["papers"]), "papers.pdf", "application/pdf")
-
-            if st.button("💾 Save"):
-                save_query(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]:
-            st.header("Gene / Variant signals")
-            for g in res["genes"]:
-                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])
-            if res["mesh_defs"]:
-                st.markdown("### MeSH definitions")
-                for d in res["mesh_defs"]:
-                    if d:
-                        st.write("-", d)
-
-        with tabs[2]:
-            st.header("Clinical trials")
-            if not res["clinical_trials"]:
-                st.info("No trials (rate-limited or none found).")
-            for t in res["clinical_trials"]:
-                st.markdown(f"**{t['NCTId'][0]}** – {t['BriefTitle'][0]}")
-                st.write(f"Phase {t.get('Phase',[''])[0]} | Status {t['OverallStatus'][0]}")
-
-        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:
-                pat = re.compile(re.escape(hl), re.I)
-                for n in nodes:
-                    n.color = "#f1c40f" if pat.search(n.label) else "#d3d3d3"
-            agraph(nodes, edges, cfg)
-
-        with tabs[4]:
-            G = build_nx([n.__dict__ for n in nodes], [e.__dict__ for e in edges])
-            st.metric("Density", f"{get_density(G):.3f}")
-            st.markdown("**Top hubs**")
-            for nid, sc in get_top_hubs(G):
-                lab = next((n.label for n in nodes if n.id == nid), nid)
-                st.write(f"- {lab}  {sc:.3f}")
-
-        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 (fixed) ───────────────────────────────────────
+        st.session_state.query_result = res
+        # Reset follow-up input
+        st.session_state.followup_input = ""
+        st.session_state.tab_index = 0
+    else:
+        res = st.session_state.get("query_result", None)
+
+    if res:
+        tabs_list = ["Results", "Genes", "Trials", "Graph", "Metrics", "Visuals"]
+        tabs = st.tabs(tabs_list, index=st.session_state.tab_index)
+        
+        for idx, name in enumerate(tabs_list):
+            with tabs[idx]:
+                st.session_state.tab_index = idx
+                if name == "Results":
+                    for i, p in enumerate(res["papers"], 1):
+                        st.markdown(f"**{i}. [{p['title']}]({p['link']})**  *{p['authors']}*")
+                        st.write(p["summary"])
+                    col1, col2 = st.columns(2)
+                    with col1:
+                        st.download_button("CSV", pd.DataFrame(res["papers"]).to_csv(index=False),
+                                            "papers.csv", "text/csv")
+                    with col2:
+                        st.download_button("PDF", _pdf(res["papers"]), "papers.pdf", "application/pdf")
+                    if st.button("💾 Save"):
+                        save_query(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"])
+
+                elif name == "Genes":
+                    st.header("Gene / Variant signals")
+                    for g in res["genes"]:
+                        st.write(f"- **{g.get('name', g.get('geneid'))}** "
+                                 f"{g.get('description', '')}")
+                    if res["gene_disease"]:
+                        st.markdown("### DisGeNET links")
+                        st.json(res["gene_disease"][:15])
+                    if res["mesh_defs"]:
+                        st.markdown("### MeSH definitions")
+                        for d in res["mesh_defs"]:
+                            if d:
+                                st.write("-", d)
+
+                elif name == "Trials":
+                    st.header("Clinical trials")
+                    if not res["clinical_trials"]:
+                        st.info("No trials (rate-limited or none found).")
+                    for t in res["clinical_trials"]:
+                        st.markdown(f"**{t['NCTId'][0]}** – {t['BriefTitle'][0]}")
+                        st.write(f"Phase {t.get('Phase',[''])[0]} | "
+                                 f"Status {t['OverallStatus'][0]}")
+
+                elif name == "Graph":
+                    nodes, edges, cfg = build_agraph(res["papers"], res["umls"], res["drug_safety"])
+                    hl = st.text_input("Highlight node:", key="hl")
+                    if hl:
+                        pat = re.compile(re.escape(hl), re.I)
+                        for n in nodes:
+                            n.color = "#f1c40f" if pat.search(n.label) else "#d3d3d3"
+                    agraph(nodes, edges, cfg)
+
+                elif name == "Metrics":
+                    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}")
+                    st.markdown("**Top hubs**")
+                    for nid, sc in get_top_hubs(G):
+                        lab = next((n.label for n in nodes if n.id == nid), nid)
+                        st.write(f"- {lab}  {sc:.3f}")
+
+                elif name == "Visuals":
+                    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 persistently under tabs
         st.markdown("---")
-        follow = st.text_input("Ask follow-up question:", key="followup_input")  # ✅ UPDATED
+        follow = st.text_input("Ask follow‑up question:",
+                               value=st.session_state.followup_input,
+                               key="followup_input")
         if st.button("Ask AI"):
-            if follow.strip():  # ✅ UPDATED
+            st.session_state.followup_input = follow
+            if follow.strip():
                 with st.spinner("Generating AI response..."):
-                    ans = asyncio.run(answer_ai_question(follow, context=query, llm=llm))
+                    ans = asyncio.run(answer_ai_question(
+                        follow, context=query, llm=llm))
                 st.write(ans["answer"])
             else:
-                st.warning("Please type a follow-up question before submitting.")  # ✅ UPDATED
-
+                st.warning("Please type a follow-up question before submitting.")
     else:
         st.info("Enter a question and press **Run Search 🚀**")
 
-# entry-point
 if __name__ == "__main__":
     render_ui()