Update app.py

app.py

@@ -1,6 +1,6 @@
-# app.py
-import asyncio, re
-from pathlib import Path
+# ── app.py ─────────────────────────────────────────────────────────────
+import asyncio
+import re
 
 import streamlit as st
 import pandas as pd
@@ -8,53 +8,63 @@ 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.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
+from mcp.graph_metrics import build_nx, get_top_hubs, get_density
+from mcp.protocols import draft_protocol
 
+# ── Streamlit setup ────────────────────────────────────────────────────
 st.set_page_config(layout="wide", page_title="MedGenesis AI")
 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)
-query= st.text_input("Enter biomedical question")
+llm = st.radio("LLM engine", ["openai", "gemini"], horizontal=True)
+query = st.text_input("Enter biomedical question")
+
+# PDF helper
 
 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")
+    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','')}
+{p.get('summary','')}
+{p.get('link','')}"
+        pdf.multi_cell(0, 6, body); pdf.ln(3)
+    return pdf.output(dest="S").encode("latin-1", errors="replace")
 
+# Run search
 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 not res:
     st.info("Enter a query and press Run Search")
     st.stop()
 
+# Tabs
+tabs = st.tabs(["Results", "Graph", "Clusters", "Variants", "Trials", "Metrics", "Visuals", "Protocols"])
+
 # ── Results tab
-tabs = st.tabs(["Results","Graph","Variants","Trials","Metrics","Visuals"])
 with tabs[0]:
-    for i,p in enumerate(res["papers"],1):
+    for i, p in enumerate(res["papers"], 1):
         st.markdown(f"**{i}. [{p['title']}]({p['link']})**")
         st.write(p["summary"])
-    c1,c2 = st.columns(2)
+    c1, c2 = st.columns(2)
     c1.download_button("CSV", pd.DataFrame(res["papers"]).to_csv(index=False),
-                       "papers.csv","text/csv")
+                       "papers.csv", "text/csv")
     c2.download_button("PDF", _make_pdf(res["papers"]),
-                       "papers.pdf","application/pdf")
+                       "papers.pdf", "application/pdf")
     st.subheader("AI summary"); st.info(res["ai_summary"])
 
 # ── Graph tab
 with tabs[1]:
-    nodes,edges,cfg = build_agraph(
+    nodes, edges, cfg = build_agraph(
         res["papers"], res["umls"], res["drug_safety"], res["umls_relations"]
     )
     hl = st.text_input("Highlight node:", key="hl")
@@ -64,40 +74,88 @@ with tabs[1]:
             n.color = "#f1c40f" if pat.search(n.label) else n.color
     agraph(nodes, edges, cfg)
 
-# ── Variants tab
+# ── Clusters tab
 with tabs[2]:
-    if res["variants"]:
+    clusters = res.get("clusters", [])
+    if clusters:
+        df = pd.DataFrame({
+            "title": [p['title'] for p in res['papers']],
+            "cluster": clusters
+        })
+        st.write("### Paper Clusters")
+        for c in sorted(set(clusters)):
+            st.write(f"**Cluster {c}**")
+            for t in df[df['cluster'] == c]['title']:
+                st.write(f"- {t}")
+    else:
+        st.info("No clusters to show.")
+
+# ── Variants tab
+with tabs[3]:
+    if res.get("variants"):
         st.json(res["variants"])
     else:
-        st.warning("No variants found. Try ‘TP53’ or ‘BRCA1’.")
+        st.warning("No variants found. Try 'TP53' or 'BRCA1'.")
 
 # ── Trials tab
-with tabs[3]:
-    if res["clinical_trials"]:
+with tabs[4]:
+    if res.get("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])
+with tabs[5]:
+    nodes_dicts = [n.__dict__ for n in nodes]
+    edges_dicts = [e.__dict__ for e in edges]
+    G = build_nx(nodes_dicts, edges_dicts)
     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)
+    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"])
+with tabs[6]:
+    years = [p.get("published","")[:4] for p in res["papers"] if p.get("published")]
+    if years:
+        st.plotly_chart(px.histogram(years, nbins=10, title="Publication Year"))
+
+# ── Protocols tab
+with tabs[7]:
+    proto_q = st.text_input("Enter hypothesis for protocol:", key="proto_q")
+    if st.button("Draft Protocol") and proto_q.strip():
+        with st.spinner("Generating protocol…"):
+            proto = asyncio.run(draft_protocol(
+                proto_q, context=res['ai_summary'], llm=llm
+            ))
+        st.subheader("Experimental Protocol")
+        st.write(proto)
+ ───────────────────────────────────────────────────
+# In import section:
+from mcp.embeddings import embed_texts, cluster_embeddings
+from mcp.protocols import draft_protocol
+
+# After creating tabs = st.tabs([...,'Clusters','Protocols']):
+with tabs[-2]:  # second last tab = Clusters
+    if res.get('clusters'):
+        df = pd.DataFrame({
+            'title': [p['title'] for p in res['papers']],
+            'cluster': res['clusters']
+        })
+        st.write("### Paper Clusters")
+        for c in sorted(set(res['clusters'])):
+            st.write(f"**Cluster {c}**")
+            for t in df[df['cluster']==c]['title']:
+                st.write(f"- {t}")
+    else:
+        st.info("No clusters to show.")
+
+with tabs[-1]:  # last tab = Protocols
+    proto_q = st.text_input("Enter hypothesis for protocol:", key="proto_q")
+    if st.button("Draft Protocol") and proto_q.strip():
+        with st.spinner("Generating protocol…"):
+            proto = asyncio.run(draft_protocol(
+                proto_q, context=res['ai_summary'], llm=llm))
+        st.subheader("Experimental Protocol")
+        st.write(proto)