Update app.py

app.py

@@ -1,134 +1,148 @@
 # app.py
 
-import os
-import streamlit as st
-import asyncio
+import asyncio, os, re
 from pathlib import Path
+
+import streamlit as st
 import pandas as pd
-from fpdf import FPDF
 import plotly.express as px
-import re
+from fpdf import FPDF
+from streamlit_agraph import agraph
 
 from mcp.orchestrator import orchestrate_search, answer_ai_question
-from mcp.schemas import UnifiedSearchInput, UnifiedSearchResult
-from mcp.workspace import get_workspace, save_query
+from mcp.workspace    import get_workspace, save_query
 from mcp.knowledge_graph import build_agraph
-from streamlit_agraph import agraph
 
 ROOT = Path(__file__).parent
 LOGO = ROOT / "assets" / "logo.png"
 
+# ------------------------- helpers -------------------------
 def generate_pdf(papers):
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
+    pdf = FPDF(); pdf.add_page(); pdf.set_font("Arial", size=12)
     pdf.cell(200, 10, "MedGenesis AI - Search Results", ln=True, align="C")
     pdf.ln(10)
     for i, p in enumerate(papers, 1):
         pdf.set_font("Arial", "B", 12)
         pdf.multi_cell(0, 10, f"{i}. {p['title']}")
         pdf.set_font("Arial", "", 10)
-        pdf.multi_cell(0, 8, f"Authors: {p['authors']}\nLink: {p['link']}\nSummary: {p['summary']}\n")
+        pdf.multi_cell(0, 8,
+            f"Authors: {p['authors']}\nLink: {p['link']}\nSummary: {p['summary']}\n")
         pdf.ln(2)
     return pdf.output(dest="S").encode("latin-1")
 
+# ------------------------- UI -------------------------
 def render_ui():
     st.set_page_config(page_title="MedGenesis AI", layout="wide")
 
-    # Workspace in sidebar
+    # ----- Sidebar ----------
     with st.sidebar:
         st.header("🗂️ Workspace")
-        ws = get_workspace()
-        if ws:
-            for i, item in enumerate(ws, 1):
-                with st.expander(f"{i}. {item['query']}"):
-                    st.write("**AI Summary:**", item["result"]["ai_summary"])
-                    st.write("**Top Paper:**", item["result"]["papers"][0]["title"] if item["result"]["papers"] else "None")
-                    df = pd.DataFrame(item["result"]["papers"])
-                    st.download_button("📥 CSV", df.to_csv(index=False), f"query_{i}.csv", "text/csv")
-        else:
-            st.info("Run & save searches here.")
-
-    # Header and logo
+        for i, it in enumerate(get_workspace(), 1):
+            with st.expander(f"{i}. {it['query']}"):
+                st.write("**AI Summary:**", it["result"]["ai_summary"])
+                df = pd.DataFrame(it["result"]["papers"])
+                st.download_button("CSV", df.to_csv(index=False), f"ws_{i}.csv", "text/csv")
+        if not get_workspace():
+            st.info("Run and save searches to populate workspace.")
+
+    # ----- Header ----------
     col1, col2 = st.columns([0.15, 0.85])
     with col1:
-        if LOGO.exists():
-            st.image(str(LOGO), width=100)
+        if LOGO.exists(): st.image(str(LOGO), width=100)
     with col2:
-        st.markdown("## 🧬 MedGenesis AI\n*Unified Intelligence from PubMed, ArXiv, OpenFDA, UMLS, GPT‑4o*")
-        st.caption("Built by Oluwafemi Idiakhoa • Hugging Face Spaces")
+        st.markdown("## 🧬 MedGenesis AI")
+        st.write("*Unified PubMed, ArXiv, OpenFDA, UMLS, NCBI, DisGeNET, ClinicalTrials & GPT-4o*")
 
     st.markdown("---")
-    query = st.text_input("🔍 Enter biomedical research question:", placeholder="e.g. CRISPR treatments for glioblastoma")
+    query = st.text_input("🔍 Biomedical research question:", placeholder="e.g. CRISPR glioblastoma")
 
-    results = None
-    if st.button("Run Search 🚀"):
-        with st.spinner("Gathering and analyzing data..."):
+    # ------------- Search -------------
+    if st.button("Run Search 🚀") and query:
+        with st.spinner("Collecting literature & biomedical data…"):
             results = asyncio.run(orchestrate_search(query))
-            st.success("✅ Search complete")
+        st.success("Completed!")
 
-    if results:
-        tabs = st.tabs(["📝 Results", "🗺️ Knowledge Graph", "📊 Visualizations"])
+        # -------- Tabs ---------
+        tabs = st.tabs([
+            "📝 Results", "🧬 Genes & Variants", "💊 Clinical Trials",
+            "🗺️ Graph", "📊 Visuals"
+        ])
 
-        # Tab 1: Results
+        # -- Results ----------
         with tabs[0]:
-            st.header("📚 Most Relevant Papers")
+            st.header("📚 Top Papers")
             for i, p in enumerate(results["papers"], 1):
                 st.markdown(f"**{i}. [{p['title']}]({p['link']})**  \n*{p['authors']}* ({p['source']})")
-                st.markdown(f"<div style='color:gray'>{p['summary']}</div>", unsafe_allow_html=True)
+                st.markdown(f"<span style='color:gray'>{p['summary']}</span>", unsafe_allow_html=True)
+
             if st.button("Save to Workspace"):
-                save_query(query, results)
-                st.success("Saved to workspace!")
+                save_query(query, results); st.success("Saved!")
+
             df = pd.DataFrame(results["papers"])
-            st.download_button("📥 Download CSV", df.to_csv(index=False), "results.csv", "text/csv")
-            pdf = generate_pdf(results["papers"])
-            st.download_button("📄 Download PDF", pdf, "results.pdf", "application/pdf")
+            st.download_button("CSV", df.to_csv(index=False), "results.csv", "text/csv")
+            st.download_button("PDF", generate_pdf(results["papers"]), "results.pdf", "application/pdf")
+
             st.subheader("🧠 UMLS Concepts")
             for c in results["umls"]:
                 if c.get("cui"):
-                    st.markdown(f"- **{c['name']}** (CUI: `{c['cui']}`): {c.get('definition','No definition')}")
+                    st.markdown(f"- **{c['name']}** (`{c['cui']}`)")
+
             st.subheader("💊 Drug Safety (OpenFDA)")
             for d in results["drug_safety"]:
                 st.json(d)
-            st.subheader("🤖 AI-Powered Summary")
+
+            st.subheader("🤖 AI Summary")
             st.info(results["ai_summary"])
 
-        # Tab 2: Knowledge Graph (with Node Search Highlight)
+        # -- Genes & Variants ----------
         with tabs[1]:
-            st.header("🗺️ Knowledge Graph Explorer")
-            search_term = st.text_input("🔎 Highlight nodes containing:", value="")
+            st.header("🧬 Gene Associations (NCBI / DisGeNET)")
+            for g in results["genes"]:
+                st.write(f"- **{g.get('name', g.get('geneid'))}** – {g.get('description','')}")
+            if results["gene_disease"]:
+                st.markdown("#### DisGeNET Disease → Gene links")
+                st.json(results["gene_disease"][:15])
+            if results["mesh_definitions"]:
+                st.markdown("#### MeSH Definitions")
+                for d in results["mesh_definitions"]:
+                    if d: st.write(f"- {d}")
+
+        # -- Clinical Trials ----------
+        with tabs[2]:
+            st.header("💊 Registered Clinical Trials")
+            for t in results["clinical_trials"]:
+                st.markdown(f"**{t['NCTId'][0]}** – {t['BriefTitle'][0]}")
+                st.write(f"Condition: {', '.join(t['Condition'])} | Phase: {t.get('Phase',[None])[0]} | Status: {t['OverallStatus'][0]}")
+
+        # -- Graph ----------
+        with tabs[3]:
+            search_term = st.text_input("Highlight node containing:", key="graphsearch")
             try:
-                nodes, edges, config = build_agraph(results["papers"], results["umls"], results["drug_safety"])
-                if search_term.strip():
-                    pattern = re.compile(re.escape(search_term), re.IGNORECASE)
-                    for node in nodes:
-                        if pattern.search(node.label) or (hasattr(node, "tooltip") and pattern.search(getattr(node, "tooltip", ""))):
-                            node.color = "#f1c40f"
-                            node.size = max(node.size, 30)
-                        else:
-                            node.color = "#ddd"
-                agraph(nodes=nodes, edges=edges, config=config)
+                nodes, edges, cfg = build_agraph(results["papers"], results["umls"], results["drug_safety"])
+                if search_term:
+                    pat = re.compile(re.escape(search_term), re.I)
+                    for n in nodes:
+                        if pat.search(n.label): n.color, n.size = "#f1c40f", max(n.size, 30)
+                        else: n.color = "#ddd"
+                agraph(nodes=nodes, edges=edges, config=cfg)
             except Exception as e:
-                st.warning("Knowledge graph unavailable: " + str(e))
+                st.error(f"Graph error: {e}")
 
-        # Tab 3: Visualizations
-        with tabs[2]:
-            pub_years = [p["published"] for p in results["papers"] if p.get("published")]
-            if pub_years:
-                fig = px.histogram(pub_years, nbins=10, title="Publication Year Distribution")
-                st.plotly_chart(fig)
+        # -- Visualizations ----------
+        with tabs[4]:
+            yrs = [p["published"] for p in results["papers"] if p.get("published")]
+            if yrs: st.plotly_chart(px.histogram(yrs, nbins=10, title="Publication Year"))
 
-    # Follow-up AI Q&A
-    st.markdown("---")
-    st.subheader("💬 Ask a Follow-up Question")
-    follow_up = st.text_input("Based on above:", placeholder="e.g. What's the most promising therapy?")
-    if st.button("Ask AI"):
-        with st.spinner("Analyzing..."):
-            ai_ans = asyncio.run(answer_ai_question(follow_up, context=query))
-            st.write(ai_ans.get("answer", ai_ans))
+        # -- Follow-up Q&A ----------
+        st.markdown("---")
+        q = st.text_input("Ask follow-up:", key="follow")
+        if st.button("Ask AI"):
+            ans = asyncio.run(answer_ai_question(q, context=query))
+            st.write(ans["answer"])
 
-    st.markdown("---")
-    st.caption("✨ Built by Oluwafemi Idiakhoa • Powered by FastAPI, Streamlit, OpenAI, UMLS, OpenFDA, NCBI")
+    else:
+        st.info("Enter a question and press **Run Search 🚀**")
 
+# ------------- Run -------------
 if __name__ == "__main__":
     render_ui()