Update app.py

app.py

@@ -13,7 +13,7 @@ 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.knowledge_graph import build_agraph
-from streamlit_agraph import agraph, Node, Edge, Config
+from streamlit_agraph import agraph
 
 ROOT = Path(__file__).parent
 LOGO = ROOT / "assets" / "logo.png"
@@ -35,17 +35,21 @@ def generate_pdf(papers):
 def render_ui():
     st.set_page_config(page_title="MedGenesis AI", layout="wide")
 
+    # Workspace in sidebar
     with st.sidebar:
         st.header("🗂️ Workspace")
-        for i, item in enumerate(get_workspace(), 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"workspace_{i}.csv", "text/csv")
-        if not get_workspace():
-            st.info("Run and save searches here.")
-
+        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
     col1, col2 = st.columns([0.15, 0.85])
     with col1:
         if LOGO.exists():
@@ -55,69 +59,76 @@ def render_ui():
         st.caption("Built by Oluwafemi Idiakhoa • Hugging Face Spaces")
 
     st.markdown("---")
-    query = st.text_input("🔍 Enter research question:", placeholder="e.g. CRISPR treatments for glioblastoma")
+    query = st.text_input("🔍 Enter biomedical research question:", placeholder="e.g. CRISPR treatments for glioblastoma")
 
     results = None
     if st.button("Run Search 🚀"):
-        with st.spinner("Analyzing..."):
+        with st.spinner("Gathering and analyzing data..."):
             results = asyncio.run(orchestrate_search(query))
-            st.success("Search complete")
+            st.success("✅ Search complete")
 
     if results:
         tabs = st.tabs(["📝 Results", "🗺️ Knowledge Graph", "📊 Visualizations"])
 
+        # Tab 1: Results
         with tabs[0]:
-            st.header("📚 Papers")
+            st.header("📚 Most Relevant Papers")
             for i, p in enumerate(results["papers"], 1):
-                st.markdown(f"**{i}. [{p['title']}]({p['link']})** *{p['authors']}*")
+                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)
             if st.button("Save to Workspace"):
-                save_query(query, results); st.success("Saved!")
+                save_query(query, results)
+                st.success("Saved to workspace!")
             df = pd.DataFrame(results["papers"])
-            st.download_button("📥 CSV", df.to_csv(index=False), "results.csv", "text/csv")
+            st.download_button("📥 Download CSV", df.to_csv(index=False), "results.csv", "text/csv")
             pdf = generate_pdf(results["papers"])
-            st.download_button("📄 PDF", pdf, "results.pdf", "application/pdf")
+            st.download_button("📄 Download PDF", pdf, "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','')}")
-            st.subheader("💊 Drug Safety")
+                    st.markdown(f"- **{c['name']}** (CUI: `{c['cui']}`): {c.get('definition','No definition')}")
+            st.subheader("💊 Drug Safety (OpenFDA)")
             for d in results["drug_safety"]:
                 st.json(d)
-            st.subheader("🤖 AI Summary")
+            st.subheader("🤖 AI-Powered Summary")
             st.info(results["ai_summary"])
 
+        # Tab 2: Knowledge Graph (with Node Search Highlight)
         with tabs[1]:
-            st.header("🔎 Knowledge Graph")
-            search_term = st.text_input("Highlight node containing text:", value="")
+            st.header("🗺️ Knowledge Graph Explorer")
+            search_term = st.text_input("🔎 Highlight nodes containing:", value="")
             try:
                 nodes, edges, config = build_agraph(results["papers"], results["umls"], results["drug_safety"])
-                if search_term:
-                    pat = re.compile(re.escape(search_term), re.IGNORECASE)
+                if search_term.strip():
+                    pattern = re.compile(re.escape(search_term), re.IGNORECASE)
                     for node in nodes:
-                        if pat.search(node.label) or (hasattr(node, "tooltip") and pat.search(getattr(node, "tooltip", ""))):
-                            node.color = "#f1c40f"; node.size = max(node.size, 30)
+                        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)
             except Exception as e:
-                st.warning("Graph unavailable: " + str(e))
+                st.warning("Knowledge graph unavailable: " + str(e))
 
+        # Tab 3: Visualizations
         with tabs[2]:
-            years = [p["published"] for p in results["papers"] if p.get("published")]
-            if years:
-                st.plotly_chart(px.histogram(years, title="Pub Year"))
+            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)
 
+    # Follow-up AI Q&A
     st.markdown("---")
-    st.subheader("💬 Ask AI a Follow-up")
-    fq = st.text_input("Ask based on above:", "")
+    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("Thinking..."):
-            ans = asyncio.run(answer_ai_question(fq, context=query))
-            st.write(ans.get("answer", ans))
+        with st.spinner("Analyzing..."):
+            ai_ans = asyncio.run(answer_ai_question(follow_up, context=query))
+            st.write(ai_ans.get("answer", ai_ans))
 
     st.markdown("---")
-    st.caption("✨ Built by Oluwafemi Idiakhoa • Powered by Streamlit & AI")
+    st.caption("✨ Built by Oluwafemi Idiakhoa • Powered by FastAPI, Streamlit, OpenAI, UMLS, OpenFDA, NCBI")
 
 if __name__ == "__main__":
     render_ui()