Update app.py

app.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
-# MedGenesis AI · CPU-only Streamlit app (OpenAI / Gemini)
+# MedGenesis AI · CPU-only Streamlit front-end (OpenAI / Gemini)
 
-import os, pathlib
-import asyncio, re
+from __future__ import annotations
+import os, pathlib, asyncio, re
 from pathlib import Path
 
 import streamlit as st
@@ -11,25 +11,29 @@ import plotly.express as px
 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.orchestrator    import orchestrate_search, answer_ai_question
+from mcp.workspace       import get_workspace, save_query, clear_workspace
 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_utils     import build_nx, get_top_hubs, get_density
+from mcp.alerts          import check_alerts
 
-# ── Streamlit telemetry dir fix ───────────────────────────────────────
-os.environ["STREAMLIT_DATA_DIR"] = "/tmp/.streamlit"
-os.environ["XDG_STATE_HOME"] = "/tmp"
+# ── 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"
 
+# -------------------------------------------------------------------#
+# Utility helpers                                                    #
+# -------------------------------------------------------------------#
 def _latin1_safe(txt: str) -> str:
     return txt.encode("latin-1", "replace").decode("latin-1")
 
-def _pdf(papers):
+
+def _pdf(papers: list[dict]) -> bytes:
     pdf = FPDF()
     pdf.set_auto_page_break(auto=True, margin=15)
     pdf.add_page()
@@ -45,34 +49,41 @@ def _pdf(papers):
         pdf.ln(1)
     return pdf.output(dest="S").encode("latin-1", "replace")
 
-def _workspace_sidebar():
+
+def _workspace_sidebar() -> None:
     with st.sidebar:
         st.header("🗂️ Workspace")
         ws = get_workspace()
         if not ws:
             st.info("Run a search then press **Save** to populate this list.")
             return
+        if st.button("Clear workspace 🗑️"):
+            clear_workspace()
+            st.experimental_rerun()
         for i, item in enumerate(ws, 1):
             with st.expander(f"{i}. {item['query']}"):
                 st.write(item["result"]["ai_summary"])
 
-def render_ui():
+
+# -------------------------------------------------------------------#
+# Streamlit main UI                                                  #
+# -------------------------------------------------------------------#
+def render_ui() -> None:
     st.set_page_config("MedGenesis AI", layout="wide")
 
-    # 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 = ""
+    # ── session_state bootstrap ────────────────────────────────────
+    for key, default in {
+        "query_result"      : None,
+        "followup_input"    : "",
+        "followup_response" : None,
+        "last_query"        : "",
+        "last_llm"          : "",
+    }.items():
+        st.session_state.setdefault(key, default)
 
     _workspace_sidebar()
 
+    # ── header ─────────────────────────────────────────────────────
     c1, c2 = st.columns([0.15, 0.85])
     with c1:
         if LOGO.exists():
@@ -81,13 +92,14 @@ def render_ui():
         st.markdown("## 🧬 **MedGenesis AI**")
         st.caption("Multi-source biomedical assistant · OpenAI / Gemini")
 
-    llm = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
-    query = st.text_input("Enter biomedical question", placeholder="e.g. CRISPR glioblastoma therapy")
+    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():
+    # ── alerts for saved queries ───────────────────────────────────
+    if ws := get_workspace():
         try:
-            news = asyncio.run(check_alerts([w["query"] for w in get_workspace()]))
+            news = asyncio.run(check_alerts([w["query"] for w in ws]))
             if news:
                 with st.sidebar:
                     st.subheader("🔔 New papers")
@@ -96,113 +108,143 @@ def render_ui():
         except Exception:
             pass
 
-    # Run search button
-    if st.button("Run Search 🚀") and query:
+    # ── primary search trigger ─────────────────────────────────────
+    if st.button("Run Search 🚀") and query.strip():
         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 = ""
-        st.session_state.followup_response = None
+        st.session_state.update({
+            "query_result"      : res,
+            "last_query"        : query,
+            "last_llm"          : llm,
+            "followup_input"    : "",
+            "followup_response" : None,
+        })
 
     res = st.session_state.query_result
-
-    if res:
-        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(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]:
-            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]:
-            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}")
-
-        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("---")
-        st.text_input("Ask follow‑up question:", key="followup_input")
-        def handle_followup():
-            follow = st.session_state.followup_input
-            if follow.strip():
-                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.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:
+    if not res:
         st.info("Enter a question and press **Run Search 🚀**")
+        return
+
+    # ----------------------------------------------------------------#
+    # Tabs                                                            #
+    # ----------------------------------------------------------------#
+    tabs = st.tabs(["Results", "Genes", "Trials", "Graph", "Metrics", "Visuals"])
+
+    # 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"])
+
+        c1, c2 = st.columns(2)
+        with c1:
+            st.download_button("CSV",
+                               pd.DataFrame(res["papers"]).to_csv(index=False),
+                               "papers.csv", "text/csv")
+        with c2:
+            st.download_button("PDF", _pdf(res["papers"]),
+                               "papers.pdf", "application/pdf")
+
+        if st.button("💾 Save this result"):
+            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"])
+
+    # Genes -----------------------------------------------------------
+    with tabs[1]:
+        st.header("Gene / Variant signals")
+        if not res["genes"]:
+            st.info("No gene hits (rate-limited or none found).")
+        for g in res["genes"]:
+            st.write(f"- **{g.get('symbol', g.get('name', ''))}**  "
+                     f"{g.get('summary', '')[:120]}…")
+
+        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)
+
+    # Trials ----------------------------------------------------------
+    with tabs[2]:
+        st.header("Clinical trials")
+        trials = res["clinical_trials"]
+        if not trials:
+            st.info("No trials (rate-limited or none found).")
+        for t in trials:
+            st.markdown(f"**{t['nctId']}** – {t['briefTitle']}")
+            st.write(f"Phase {t.get('phase','')} | Status {t.get('status')}")
+
+    # Graph -----------------------------------------------------------
+    with tabs[3]:
+        nodes, edges, cfg = build_agraph(
+            res["papers"], res["umls"], res["drug_safety"],
+            res["genes"], res["clinical_trials"], res.get("ot_associations", [])
+        )
+        hl = st.text_input("Highlight node:")
+        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)
+
+    # Metrics ---------------------------------------------------------
+    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}")
+
+    # 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                                                 #
+    # ----------------------------------------------------------------#
+    st.markdown("---")
+    st.text_input("Ask follow-up question:",
+                  key="followup_input",
+                  placeholder="e.g. Any phase III trials recruiting now?")
+
+    def _on_ask() -> None:
+        q = st.session_state.followup_input.strip()
+        if not q:
+            st.warning("Please type a question first.")
+            return
+        with st.spinner("Querying LLM …"):
+            ans = asyncio.run(
+                answer_ai_question(q,
+                                   context=st.session_state.last_query,
+                                   llm=st.session_state.last_llm))
+        st.session_state.followup_response = ans["answer"]
+
+    st.button("Ask AI", on_click=_on_ask)
+    if st.session_state.followup_response:
+        st.write(st.session_state.followup_response)
+
 
+# -------------------------------------------------------------------#
 if __name__ == "__main__":
     render_ui()