Update app.py

app.py

@@ -1,12 +1,14 @@
 #!/usr/bin/env python3
 """
-MedGenesis AI – Streamlit UI (v3, June 2025)
+MedGenesis AI – Streamlit UI  (v3.1 • June 2025)
 
 • Dual-LLM selector (OpenAI | Gemini)
-• Tabs: Results | Genes | Trials | Variants | Graph | Metrics | Visuals
+• Tabs:
+      Results | Genes | Trials | Variants | Graph | Metrics | Visuals
 • Robust PDF export (all Unicode → Latin-1 safe)
-• Null-safe handling of any RuntimeError / HTTPStatusError objects that
-  slip through the async pipeline.
+• Null-safe handling of RuntimeError / HTTPStatusError placeholders
+• Metrics tab now converts Edge objects → {'source', 'target'} safely,
+  preventing the KeyError you just saw.
 """
 
 from __future__ import annotations
@@ -16,7 +18,7 @@ from pathlib import Path
 import streamlit as st
 import pandas as pd
 import plotly.express as px
-from streamlit_agraph import agraph
+from streamlit_agraph import agraph, Node, Edge
 from fpdf import FPDF
 
 from mcp.orchestrator    import orchestrate_search, answer_ai_question
@@ -65,7 +67,7 @@ def _workspace_sidebar():
             with st.expander(f"{i}. {item['query']}"):
                 st.write(item["result"]["ai_summary"])
 
-# ── Main UI ─────────────────────────────────────────────────────────
+# ── Main UI ──────────────────────────────────────────────────────────
 def render_ui() -> None:
     st.set_page_config("MedGenesis AI", layout="wide")
 
@@ -112,7 +114,7 @@ def render_ui() -> None:
         st.info("Enter a question and press **Run Search 🚀**")
         return
 
-    # Guarantee all expected keys exist
+    # Guarantee keys
     for k in (
         "papers", "umls", "drug_safety", "genes", "mesh_defs",
         "gene_disease", "clinical_trials", "variants"
@@ -125,7 +127,7 @@ def render_ui() -> None:
         "Graph", "Metrics", "Visuals"
     ])
 
-    # ---- Results ----------------------------------------------------
+    # Results tab -----------------------------------------------------
     with tabs[0]:
         st.subheader("Literature")
         for i, p in enumerate(res["papers"], 1):
@@ -158,7 +160,7 @@ def render_ui() -> None:
         st.subheader("AI summary")
         st.info(res["ai_summary"])
 
-    # ---- Genes ------------------------------------------------------
+    # Genes tab -------------------------------------------------------
     with tabs[1]:
         st.header("Gene / Variant signals")
         clean = [g for g in res["genes"] if isinstance(g, dict)]
@@ -179,7 +181,7 @@ def render_ui() -> None:
                 if d:
                     st.write("-", d)
 
-    # ---- Trials -----------------------------------------------------
+    # Trials tab ------------------------------------------------------
     with tabs[2]:
         st.header("Clinical trials")
         if not res["clinical_trials"]:
@@ -189,7 +191,7 @@ def render_ui() -> None:
                 st.markdown(f"**{t['nctId']}** – {t['briefTitle']}")
                 st.write(f"Phase {t.get('phase')} | Status {t.get('status')}")
 
-    # ---- Variants ---------------------------------------------------
+    # Variants tab ----------------------------------------------------
     with tabs[3]:
         st.header("Cancer variants (cBioPortal)")
         if not res["variants"]:
@@ -197,7 +199,7 @@ def render_ui() -> None:
         else:
             st.json(res["variants"][:50])
 
-    # ---- Graph ------------------------------------------------------
+    # Graph tab -------------------------------------------------------
     with tabs[4]:
         nodes, edges, cfg = build_agraph(
             res["papers"], res["umls"], res["drug_safety"]
@@ -209,11 +211,19 @@ def render_ui() -> None:
                 n.color = "#f1c40f" if pat.search(n.label) else "#d3d3d3"
         agraph(nodes, edges, cfg)
 
-    # ---- Metrics ----------------------------------------------------
+    # Metrics tab -----------------------------------------------------
     with tabs[5]:
+        # Convert Edge objects → dicts with guaranteed 'source'/'target'
+        edge_dicts = [
+            {"source": getattr(e, "source", getattr(e, "from", "")),
+             "target": getattr(e, "target", getattr(e, "to", ""))}
+            for e in edges if isinstance(e, Edge)
+            if getattr(e, "source", getattr(e, "from", None))
+            and getattr(e, "target", getattr(e, "to", None))
+        ]
         G = build_nx(
             [n.__dict__ for n in nodes],
-            [e.__dict__ for e in edges],
+            edge_dicts,
         )
         st.metric("Density", f"{get_density(G):.3f}")
         st.markdown("**Top hubs**")
@@ -221,14 +231,14 @@ def render_ui() -> None:
             lab = next((n.label for n in nodes if n.id == nid), nid)
             st.write(f"- {lab}  {sc:.3f}")
 
-    # ---- Visuals ----------------------------------------------------
+    # Visuals tab -----------------------------------------------------
     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=12,
                                          title="Publication Year"))
 
-    # ---- Follow-up QA ----------------------------------------------
+    # Follow-up QA ----------------------------------------------------
     st.markdown("---")
     st.text_input("Ask follow-up question:", key="followup_input")