Update app.py

app.py

@@ -4,35 +4,29 @@ import os
 import streamlit as st
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-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_knowledge_graph
+import asyncio
 from pathlib import Path
 import pandas as pd
 from fpdf import FPDF
-import asyncio
 import plotly.express as px
 import streamlit.components.v1 as components
 
-ROOT_DIR = Path(__file__).resolve().parent
-LOGO_PATH = ROOT_DIR / "assets" / "logo.png"
+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
 
-# --- FASTAPI BACKEND ---
+# Paths
+ROOT = Path(__file__).parent
+LOGO = ROOT / "assets" / "logo.png"
 
+# FastAPI setup
 api = FastAPI(
     title="MedGenesis MCP Server",
     version="2.0.0",
-    description="MedGenesis AI unifies PubMed, ArXiv, OpenFDA, UMLS, and GPT-4o into a single biomedical intelligence platform."
-)
-
-api.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+    description="MedGenesis AI – Unified Intelligence from PubMed, ArXiv, OpenFDA, UMLS, GPT-4o"
 )
+api.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
 @api.post("/unified_search", response_model=UnifiedSearchResult)
 async def unified_search_endpoint(data: UnifiedSearchInput):
@@ -42,174 +36,117 @@ async def unified_search_endpoint(data: UnifiedSearchInput):
 async def ask_ai_endpoint(question: str, context: str = ""):
     return await answer_ai_question(question, context)
 
-# --- PDF Export Utility ---
-
+# PDF Export
 def generate_pdf(papers):
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-    pdf.cell(200, 10, txt="MedGenesis AI - Search Results", ln=True, align='C')
+    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, paper in enumerate(papers, 1):
-        pdf.set_font("Arial", style="B", size=12)
-        pdf.multi_cell(0, 10, f"{i}. {paper['title']}")
-        pdf.set_font("Arial", style="", size=10)
-        pdf.multi_cell(0, 8, f"Authors: {paper['authors']}\nLink: {paper['link']}\nSummary: {paper['summary']}\n")
+    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.ln(2)
-    pdf_out = pdf.output(dest='S').encode('latin-1')
-    return pdf_out
-
-# --- STREAMLIT UI ---
+    return pdf.output(dest="S").encode("latin-1")
 
+# UI
 def render_ui():
     st.set_page_config(page_title="MedGenesis AI", layout="wide")
 
-    # --- SIDEBAR WORKSPACE ---
+    # Sidebar workspace
     with st.sidebar:
-        st.header("🗂️ Your Workspace")
-        saved_queries = get_workspace()
-        if saved_queries:
-            for i, item in enumerate(saved_queries, 1):
+        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("**First Paper:**", item["result"]["papers"][0]["title"] if item["result"]["papers"] else "None")
+                    st.write("**Top Paper:**", item["result"]["papers"][0]["title"] if item["result"]["papers"] else "None")
                     df = pd.DataFrame(item["result"]["papers"])
-                    st.download_button(
-                        label="📥 Download as CSV",
-                        data=df.to_csv(index=False),
-                        file_name=f"workspace_query_{i}.csv",
-                        mime="text/csv",
-                    )
+                    st.download_button("📥 CSV", df.to_csv(index=False), f"query_{i}.csv", "text/csv")
         else:
-            st.info("Run a search and save it here!")
+            st.info("Run & save searches here.")
 
-    # --- MAIN APP HEADER ---
+    # Header
     col1, col2 = st.columns([0.15, 0.85])
     with col1:
-        if LOGO_PATH.exists():
-            st.image(str(LOGO_PATH), width=100)
-        else:
-            st.markdown("🧬")
+        if LOGO.exists():
+            st.image(str(LOGO), width=100)
     with col2:
-        st.markdown("""
-            ## 🧬 **MedGenesis AI** – Biomedical Research Reimagined  
-            *Unified Intelligence from PubMed, ArXiv, OpenFDA, UMLS, and GPT-4o*
-        """)
-        st.caption("Created by Oluwafemi Idiakhoa | Hugging Face Spaces")
+        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("---")
-
-    # Unified Semantic Search
-    st.subheader("🔍 Unified Semantic Search")
-    query = st.text_input("Enter your biomedical research question:", placeholder="e.g. New treatments for glioblastoma using CRISPR")
+    query = st.text_input("🔍 Enter biomedical research question:", placeholder="e.g. Glioblastoma CRISPR treatments")
 
     results = None
     if st.button("Run Search 🚀"):
-        with st.spinner("Thinking... Gathering and analyzing data across 5 systems..."):
+        with st.spinner("Fetching & synthesizing..."):
             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 ---
+        # Tab 1: Results
         with tabs[0]:
-            for i, paper in enumerate(results["papers"], 1):
-                st.markdown(f"**{i}. [{paper['title']}]({paper['link']})**  \n*{paper['authors']}* ({paper['source']})")
-                st.markdown(f"<div style='font-size: 0.9em; color: gray'>{paper['summary']}</div>", unsafe_allow_html=True)
-
-            # Save to workspace (user-initiated, for clarity)
-            if st.button("Save this search to Workspace"):
-                save_query(query, results)
-                st.success("Saved to your workspace!")
-
-            # Export as CSV
-            if results["papers"]:
-                df = pd.DataFrame(results["papers"])
-                csv = df.to_csv(index=False)
-                st.download_button(
-                    label="📥 Download results as CSV",
-                    data=csv,
-                    file_name="medgenesis_results.csv",
-                    mime="text/csv",
-                )
-
-            # Export as PDF
-            if results["papers"]:
-                pdf_bytes = generate_pdf(results["papers"])
-                st.download_button(
-                    label="📄 Download results as PDF",
-                    data=pdf_bytes,
-                    file_name="medgenesis_results.pdf",
-                    mime="application/pdf",
-                )
-
-            # UMLS Concepts
-            st.markdown("### 🧠 Biomedical Concept Enrichment (UMLS)")
-            for concept in results["umls"]:
-                if concept["cui"]:
-                    st.markdown(f"🔹 **{concept['name']}** (CUI: `{concept['cui']}`): {concept['definition'] or 'No definition available'}")
-
-            # Drug Safety
-            st.markdown("### 💊 Drug Safety Insights (OpenFDA)")
-            for drug_report in results["drug_safety"]:
-                if drug_report:
-                    st.json(drug_report)
-
-            # AI Summary
-            st.markdown("### 🤖 AI-Powered Summary")
-            st.info(results["ai_summary"])
+            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)
+
+            if st.button("Save to Workspace"):
+                save_query(query, results); st.success("Saved!")
 
-            # Suggested Reading
-            st.markdown("### 📖 Suggested Links")
-            for link in results["suggested_reading"]:
-                st.write(f"- {link}")
+            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")
 
-        # --- TAB 2: Knowledge Graph ---
+            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.subheader("💊 Drug Safety (OpenFDA)")
+            for d in results["drug_safety"]:
+                st.json(d)
+
+            st.subheader("🤖 AI Summary")
+            st.info(results["ai_summary"])
+
+        # Tab 2: Knowledge Graph
         with tabs[1]:
-            st.markdown("#### Explore Connections")
+            st.header("🗺️ Knowledge Graph Explorer")
             try:
-                kg_html_path = build_knowledge_graph(results["papers"], results["umls"], results["drug_safety"])
-                with open(kg_html_path, 'r', encoding='utf-8') as f:
-                    kg_html = f.read()
-                components.html(kg_html, height=600)
+                nodes, edges, config = build_agraph(results["papers"], results["umls"], results["drug_safety"])
+                from streamlit_agraph import agraph
+                agraph(nodes=nodes, edges=edges, config=config)
             except Exception as e:
-                st.warning("⚠️ Knowledge graph visualization is unavailable. "
-                           "Please ensure 'pyvis' and 'jinja2' are installed. "
-                           f"Error: {e}")
+                st.warning("Knowledge graph unavailable: " + str(e))
 
-        # --- TAB 3: Visualizations ---
+        # 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)
-            # Placeholder for more charts
 
-    # Follow-up AI Q&A
+    # Follow-up Q&A
     st.markdown("---")
-    st.subheader("💬 Ask AI a Follow-up Question")
-    follow_up = st.text_input("What do you want to ask based on the above?", placeholder="e.g. What's the most promising therapy?")
+    st.subheader("💬 Ask Follow-up AI Question")
+    follow_up = st.text_input("Based on above:", placeholder="e.g. Most promising therapy?")
     if st.button("Ask AI"):
-        with st.spinner("Analyzing and responding..."):
-            ai_answer = asyncio.run(answer_ai_question(follow_up, context=query))
-            st.success("AI's Response:")
-            st.write(ai_answer["answer"])
+        with st.spinner("Thinking..."):
+            ai_ans = asyncio.run(answer_ai_question(follow_up, context=query))
+            st.write(ai_ans.get("answer", ai_ans))
 
-    # Footer
     st.markdown("---")
-    st.markdown(
-        "<div style='text-align: center; font-size: 0.9em;'>"
-        "✨ Built with ❤️ by <strong>Oluwafemi Idiakhoa</strong> • Powered by FastAPI, Streamlit, Hugging Face, OpenAI, UMLS, OpenFDA, and NCBI</div>",
-        unsafe_allow_html=True
-    )
-
-# --- MAIN ENTRY ---
+    st.caption("✨ Built by Oluwafemi Idiakhoa • Powered by FastAPI, Streamlit, OpenAI, UMLS, OpenFDA, NCBI")
 
+# Entry
 if __name__ == "__main__":
     import sys
     if "runserver" in sys.argv:
-        import uvicorn
-        uvicorn.run(api, host="0.0.0.0", port=7860)
+        import uvicorn; uvicorn.run(api, host="0.0.0.0", port=7860)
     else:
         render_ui()