app.py

# app.py

import os
import streamlit as st
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import asyncio
from pathlib import Path
import pandas as pd
from fpdf import FPDF
import plotly.express as px
import re

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

ROOT = Path(__file__).parent
LOGO = ROOT / "assets" / "logo.png"

api = FastAPI(
    title="MedGenesis MCP Server",
    version="2.0.0",
    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):
    return await orchestrate_search(data.query)

@api.post("/ask_ai")
async def ask_ai_endpoint(question: str, context: str = ""):
    return await answer_ai_question(question, context)

def generate_pdf(papers):
    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.ln(2)
    return pdf.output(dest="S").encode("latin-1")

def render_ui():
    st.set_page_config(page_title="MedGenesis AI", layout="wide")

    # Sidebar workspace
    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
    col1, col2 = st.columns([0.15, 0.85])
    with col1:
        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("---")
    query = st.text_input("🔍 Enter biomedical research question:", placeholder="e.g. Glioblastoma CRISPR treatments")

    results = None
    if st.button("Run Search 🚀"):
        with st.spinner("Fetching & synthesizing..."):
            results = asyncio.run(orchestrate_search(query))
            st.success("✅ Search complete")

    if results:
        tabs = st.tabs(["📝 Results", "🗺️ Knowledge Graph", "📊 Visualizations"])

        # Tab 1: Results
        with tabs[0]:
            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!")

            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.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 Search & Highlight
        with tabs[1]:
            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"])
                # Highlight logic
                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)
            except Exception as e:
                st.warning("Knowledge graph unavailable: " + str(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)

    # Follow-up Q&A
    st.markdown("---")
    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("Thinking..."):
            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 FastAPI, Streamlit, OpenAI, UMLS, OpenFDA, NCBI")

if __name__ == "__main__":
    import sys
    if "runserver" in sys.argv:
        import uvicorn; uvicorn.run(api, host="0.0.0.0", port=7860)
    else:
        render_ui()