app.py

# app.py

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 pathlib import Path
import pandas as pd
from fpdf import FPDF
import asyncio

ROOT_DIR = Path(__file__).resolve().parent
LOGO_PATH = ROOT_DIR / "assets" / "logo.png"

# --- FASTAPI BACKEND ---

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=["*"],
)

@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)

# --- PDF Export Utility ---

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.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")
        pdf.ln(2)
    pdf_out = pdf.output(dest='S').encode('latin-1')
    return pdf_out

# --- STREAMLIT UI ---

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

    # Header with logo and branding
    col1, col2 = st.columns([0.15, 0.85])
    with col1:
        if LOGO_PATH.exists():
            st.image(str(LOGO_PATH), width=100)
        else:
            st.markdown("🧬")
    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("---")

    # 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")

    results = None
    if st.button("Run Search 🚀"):
        with st.spinner("Thinking... Gathering and analyzing data across 5 systems..."):
            results = asyncio.run(orchestrate_search(query))
            st.success("Search complete! 🎉")

    if results:
        # Papers
        st.markdown("### 📚 Most Relevant Papers")
        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)

        # 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"])

        # Suggested Reading
        st.markdown("### 📖 Suggested Links")
        for link in results["suggested_reading"]:
            st.write(f"- {link}")

    # Follow-up AI 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?")
    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"])

    # 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 ---

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()