File size: 1,546 Bytes
d5e0cb0 1e92bb0 e202a39 1e92bb0 d5e0cb0 1e92bb0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
# mcp/knowledge_graph.py
from streamlit_agraph import Node, Edge, Config
def build_agraph(res: Dict) -> (list, list, Config):
nodes, edges = [], []
# add each paper as a node
for i,p in enumerate(res["papers"]):
nid = f"paper_{i}"
nodes.append(Node(id=nid, label=p["title"], size=20, color="#0984e3"))
# connect to AI summary?
# add UMLS concepts
for u in res["umls"]:
cid = f"cui_{u['cui']}"
label = f"{u['name']} ({u['cui']})"
nodes.append(Node(id=cid, label=label, size=25, color="#00b894"))
# connect concept → first paper
edges.append(Edge(source=cid, target="paper_0", label="mentioned_in"))
# genes
g = res.get("gene",{})
if g:
gid = "gene_node"
nodes.append(Node(id=gid, label=g.get("symbol",g.get("name","gene")), color="#d63031"))
edges.append(Edge(source=gid, target="cui_"+res["umls"][0]["cui"], label="related"))
# variants
for v in res["variants"]:
vid = f"var_{v['mutationId']}"
nodes.append(Node(id=vid, label=v["mutationId"], color="#fdcb6e", size=15))
edges.append(Edge(source=vid, target=gid, label="affects"))
# trials
for t in res["trials"]:
tid = t["NCTId"][0]
nodes.append(Node(id=tid, label=tid, color="#6c5ce7"))
edges.append(Edge(source=tid, target=gid, label="studies"))
cfg = Config(width="100%", height="600", directed=True,
nodeHighlightBehavior=True, highlightColor="#fdcb6e")
return nodes, edges, cfg
|