Update app.py

app.py

@@ -1,41 +1,103 @@
 # app.py
-import asyncio, streamlit as st, pandas as pd
-from mcp.orchestrator import orchestrate_search
+import asyncio, re
+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.knowledge_graph import build_agraph
+from mcp.graph_metrics   import build_nx, get_top_hubs, get_density
 
 st.set_page_config(layout="wide", page_title="MedGenesis AI")
-if "res" not in st.session_state: st.session_state.res = None
+if "res" not in st.session_state:
+    st.session_state.res = None
 
 st.title("🧬 MedGenesis AI")
-llm = st.radio("LLM engine", ["openai","gemini"], horizontal=True)
-q = st.text_input("Enter biomedical question")
-if st.button("Run Search") and q:
-    with st.spinner("Fetching data…"):
-        st.session_state.res = asyncio.run(orchestrate_search(q, llm=llm))
+llm   = st.radio("LLM engine", ["openai","gemini"], horizontal=True)
+query= st.text_input("Enter biomedical question")
+
+def _make_pdf(papers):
+    pdf = FPDF(); pdf.add_page(); pdf.set_font("Helvetica",size=12)
+    pdf.cell(0,10,"MedGenesis AI – Results",ln=True,align="C"); pdf.ln(5)
+    for i,p in enumerate(papers,1):
+        pdf.set_font("Helvetica","B",11); pdf.multi_cell(0,7,f"{i}. {p.get('title','')}")
+        pdf.set_font("Helvetica",size=9)
+        body = f"{p.get('authors','')}\n{p.get('summary','')}\n{p.get('link','')}"
+        pdf.multi_cell(0,6,body); pdf.ln(3)
+    return pdf.output(dest="S").encode("latin-1",errors="replace")
 
+if st.button("Run Search 🚀") and query:
+    with st.spinner("Gathering data…"):
+        st.session_state.res = asyncio.run(orchestrate_search(query, llm))
 res = st.session_state.res
-if res:
-    st.subheader("🔬 Papers")
-    for p in res["papers"]:
-        st.markdown(f"**[{p['title']}]({p['link']})**  –  {p['authors']}")
+
+if not res:
+    st.info("Enter a query and press Run Search")
+    st.stop()
+
+# ── Results tab
+tabs = st.tabs(["Results","Graph","Variants","Trials","Metrics","Visuals"])
+with tabs[0]:
+    for i,p in enumerate(res["papers"],1):
+        st.markdown(f"**{i}. [{p['title']}]({p['link']})**")
         st.write(p["summary"])
-    st.subheader("💡 AI Summary")
-    st.info(res["ai_summary"])
-
-    tabs = st.tabs(["Graph","Variants","Trials"])
-    with tabs[0]:
-        from mcp.knowledge_graph import build_agraph
-        nodes, edges, cfg = build_agraph(res)
-        from streamlit_agraph import agraph
-        agraph(nodes, edges, cfg)
-    with tabs[1]:
-        if res["variants"]:
-            st.json(res["variants"])
-        else:
-            st.warning("No variants found. Try TP53 or BRCA1.")
-    with tabs[2]:
-        if res["trials"]:
-            st.json(res["trials"])
-        else:
-            st.warning("No trials. Try a disease e.g. ‘Breast Neoplasms’ or a drug.")
-else:
-    st.info("Enter a query and press Run Search.")
+    c1,c2 = st.columns(2)
+    c1.download_button("CSV", pd.DataFrame(res["papers"]).to_csv(index=False),
+                       "papers.csv","text/csv")
+    c2.download_button("PDF", _make_pdf(res["papers"]),
+                       "papers.pdf","application/pdf")
+    st.subheader("AI summary"); st.info(res["ai_summary"])
+
+# ── Graph tab
+with tabs[1]:
+    nodes,edges,cfg = build_agraph(
+        res["papers"], res["umls"], res["drug_safety"], res["umls_relations"]
+    )
+    hl = st.text_input("Highlight node:", key="hl")
+    if hl:
+        pat = re.compile(re.escape(hl), re.I)
+        for n in nodes:
+            n.color = "#f1c40f" if pat.search(n.label) else n.color
+    agraph(nodes, edges, cfg)
+
+# ── Variants tab
+with tabs[2]:
+    if res["variants"]:
+        st.json(res["variants"])
+    else:
+        st.warning("No variants found. Try ‘TP53’ or ‘BRCA1’.")
+
+# ── Trials tab
+with tabs[3]:
+    if res["clinical_trials"]:
+        st.json(res["clinical_trials"])
+    else:
+        st.warning("No trials found. Try a disease or drug.")
+
+# ── Metrics tab
+with tabs[4]:
+    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):
+        lbl = next((n.label for n in nodes if n.id==nid), nid)
+        st.write(f"- {lbl}: {sc:.3f}")
+
+# ── Visuals tab
+with tabs[5]:
+    yrs = [p.get("published","")[:4] for p in res["papers"] if p.get("published")]
+    if yrs:
+        st.plotly_chart(px.histogram(yrs,nbins=10,title="Publication Year"))
+
+# ── Follow-up QA
+st.markdown("---")
+q = st.text_input("Ask follow-up question:", key="followup_input")
+if st.button("Ask AI"):
+    with st.spinner("Querying LLM…"):
+        ans = asyncio.run(answer_ai_question(
+            q, context=res["ai_summary"], llm=llm))
+        st.write(ans["answer"])