Update app.py

app.py

@@ -9,18 +9,17 @@ from pathlib import Path
 import pandas as pd
 from fpdf import FPDF
 import plotly.express as px
-import streamlit.components.v1 as components
+import re
 
 from mcp.orchestrator import orchestrate_search, answer_ai_question
 from mcp.schemas import UnifiedSearchInput, UnifiedSearchResult
 from mcp.workspace import get_workspace, save_query
 from mcp.knowledge_graph import build_agraph
+from streamlit_agraph import agraph, Node, Edge, Config
 
-# Paths
 ROOT = Path(__file__).parent
 LOGO = ROOT / "assets" / "logo.png"
 
-# FastAPI setup
 api = FastAPI(
     title="MedGenesis MCP Server",
     version="2.0.0",
@@ -36,7 +35,6 @@ async def unified_search_endpoint(data: UnifiedSearchInput):
 async def ask_ai_endpoint(question: str, context: str = ""):
     return await answer_ai_question(question, context)
 
-# PDF Export
 def generate_pdf(papers):
     pdf = FPDF(); pdf.add_page(); pdf.set_font("Arial", size=12)
     pdf.cell(200, 10, "MedGenesis AI - Search Results", ln=True, align="C")
@@ -48,7 +46,6 @@ def generate_pdf(papers):
         pdf.ln(2)
     return pdf.output(dest="S").encode("latin-1")
 
-# UI
 def render_ui():
     st.set_page_config(page_title="MedGenesis AI", layout="wide")
 
@@ -114,12 +111,21 @@ def render_ui():
             st.subheader("🤖 AI Summary")
             st.info(results["ai_summary"])
 
-        # Tab 2: Knowledge Graph
+        # Tab 2: Knowledge Graph with Search & Highlight
         with tabs[1]:
             st.header("🗺️ Knowledge Graph Explorer")
+            search_term = st.text_input("🔎 Highlight nodes containing:", value="")
             try:
                 nodes, edges, config = build_agraph(results["papers"], results["umls"], results["drug_safety"])
-                from streamlit_agraph import agraph
+                # Highlight logic
+                if search_term.strip():
+                    pattern = re.compile(re.escape(search_term), re.IGNORECASE)
+                    for node in nodes:
+                        if pattern.search(node.label) or (hasattr(node, "tooltip") and pattern.search(getattr(node, "tooltip", ""))):
+                            node.color = "#f1c40f"
+                            node.size = max(node.size, 30)
+                        else:
+                            node.color = "#ddd"
                 agraph(nodes=nodes, edges=edges, config=config)
             except Exception as e:
                 st.warning("Knowledge graph unavailable: " + str(e))
@@ -143,7 +149,6 @@ def render_ui():
     st.markdown("---")
     st.caption("✨ Built by Oluwafemi Idiakhoa • Powered by FastAPI, Streamlit, OpenAI, UMLS, OpenFDA, NCBI")
 
-# Entry
 if __name__ == "__main__":
     import sys
     if "runserver" in sys.argv: