app.py

#!/usr/bin/env python3
"""
MedGenesis AI – Streamlit front-end (v3)
--------------------------------------
Supports **OpenAI** and **Gemini** engines and the enriched backend
payload introduced in orchestrator v3:
    • papers, umls, drug_safety, genes, mesh_defs, gene_disease,
      clinical_trials, variants, ai_summary
Tabs:
    Results | Genes | Trials | Variants | Graph | Metrics | Visuals
"""

import os
import pathlib
import asyncio
from pathlib import Path

import streamlit as st
import pandas as pd
import plotly.express as px
from fpdf import FPDF
from streamlit_agraph import agraph

from mcp.orchestrator import orchestrate_search, answer_ai_question
from mcp.workspace import get_workspace, save_query
from mcp.knowledge_graph import build_agraph
from mcp.graph_metrics import build_nx, get_top_hubs, get_density
from mcp.alerts import check_alerts

# Streamlit telemetry directory → /tmp
os.environ.update({
    "STREAMLIT_DATA_DIR": "/tmp/.streamlit",
    "XDG_STATE_HOME": "/tmp",
    "STREAMLIT_BROWSER_GATHERUSAGESTATS": "false",
})
pathlib.Path("/tmp/.streamlit").mkdir(parents=True, exist_ok=True)

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


def _latin1_safe(txt: str) -> str:
    """Coerce UTF-8 → Latin-1 with replacement (for FPDF)."""
    return txt.encode("latin-1", "replace").decode("latin-1")


def _pdf(papers: list[dict]) -> bytes:
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
    pdf.set_font("Helvetica", size=11)
    pdf.cell(200, 8, _latin1_safe("MedGenesis AI – Results"), ln=True, align="C")
    pdf.ln(3)
    for i, p in enumerate(papers, 1):
        pdf.set_font("Helvetica", "B", 11)
        pdf.multi_cell(0, 7, _latin1_safe(f"{i}. {p.get('title','')}"))
        pdf.set_font("Helvetica", "", 9)
        body = (
            f"{p.get('authors','')}\n"
            f"{p.get('summary','')}\n"
            f"{p.get('link','')}\n"
        )
        pdf.multi_cell(0, 6, _latin1_safe(body))
        pdf.ln(1)
    return pdf.output(dest="S").encode("latin-1", "replace")


def _workspace_sidebar():
    with st.sidebar:
        st.header("🗂️ Workspace")
        ws = get_workspace()
        if not ws:
            st.info("Run a search then press **Save** to populate this list.")
            return
        for i, item in enumerate(ws, 1):
            with st.expander(f"{i}. {item['query']}"):
                st.write(item['result']['ai_summary'])


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

    # Session-state defaults
    defaults = dict(
        query_result=None,
        followup_input="",
        followup_response=None,
        last_query="",
        last_llm="openai",
    )
    for k, v in defaults.items():
        st.session_state.setdefault(k, v)

    _workspace_sidebar()

    # Header
    col1, col2 = st.columns([0.15, 0.85])
    with col1:
        if LOGO.exists():
            st.image(str(LOGO), width=105)
    with col2:
        st.markdown("## 🧬 **MedGenesis AI**")
        st.caption("Multi-source biomedical assistant · OpenAI / Gemini")

    # Controls
    engine = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
    query = st.text_input("Enter biomedical question", placeholder="e.g. CRISPR glioblastoma therapy")

    # Alerts
    if get_workspace():
        try:
            alerts = asyncio.run(check_alerts([w["query"] for w in get_workspace()]))
            if alerts:
                with st.sidebar:
                    st.subheader("🔔 New papers")
                    for q, lnks in alerts.items():
                        st.write(f"**{q}** – {len(lnks)} new")
        except Exception:
            pass

    # Run Search
    if st.button("Run Search 🚀") and query:
        with st.spinner("Collecting literature & biomedical data …"):
            res = asyncio.run(orchestrate_search(query, llm=engine))
        st.session_state.update(
            query_result=res,
            last_query=query,
            last_llm=engine,
            followup_input="",
            followup_response=None,
        )
        st.success(f"Completed with **{res['llm_used'].title()}**")

    res = st.session_state.query_result
    if not res:
        st.info("Enter a question and press **Run Search 🚀**")
        return

    # Tabs
    tabs = st.tabs(["Results", "Genes", "Trials", "Variants", "Graph", "Metrics", "Visuals"])

    # --- Results tab ---
    with tabs[0]:
        st.subheader("Literature")
        for i, p in enumerate(res['papers'], 1):
            st.markdown(f"**{i}. [{p.get('title','')}]({p.get('link','')})**  *{p.get('authors','')}*")
            st.write(p.get('summary',''))
        c1, c2 = st.columns(2)
        with c1:
            st.download_button("CSV", pd.DataFrame(res['papers']).to_csv(index=False), "papers.csv", "text/csv")
        with c2:
            st.download_button("PDF", _pdf(res['papers']), "papers.pdf", "application/pdf")
        if st.button("💾 Save"):
            save_query(st.session_state.last_query, res)
            st.success("Saved to workspace")

        st.subheader("UMLS concepts")
        for c in res['umls']:
            if c.get('cui'):
                st.write(f"- **{c.get('name','')}** ({c.get('cui')})")

        st.subheader("OpenFDA safety signals")
        for d in res['drug_safety']:
            st.json(d)

        st.subheader("AI summary")
        st.info(res['ai_summary'])

    # --- Genes tab ---
    with tabs[1]:
        st.header("Gene / Variant signals")
        valid_genes = [g for g in res['genes'] if isinstance(g, dict)]
        if valid_genes:
            for g in valid_genes:
                sym = g.get('symbol') or g.get('name') or ''
                st.write(f"- **{sym}**")
        else:
            st.info("No gene signals returned.")

        mesh_list = [d for d in res['mesh_defs'] if isinstance(d, str) and d]
        if mesh_list:
            st.markdown("### MeSH definitions")
            for d in mesh_list:
                st.write(f"- {d}")

        gene_disease = [d for d in res['gene_disease'] if isinstance(d, dict)]
        if gene_disease:
            st.markdown("### DisGeNET links")
            st.json(gene_disease[:15])

    # --- Trials tab ---
    with tabs[2]:
        st.header("Clinical trials")
        trials = res['clinical_trials']
        if not trials:
            st.info(
                "No trials found. Try a disease name (e.g. ‘Breast Neoplasms’) "
                "or specific drug (e.g. ‘Pembrolizumab’)."
            )
        else:
            for t in trials:
                st.markdown(
                    f"**{t.get('nctId','')}** – {t.get('briefTitle','')}  "
                    f"Phase {t.get('phase','?')} | Status {t.get('status','?')}"
                )

    # --- Variants tab ---
    with tabs[3]:
        st.header("Cancer variants (cBioPortal)")
        variants = res['variants']
        if not variants:
            st.info(
                "No variants found. Try a well-known gene symbol like ‘TP53’ or ‘BRCA1’."
            )
        else:
            st.json(variants[:30])

    # --- Graph tab ---
    with tabs[4]:
        nodes, edges, cfg = build_agraph(res['papers'], res['umls'], res['drug_safety'])
        agraph(nodes, edges, cfg)

    # --- Metrics tab ---
    with tabs[5]:
        G = build_nx([n.__dict__ for n in nodes], [e.__dict__ for e in edges])
        st.metric("Density", f"{get_density(G):.3f}")
        st.markdown("**Top hubs**")
        for nid, sc in get_top_hubs(G):
            lab = next((n.label for n in nodes if n.id == nid), nid)
            st.write(f"- {lab}  {sc:.3f}")

    # --- Visuals tab ---
    with tabs[6]:
        years = [p.get('published') for p in res['papers'] if p.get('published')]
        if years:
            st.plotly_chart(px.histogram(years, nbins=12, title="Publication Year"))

    # Follow-up QA (outside tabs)
    st.markdown("---")
    input_col, button_col = st.columns([4, 1])
    with input_col:
        followup = st.text_input("Ask follow-up question:", key="followup_input")
    with button_col:
        if st.button("Ask AI"):
            if followup.strip():
                with st.spinner("Querying LLM …"):
                    ans = asyncio.run(
                        answer_ai_question(
                            question=followup,
                            context=st.session_state.last_query,
                            llm=st.session_state.last_llm,
                        )
                    )
                    st.session_state.followup_response = ans.get('answer', '')
            else:
                st.warning("Please type a question first.")
    if st.session_state.followup_response:
        st.write(st.session_state.followup_response)


if __name__ == "__main__":
    render_ui()