Update app.py

app.py

@@ -1,6 +1,21 @@
-# app.py - MedGenesis AI Streamlit app (OpenAI/Gemini)
-
-import os, pathlib, asyncio, re
+#!/usr/bin/env python3
+"""
+MedGenesis AI – Streamlit front‑end (v3)
+--------------------------------------
+Supports **OpenAI** and **Gemini** engines and the enriched backend
+payload introduced in orchestrator v3:
+    • papers, umls, drug_safety, genes, mesh_defs, gene_disease,
+      clinical_trials, variants, ai_summary
+Tabs:
+    Results | Genes | Trials | Variants | Graph | Metrics | Visuals
+"""
+
+##############################################################################
+# Std‑lib / third‑party
+##############################################################################
+import os
+import pathlib
+import asyncio
 from pathlib import Path
 
 import streamlit as st
@@ -9,25 +24,38 @@ 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.knowledge_graph import build_agraph
 from mcp.graph_metrics import build_nx, get_top_hubs, get_density
 from mcp.alerts import check_alerts
 
-# --- Fix Streamlit temp dir ---
-os.environ["STREAMLIT_DATA_DIR"] = "/tmp/.streamlit"
-os.environ["XDG_STATE_HOME"] = "/tmp"
-os.environ["STREAMLIT_BROWSER_GATHERUSAGESTATS"] = "false"
+# ---------------------------------------------------------------------------
+#  Streamlit telemetry directory → /tmp
+# ---------------------------------------------------------------------------
+os.environ.update({
+    "STREAMLIT_DATA_DIR": "/tmp/.streamlit",
+    "XDG_STATE_HOME"    : "/tmp",
+    "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:
+    """Coerce UTF‑8 → Latin‑1 with replacement (for FPDF)."""
     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()
@@ -36,13 +64,16 @@ def _pdf(papers):
     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.get('title', '')}"))
+        pdf.multi_cell(0, 7, _latin1_safe(f"{i}. {p.get('title','')}"))
         pdf.set_font("Helvetica", "", 9)
-        body = f"{p.get('authors','')}\n{p.get('summary','')}\n{p.get('link','')}\n"
+        body = f"{p.get('authors','')}
+{p.get('summary','')}
+{p.get('link','')}\n"
         pdf.multi_cell(0, 6, _latin1_safe(body))
         pdf.ln(1)
     return pdf.output(dest="S").encode("latin-1", "replace")
 
+
 def _workspace_sidebar():
     with st.sidebar:
         st.header("🗂️ Workspace")
@@ -52,141 +83,140 @@ def _workspace_sidebar():
             return
         for i, item in enumerate(ws, 1):
             with st.expander(f"{i}. {item['query']}"):
-                st.write(item["result"].get("ai_summary", ""))
+                st.write(item['result']['ai_summary'])
+
+##############################################################################
+# Main UI renderer
+##############################################################################
 
 def render_ui():
     st.set_page_config("MedGenesis AI", layout="wide")
 
-    # Session state
-    for k, v in [
-        ("query_result", None), ("followup_input", ""),
-        ("followup_response", None), ("last_query", ""), ("last_llm", "")
-    ]:
-        if k not in st.session_state:
-            st.session_state[k] = v
+    # Session‑state defaults
+    defaults = dict(
+        query_result=None,
+        followup_input="",
+        followup_response=None,
+        last_query="",
+        last_llm="openai",
+    )
+    for k, v in defaults.items():
+        st.session_state.setdefault(k, v)
 
     _workspace_sidebar()
-    c1, c2 = st.columns([0.15, 0.85])
-    with c1:
+
+    # Header
+    col1, col2 = st.columns([0.15, 0.85])
+    with col1:
         if LOGO.exists():
             st.image(str(LOGO), width=105)
-    with c2:
+    with col2:
         st.markdown("## 🧬 **MedGenesis AI**")
-        st.caption("Multi-source biomedical assistant · OpenAI / Gemini")
+        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")
+    # Controls
+    engine = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
+    query  = st.text_input("Enter biomedical question", placeholder="e.g. CRISPR glioblastoma therapy")
 
     # Alerts
-    wsq = get_workspace()
-    if wsq:
+    if get_workspace():
         try:
-            news = asyncio.run(check_alerts([w["query"] for w in wsq]))
-            if news:
+            alerts = asyncio.run(check_alerts([w["query"] for w in get_workspace()]))
+            if alerts:
                 with st.sidebar:
                     st.subheader("🔔 New papers")
-                    for q, lnks in news.items():
+                    for q, lnks in alerts.items():
                         st.write(f"**{q}** – {len(lnks)} new")
         except Exception:
             pass
 
+    # Run 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.get('llm_used','LLM').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
+            res = asyncio.run(orchestrate_search(query, llm=engine))
+        st.session_state.update(
+            query_result=res,
+            last_query=query,
+            last_llm=engine,
+            followup_input="",
+            followup_response=None,
+        )
+        st.success(f"Completed with **{res['llm_used'].title()}**")
 
     res = st.session_state.query_result
     if not res:
         st.info("Enter a question and press **Run Search 🚀**")
         return
 
+    # Tabs
     tabs = st.tabs(["Results", "Genes", "Trials", "Variants", "Graph", "Metrics", "Visuals"])
-    # --------------- Results Tab ---------------
+
+    # --- Results tab
     with tabs[0]:
-        for i, p in enumerate(res.get("papers", []), 1):
+        st.subheader("Literature")
+        for i, p in enumerate(res['papers'], 1):
             st.markdown(f"**{i}. [{p.get('title','')}]({p.get('link','')})**  *{p.get('authors','')}*")
-            st.write(p.get("summary", ""))
-        col1, col2 = st.columns(2)
-        with col1:
-            st.download_button("CSV", pd.DataFrame(res.get("papers", [])).to_csv(index=False),
-                               "papers.csv", "text/csv")
-        with col2:
-            st.download_button("PDF", _pdf(res.get("papers", [])), "papers.pdf", "application/pdf")
+            st.write(p.get('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"):
             save_query(st.session_state.last_query, res)
             st.success("Saved to workspace")
+
         st.subheader("UMLS concepts")
-        for c in res.get("umls", []):
-            if isinstance(c, dict) and c.get("cui"):
+        for c in res['umls']:
+            if c.get('cui'):
                 st.write(f"- **{c.get('name','')}** ({c.get('cui')})")
+
         st.subheader("OpenFDA safety signals")
-        st.json(res.get("drug_safety", []))
+        for d in res['drug_safety']:
+            st.json(d)
+
         st.subheader("AI summary")
-        st.info(res.get("ai_summary", ""))
+        st.info(res['ai_summary'])
 
-    # --------------- Genes Tab ---------------
+    # --- Genes tab
     with tabs[1]:
         st.header("Gene / Variant signals")
-        genes = res.get("genes", [])
-        if not genes:
-            st.info("No gene hits (rate-limited or none found).")
-        else:
-            for g in genes:
-                if isinstance(g, dict):
-                    lab = g.get("name") or g.get("symbol") or g.get("geneid")
-                    st.write(f"- **{lab}** {g.get('description','')}")
-        if res.get("gene_disease"):
-            st.markdown("### DisGeNET associations")
-            st.json(res.get("gene_disease")[:15])
-        if res.get("mesh_defs"):
+        for g in res['genes']:
+            sym = g.get('symbol') or g.get('name') or ''
+            st.write(f"- **{sym}**")
+        if res['mesh_defs']:
             st.markdown("### MeSH definitions")
-            for d in res["mesh_defs"]:
-                if d:
-                    st.write("-", d)
+            for d in res['mesh_defs']:
+                st.write(f"- {d}")
+        if res['gene_disease']:
+            st.markdown("### DisGeNET links")
+            st.json(res['gene_disease'][:15])
 
-    # --------------- Trials Tab ---------------
+    # --- Trials tab
     with tabs[2]:
         st.header("Clinical trials")
-        trials = res.get("clinical_trials", [])
+        trials = res['clinical_trials']
         if not trials:
-            st.info("No trials (rate-limited or none found).")
-        else:
-            for t in trials:
-                nct = t.get("nctId") or (t.get("NCTId", [""])[0] if isinstance(t.get("NCTId"), list) else "")
-                title = t.get("briefTitle") or (t.get("BriefTitle", [""])[0] if isinstance(t.get("BriefTitle"), list) else "")
-                phase = t.get("phase") or (t.get("Phase", [""])[0] if isinstance(t.get("Phase"), list) else "")
-                status = t.get("status") or (t.get("OverallStatus", [""])[0] if isinstance(t.get("OverallStatus"), list) else "")
-                st.markdown(f"**{nct}** – {title}")
-                st.write(f"Phase {phase} | Status {status}")
-
-    # --------------- Variants Tab ---------------
+            st.info("No trials returned (rate‑limited or none found).")
+        for t in trials:
+            st.markdown(f"**{t.get('nctId','')}** – {t.get('briefTitle','')}  Phase {t.get('phase','?')} | Status {t.get('status','?')}")
+
+    # --- Variants tab
     with tabs[3]:
         st.header("Cancer variants (cBioPortal)")
-        variants = res.get("variants", [])
+        variants = res['variants']
         if not variants:
-            st.info("No variant data.")
+            st.info("No variants for this gene/profile.")
         else:
-            for v in variants:
-                st.json(v)
+            st.json(variants[:30])
 
-    # --------------- Graph Tab ---------------
+    # --- Graph tab
     with tabs[4]:
-        nodes, edges, cfg = build_agraph(res.get("papers", []), res.get("umls", []), res.get("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"
+        nodes, edges, cfg = build_agraph(res['papers'], res['umls'], res['drug_safety'])
         agraph(nodes, edges, cfg)
 
-    # --------------- Metrics Tab ---------------
+    # --- Metrics tab
     with tabs[5]:
-        nodes, edges, _ = build_agraph(res.get("papers", []), res.get("umls", []), res.get("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**")
@@ -194,28 +224,31 @@ def render_ui():
             lab = next((n.label for n in nodes if n.id == nid), nid)
             st.write(f"- {lab}  {sc:.3f}")
 
-    # --------------- Visuals Tab ---------------
+    # --- Visuals tab
     with tabs[6]:
-        years = [p.get("published", "") for p in res.get("papers", []) if p.get("published")]
+        years = [p.get('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 ---------------
+    # Follow‑up QA
     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.get("answer", "No answer.")
-        else:
-            st.session_state.followup_response = None
-    st.button("Ask AI", on_click=handle_followup)
+
+    def _on_ask():
+        q = st.session_state.followup_input
+        if not q.strip():
+            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)
 
+##############################################################################
+# Entrypoint
+##############################################################################
 if __name__ == "__main__":
     render_ui()