mcp/knowledge_graph.py

# mcp/knowledge_graph.py

from streamlit_agraph import agraph, Node, Edge, Config

def build_agraph(papers, umls, drugs):
    """
    Build nodes and edges for streamlit-agraph visualization.
    """
    nodes, edges = [], []

    # Map concepts
    for c in umls:
        if c.get("cui"):
            cid = f"concept_{c['cui']}"
            nodes.append(Node(id=cid, label=c["name"], size=25, color="#00b894"))
    
    # Map drugs
    for i, drug_report in enumerate(drugs):
        if drug_report:
            did = f"drug_{i}"
            dname = drug_report.get("drug_name", f"drug_{i}")
            nodes.append(Node(id=did, label=dname, size=25, color="#d35400"))

    # Map papers
    for i, p in enumerate(papers):
        pid = f"paper_{i}"
        nodes.append(Node(id=pid, label=f"P{i+1}", tooltip=p["title"], size=15, color="#0984e3"))
        # connect to concepts
        for c in umls:
            if c["name"].lower() in (p["title"] + p["summary"]).lower():
                edges.append(Edge(source=pid, target=f"concept_{c['cui']}", label="mentions"))
        # connect to drugs
        for j, drug_report in enumerate(drugs):
            dname = drug_report.get("drug_name", "")
            if dname and dname.lower() in (p["title"] + p["summary"]).lower():
                edges.append(Edge(source=pid, target=f"drug_{j}", label="mentions"))
    
    config = Config(
        width="100%", height="600", directed=False,
        nodeHighlightBehavior=True, highlightColor="#f0a",
        collapsible=True
    )
    return nodes, edges, config