new-graph-visualization / graph_backend.py
anurag-raapid's picture
Upload app, Docker, staging parquet only (no venv)
4b988e0 verified
Raw
History Blame Contribute Delete
48.4 kB
"""Shared UMLS graph backend (FalkorDB + subgraph builder). Used by app.py."""
import os
import threading
import time
from collections import deque
from itertools import combinations
from pathlib import Path
import pandas as pd
from falkordb import FalkorDB
from tui_definitions import format_tui_list, tui_codes, tui_display
# ─────────────────────────────────────────────────────────────────────────────
# FalkorDB connection
# ─────────────────────────────────────────────────────────────────────────────
FALKORDB_HOST = os.environ.get("FALKORDB_HOST", "localhost")
FALKORDB_PORT = int(os.environ.get("FALKORDB_PORT", "6379"))
FALKORDB_GRAPH = os.environ.get("FALKORDB_GRAPH", "clinical_graph")
DATA_DIR = Path(__file__).resolve().parent
NODES_PATH = DATA_DIR / "graph_nodes.parquet"
# Safety caps (keeps the browser responsive on dense concepts)
MAX_SAME_CUI_AUIS = 14
MAX_MRREL_EDGES = 800
MRREL_CAP_MIN_EDGES = 2_000 # only cap when the subgraph has a very large MRREL set
NODE_WARN_THRESHOLD = 2500
NODE_TRIM_TARGET = 2000
QUERY_TIMEOUT_MS = 120_000
MAX_BFS_AUIS = 15_000
# FalkorDB silently truncates any query result to RESULTSET_SIZE rows (default
# 10,000). Fetching metadata/edges for the whole BFS frontier (up to MAX_BFS_AUIS)
# therefore drops rows — including the seed's own neighbours — which disconnects
# the graph. We never display more than ~NODE_TRIM_TARGET nodes, so we pick the
# lowest-hop AUIs up to MAX_DISPLAY_AUIS *before* fetching, and chunk every
# id-keyed query below FETCH_CHUNK so no single query can be truncated.
MAX_DISPLAY_AUIS = NODE_TRIM_TARGET
FETCH_CHUNK = 5_000
# Path finder (AUI ↔ AUI via HIER_ISA + MRREL, undirected)
MAX_PATH_ENUM = 50
MAX_PATH_SEARCH_HOPS = 25
PATH_EXTRA_HOPS = 4
MAX_PATH_BFS_VISITED = 12_000
MAX_PATH_BFS_FRONTIER = 3_000
MAX_PATH_ADJ_NODES = 8_000
# FalkorDB schema: Concept nodes (aui, cui, str, sab, tui, tty); HIER_ISA + MRREL edges
REL_HIER = "HIER_ISA"
REL_MRREL = "MRREL"
NODE_LABEL = "Concept"
def _truncate(text: str, n: int) -> str:
return (text[: n - 1] + "…") if len(text) > n else text
def _aui_node_label(aui: str, s: str) -> str:
short = _truncate(s, 36)
return f"{aui}\n{short}"
def _cui_node_label(cui: str, _s: str = "") -> str:
return cui
def _prop_str(val) -> str:
if val is None:
return ""
return str(val).strip()
def _edge_props(props) -> dict:
if not props:
return {}
if isinstance(props, dict):
return props
return dict(props)
def _clamp_depth(depth: int) -> int:
return max(1, min(15, int(depth or 1)))
class _LockedGraph:
"""Serialize FalkorDB access (Dash may invoke callbacks concurrently)."""
def __init__(self, graph):
self._graph = graph
self._lock = threading.Lock()
def query(
self, cypher: str, params: dict | None = None, *, timeout: int | None = None
):
with self._lock:
return self._graph.query(cypher, params=params or {}, timeout=timeout)
print(
f"▶ Connecting to FalkorDB ({FALKORDB_HOST}:{FALKORDB_PORT}, graph={FALKORDB_GRAPH}) …",
flush=True,
)
_t = time.time()
_fdb = FalkorDB(host=FALKORDB_HOST, port=FALKORDB_PORT)
GRAPH = _LockedGraph(_fdb.select_graph(FALKORDB_GRAPH))
# Smoke test
_probe = GRAPH.query(
f"MATCH (n:{NODE_LABEL}) RETURN count(n) AS c LIMIT 1", timeout=QUERY_TIMEOUT_MS
)
_node_count = _probe.result_set[0][0] if _probe.result_set else "?"
print(f" {FALKORDB_GRAPH}: {_node_count:,} Concept nodes ({time.time() - _t:.1f}s)")
# Warm index/cache so the first user explore is less likely to time out.
_warm = GRAPH.query(
f"MATCH (n:{NODE_LABEL}) WHERE n.aui IS NOT NULL RETURN n.aui LIMIT 1",
timeout=QUERY_TIMEOUT_MS,
)
if _warm.result_set:
print(" index warm-up ok", flush=True)
# Category lives in parquet only (not stored in FalkorDB Concept nodes).
_CATEGORY_BY_AUI: dict[str, str] = {}
def _load_category_lookup() -> None:
global _CATEGORY_BY_AUI
if not NODES_PATH.exists():
print(" ⚠ graph_nodes.parquet missing — category filter disabled.", flush=True)
_CATEGORY_BY_AUI = {}
return
t0 = time.time()
df = pd.read_parquet(NODES_PATH, columns=["node_id", "category"])
_CATEGORY_BY_AUI = {
str(row.node_id).strip(): str(row.category).strip()
for row in df.itertuples(index=False)
if row.node_id
}
print(
f" category lookup: {len(_CATEGORY_BY_AUI):,} AUIs from parquet ({time.time() - t0:.1f}s)",
flush=True,
)
def _category_for_aui(aui: str) -> str:
return _CATEGORY_BY_AUI.get(aui, "")
_load_category_lookup()
print("▶ FalkorDB ready.", flush=True)
# Legacy alias (app previously used DB for DuckDB)
DB = GRAPH
# ─────────────────────────────────────────────────────────────────────────────
# Cypher helpers
# ─────────────────────────────────────────────────────────────────────────────
def _chunked(seq: list[str], size: int = FETCH_CHUNK):
for i in range(0, len(seq), size):
yield seq[i : i + size]
def _fetch_concepts_by_auis(auis: list[str]) -> list[dict]:
if not auis:
return []
out: list[dict] = []
# Chunk so a single query never returns more than RESULTSET_SIZE rows
# (one row per AUI), which FalkorDB would otherwise silently truncate.
for chunk in _chunked(list(auis)):
rows = GRAPH.query(
f"""
UNWIND $auis AS a
MATCH (n:{NODE_LABEL} {{aui: a}})
RETURN n.aui AS aui, n.cui AS cui, n.str AS str,
n.sab AS sab, n.tui AS tui, n.tty AS tty
ORDER BY aui
""",
params={"auis": chunk},
timeout=QUERY_TIMEOUT_MS,
).result_set
for row in rows:
out.append(
{
"AUI": _prop_str(row[0]),
"CUI": _prop_str(row[1]),
"STR": _prop_str(row[2]),
"SAB": _prop_str(row[3]),
"TUI": _prop_str(row[4]),
"TTY": _prop_str(row[5]),
}
)
return out
def _fetch_concepts_by_cuis(cuis: list[str]) -> list[dict]:
if not cuis:
return []
out: list[dict] = []
# One CUI can map to many AUIs, so chunk the input AND guard against the
# RESULTSET cap by shrinking the chunk if a query comes back maxed out.
for chunk in _chunked(list(cuis), max(1, FETCH_CHUNK // 4)):
rows = GRAPH.query(
f"""
UNWIND $cuis AS c
MATCH (n:{NODE_LABEL} {{cui: c}})
RETURN n.cui AS cui, n.aui AS aui, n.str AS str, n.tui AS tui
ORDER BY cui, aui
""",
params={"cuis": chunk},
timeout=QUERY_TIMEOUT_MS,
).result_set
out.extend(
{
"CUI": _prop_str(row[0]),
"AUI": _prop_str(row[1]),
"STR": _prop_str(row[2]),
"TUI": _prop_str(row[3]),
}
for row in rows
)
return out
def _bfs_auis(
seed_auis: list[str], depth: int
) -> tuple[list[str], dict[str, int], bool]:
"""
Layered BFS over HIER_ISA + MRREL (undirected).
Returns (all AUIs sorted, hop distance per AUI from seed, capped_flag).
Layered expansion avoids FalkorDB's ~10k row cap on variable-length patterns
and yields stable hop ranks for large-graph trimming. Metadata is fetched
separately (after hop-budget selection) so we never pull more rows than we
display, which would otherwise hit the RESULTSET_SIZE truncation cap.
"""
if not seed_auis:
return [], {}, False
depth = _clamp_depth(depth)
hop: dict[str, int] = {}
for a in seed_auis:
hop[a] = 0
frontier = set(seed_auis)
all_auis: set[str] = set(seed_auis)
capped = False
for d in range(1, depth + 1):
if not frontier:
break
if len(all_auis) >= MAX_BFS_AUIS:
capped = True
break
nbrs = _neighbors_one_hop_auis(frontier) - all_auis
if len(all_auis) + len(nbrs) > MAX_BFS_AUIS:
room = MAX_BFS_AUIS - len(all_auis)
nbrs = set(sorted(nbrs)[:room])
capped = True
for a in nbrs:
hop[a] = d
all_auis |= nbrs
frontier = nbrs
if capped:
break
return sorted(all_auis), hop, capped
def _select_auis_by_hop(
all_auis: list[str],
aui_hop: dict[str, int],
seed_auis: list[str],
max_auis: int,
) -> list[str]:
"""
Pick the AUIs closest to the seed (lowest hop) up to *max_auis*.
Keeping the nearest hops guarantees the seed and its immediate neighbours
survive, so the displayed graph stays connected. Selecting before fetching
also keeps every downstream query well under the RESULTSET_SIZE cap.
"""
if len(all_auis) <= max_auis:
return list(all_auis)
seed_set = set(seed_auis)
ordered = sorted(
all_auis,
key=lambda a: (0 if a in seed_set else 1, aui_hop.get(a, 10**9), a),
)
return sorted(ordered[:max_auis])
def fetch_aui_filter_meta(aui_ids: list[str]) -> dict[str, dict]:
"""Return {AUI_<id>: {sab, tty, tuis}} for sidebar filters."""
if not aui_ids:
return {}
rows = _fetch_concepts_by_auis(aui_ids)
meta: dict[str, dict] = {}
for r in rows:
meta[f"AUI_{r['AUI']}"] = {
"sab": r["SAB"],
"tty": r["TTY"],
"tuis": tui_codes(r["TUI"]),
"category": _category_for_aui(r["AUI"]),
}
return meta
# ─────────────────────────────────────────────────────────────────────────────
# Subgraph post-processing (connectivity + size cap)
# ─────────────────────────────────────────────────────────────────────────────
def _split_elements(elements: list[dict]) -> tuple[list[dict], list[dict]]:
nodes = [e for e in elements if "source" not in e["data"]]
edges = [e for e in elements if "source" in e["data"]]
return nodes, edges
def _build_adjacency(edges: list[dict]) -> dict[str, set[str]]:
adj: dict[str, set[str]] = {}
for el in edges:
s, t = el["data"]["source"], el["data"]["target"]
adj.setdefault(s, set()).add(t)
adj.setdefault(t, set()).add(s)
return adj
def _bfs_reachable(
starts: set[str], adj: dict[str, set[str]], universe: set[str]
) -> set[str]:
reachable: set[str] = set()
q: deque[str] = deque()
for s in starts:
if s in universe and s not in reachable:
reachable.add(s)
q.append(s)
while q:
u = q.popleft()
for v in adj.get(u, ()):
if v in universe and v not in reachable:
reachable.add(v)
q.append(v)
return reachable
def _bfs_distances(
starts: set[str], adj: dict[str, set[str]], universe: set[str]
) -> dict[str, int]:
dist: dict[str, int] = {}
q: deque[str] = deque()
for s in starts:
if s in universe and s not in dist:
dist[s] = 0
q.append(s)
while q:
u = q.popleft()
du = dist[u]
for v in adj.get(u, ()):
if v in universe and v not in dist:
dist[v] = du + 1
q.append(v)
return dist
def _seed_node_ids(
seed: str, seed_auis: list[str], seed_type: str, show_cui_layer: bool
) -> set[str]:
if seed_type == "CUI" and show_cui_layer:
return {f"CUI_{seed}"}
if seed_auis:
return {f"AUI_{seed_auis[0]}"}
return set()
def _subset_for_kept_auis(
nodes: list[dict], edges: list[dict], keep_aui_ids: set[str]
) -> tuple[list[dict], list[dict]]:
"""Keep selected AUI nodes and CUI nodes tied to those AUIs (plus seed CUI)."""
keep_cui: set[str] = set()
for n in nodes:
d = n["data"]
if d.get("node_type") == "AUI" and d["id"] in keep_aui_ids:
cui = d.get("cui")
if cui:
keep_cui.add(f"CUI_{cui}")
elif d.get("node_type") == "CUI" and d.get("is_seed"):
keep_cui.add(d["id"])
keep_ids = keep_aui_ids | keep_cui
nodes_out = [n for n in nodes if n["data"]["id"] in keep_ids]
edges_out = [
e
for e in edges
if e["data"]["source"] in keep_ids and e["data"]["target"] in keep_ids
]
return nodes_out, edges_out
def _trim_auis_by_hop_budget(
nodes: list[dict], edges: list[dict], aui_hop: dict[str, int], target: int
) -> tuple[list[dict], list[dict], int, set[str]]:
"""
Add AUIs in ascending hop order until the AUI+CUI subset would exceed *target*.
Ensures a larger BFS depth never keeps fewer AUIs than a smaller depth.
"""
aui_nodes = [n for n in nodes if n["data"].get("node_type") == "AUI"]
aui_nodes.sort(
key=lambda n: (aui_hop.get(n["data"].get("aui"), 10**9), n["data"]["id"])
)
kept_aui: set[str] = set()
best_nodes, best_edges = [], []
max_hop_kept = 0
for n in aui_nodes:
trial_aui = kept_aui | {n["data"]["id"]}
trial_nodes, trial_edges = _subset_for_kept_auis(nodes, edges, trial_aui)
if len(trial_nodes) > target:
break
kept_aui = trial_aui
best_nodes, best_edges = trial_nodes, trial_edges
max_hop_kept = max(max_hop_kept, aui_hop.get(n["data"].get("aui"), 0))
if not kept_aui and aui_nodes:
seed_n = aui_nodes[0]
best_nodes, best_edges = _subset_for_kept_auis(
nodes, edges, {seed_n["data"]["id"]}
)
max_hop_kept = aui_hop.get(seed_n["data"].get("aui"), 0)
return best_nodes, best_edges, max_hop_kept, kept_aui
def _apply_graph_limits(
elements: list[dict],
seed: str,
seed_auis: list[str],
seed_type: str,
show_cui_layer: bool,
aui_hop: dict[str, int] | None = None,
) -> tuple[list[dict], list[str]]:
notes: list[str] = []
pinned_auis: set[str] = set()
nodes, edges = _split_elements(elements)
all_ids = {n["data"]["id"] for n in nodes}
if not all_ids:
return elements, notes
starts = _seed_node_ids(seed, seed_auis, seed_type, show_cui_layer) & all_ids
if not starts and seed_auis:
starts = {f"AUI_{seed_auis[0]}"}
n_nodes = len(nodes)
if n_nodes > NODE_WARN_THRESHOLD and aui_hop:
best_nodes, best_edges, max_hop_kept, pinned_aui = _trim_auis_by_hop_budget(
nodes, edges, aui_hop, NODE_TRIM_TARGET
)
pinned_auis = pinned_aui
trimmed = n_nodes - len(best_nodes)
n_aui_kept = sum(1 for n in best_nodes if n["data"].get("node_type") == "AUI")
notes.append(
f"⚠ {n_nodes:,} nodes exceeds {NODE_WARN_THRESHOLD:,} — trimmed to "
f"{len(best_nodes):,} ({n_aui_kept:,} AUIs, farthest hop kept: {max_hop_kept}, "
f"removed {trimmed:,})"
)
nodes, edges = best_nodes, best_edges
all_ids = {n["data"]["id"] for n in nodes}
n_nodes = len(nodes)
elif n_nodes > NODE_WARN_THRESHOLD:
adj = _build_adjacency(edges)
dist = _bfs_distances(starts, adj, all_ids)
order_remove = sorted(
nodes,
key=lambda n: (-dist.get(n["data"]["id"], 10**9), n["data"]["id"]),
)
keep_ids = {n["data"]["id"] for n in nodes}
for n in order_remove:
if len(keep_ids) <= NODE_TRIM_TARGET:
break
keep_ids.discard(n["data"]["id"])
trimmed = n_nodes - len(keep_ids)
notes.append(
f"⚠ {n_nodes:,} nodes exceeds {NODE_WARN_THRESHOLD:,} — trimmed to "
f"{len(keep_ids):,} (removed {trimmed:,} farthest from seed)"
)
nodes = [n for n in nodes if n["data"]["id"] in keep_ids]
edges = [
e
for e in edges
if e["data"]["source"] in keep_ids and e["data"]["target"] in keep_ids
]
all_ids = keep_ids
adj = _build_adjacency(edges)
reachable2 = _bfs_reachable(starts, adj, all_ids)
if len(reachable2) < len(all_ids):
extra = len(all_ids) - len(reachable2)
notes.append(
f"removed {extra:,} node{'s' if extra != 1 else ''} after trim (no path to seed)"
)
nodes = [n for n in nodes if n["data"]["id"] in reachable2]
edges = [
e
for e in edges
if e["data"]["source"] in reachable2
and e["data"]["target"] in reachable2
]
all_ids = {n["data"]["id"] for n in nodes}
adj = _build_adjacency(edges)
reachable = _bfs_reachable(starts, adj, all_ids)
if pinned_auis:
reachable |= pinned_auis
dropped = len(all_ids) - len(reachable)
if dropped:
notes.append(
f"removed {dropped:,} disconnected node{'s' if dropped != 1 else ''}"
)
nodes = [n for n in nodes if n["data"]["id"] in reachable]
edges = [
e
for e in edges
if e["data"]["source"] in reachable and e["data"]["target"] in reachable
]
return nodes + edges, notes
def _seed_starts_from_nodes(nodes: list[dict]) -> set[str]:
cui_seeds = sorted(
n["data"]["id"]
for n in nodes
if n["data"].get("is_seed") and n["data"].get("node_type") == "CUI"
)
if cui_seeds:
return {cui_seeds[0]}
aui_seeds = sorted(
n["data"]["id"]
for n in nodes
if n["data"].get("is_seed") and n["data"].get("node_type") == "AUI"
)
if aui_seeds:
return {aui_seeds[0]}
legacy = sorted(
n["data"]["id"] for n in nodes if "seed-node" in (n.get("classes") or "")
)
return {legacy[0]} if legacy else set()
def filter_to_seed_component(
nodes: list[dict], edges: list[dict]
) -> tuple[list[dict], list[dict]]:
node_ids = {n["data"]["id"] for n in nodes}
if not node_ids:
return nodes, edges
starts = _seed_starts_from_nodes(nodes) & node_ids
if not starts:
aui_seeds = sorted(
n["data"]["id"]
for n in nodes
if n["data"].get("is_seed") and n["data"].get("node_type") == "AUI"
)
if aui_seeds:
starts = {aui_seeds[0]}
else:
return nodes, edges
adj = _build_adjacency(edges)
reachable = _bfs_reachable(starts, adj, node_ids)
nodes_out = [n for n in nodes if n["data"]["id"] in reachable]
edges_out = [
e
for e in edges
if e["data"]["source"] in reachable and e["data"]["target"] in reachable
]
return nodes_out, edges_out
def _neighbors_one_hop_auis(frontier: set[str], *, directed: bool = False) -> set[str]:
if not frontier:
return set()
frontier_list = list(frontier)
if directed:
rows = GRAPH.query(
f"""
MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]->(m:{NODE_LABEL})
WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
RETURN DISTINCT m.aui AS aui
""",
params={"frontier": frontier_list},
timeout=QUERY_TIMEOUT_MS,
).result_set
else:
rows = GRAPH.query(
f"""
MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]-(m:{NODE_LABEL})
WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
RETURN DISTINCT m.aui AS aui
""",
params={"frontier": frontier_list},
timeout=QUERY_TIMEOUT_MS,
).result_set
return {_prop_str(row[0]) for row in rows if row[0]}
def _neighbors_one_hop_predecessors(frontier: set[str]) -> set[str]:
if not frontier:
return set()
rows = GRAPH.query(
f"""
MATCH (m:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]->(n:{NODE_LABEL})
WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
RETURN DISTINCT m.aui AS aui
""",
params={"frontier": list(frontier)},
).result_set
return {_prop_str(row[0]) for row in rows if row[0]}
def _neighbors_undirected_edges(frontier: set[str]) -> list[tuple[str, str]]:
if not frontier:
return []
rows = GRAPH.query(
f"""
MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]-(m:{NODE_LABEL})
WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
RETURN DISTINCT n.aui AS u, m.aui AS v
""",
params={"frontier": list(frontier)},
).result_set
out: list[tuple[str, str]] = []
for row in rows:
u, v = _prop_str(row[0]), _prop_str(row[1])
if u and v:
out.append((u, v) if u <= v else (v, u))
return out
def _directed_edges_forward(frontier: set[str]) -> list[tuple[str, str]]:
if not frontier:
return []
rows = GRAPH.query(
f"""
MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]->(m:{NODE_LABEL})
WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
RETURN DISTINCT n.aui AS u, m.aui AS v
""",
params={"frontier": list(frontier)},
).result_set
return [(_prop_str(row[0]), _prop_str(row[1])) for row in rows if row[0] and row[1]]
def _aui_exists(aui: str) -> bool:
hit = GRAPH.query(
f"MATCH (n:{NODE_LABEL} {{aui: $a}}) RETURN n.aui LIMIT 1",
params={"a": aui},
).result_set
return bool(hit)
def _resolve_seed(seed: str) -> tuple[list[str], str] | tuple[None, None]:
"""Return (seed_auis, seed_type) or (None, None) if not found."""
hit = GRAPH.query(
f"MATCH (n:{NODE_LABEL} {{aui: $s}}) RETURN n.aui LIMIT 1",
params={"s": seed},
).result_set
if hit:
return [seed], "AUI"
rows = GRAPH.query(
f"MATCH (n:{NODE_LABEL} {{cui: $s}}) RETURN DISTINCT n.aui AS aui ORDER BY aui",
params={"s": seed},
).result_set
if rows:
return [_prop_str(row[0]) for row in rows if row[0]], "CUI"
return None, None
def _shortest_hop_distance(
aui_start: str, aui_end: str, max_hops: int, *, directed: bool = False
) -> int | None:
if aui_start == aui_end:
return 0
dist_fwd: dict[str, int] = {aui_start: 0}
dist_bwd: dict[str, int] = {aui_end: 0}
front_fwd: set[str] = {aui_start}
front_bwd: set[str] = {aui_end}
for step in range(1, max_hops + 1):
if len(dist_fwd) + len(dist_bwd) > MAX_PATH_BFS_VISITED:
return None
meet = set(dist_fwd) & set(dist_bwd)
if meet:
return min(dist_fwd[m] + dist_bwd[m] for m in meet)
if len(front_fwd) <= len(front_bwd):
if not front_fwd:
break
nbrs = _neighbors_one_hop_auis(front_fwd, directed=directed)
next_fwd: set[str] = set()
for v in nbrs:
if v not in dist_fwd:
dist_fwd[v] = step
next_fwd.add(v)
if len(dist_fwd) > MAX_PATH_BFS_VISITED:
return None
if len(next_fwd) > MAX_PATH_BFS_FRONTIER:
return None
front_fwd = next_fwd
else:
if not front_bwd:
break
nbrs = (
_neighbors_one_hop_predecessors(front_bwd)
if directed
else _neighbors_one_hop_auis(front_bwd)
)
next_bwd: set[str] = set()
for v in nbrs:
if v not in dist_bwd:
dist_bwd[v] = step
next_bwd.add(v)
if len(dist_bwd) > MAX_PATH_BFS_VISITED:
return None
if len(next_bwd) > MAX_PATH_BFS_FRONTIER:
return None
front_bwd = next_bwd
return None
def _build_local_adjacency(
seed_auis: set[str], max_hops: int, *, directed: bool = False
) -> tuple[set[str], dict[str, set[str]]]:
visited: set[str] = set(seed_auis)
adj: dict[str, set[str]] = {a: set() for a in seed_auis}
frontier = set(seed_auis)
for _ in range(max_hops):
if not frontier or len(visited) >= MAX_PATH_BFS_VISITED:
break
new_frontier: set[str] = set()
if directed:
for u, v in _directed_edges_forward(frontier):
if len(visited) >= MAX_PATH_BFS_VISITED:
break
adj.setdefault(u, set()).add(v)
adj.setdefault(v, set())
if v not in visited:
visited.add(v)
new_frontier.add(v)
else:
for u, v in _neighbors_undirected_edges(frontier):
if len(visited) >= MAX_PATH_BFS_VISITED:
break
adj.setdefault(u, set()).add(v)
adj.setdefault(v, set()).add(u)
for n in (u, v):
if n not in visited:
visited.add(n)
adj.setdefault(n, set())
new_frontier.add(n)
if len(visited) >= MAX_PATH_BFS_VISITED:
break
if (
len(visited) >= MAX_PATH_BFS_VISITED
or len(new_frontier) > MAX_PATH_BFS_FRONTIER
):
break
frontier = new_frontier
return visited, adj
def _enumerate_simple_paths(
aui_start: str,
aui_end: str,
adj: dict[str, set[str]],
max_depth: int,
max_paths: int,
) -> list[list[str]]:
paths: list[list[str]] = []
def dfs(curr: str, visited: set[str], trail: list[str]) -> None:
if len(paths) >= max_paths:
return
if curr == aui_end:
paths.append(trail[:])
return
if len(trail) >= max_depth:
return
for nxt in sorted(adj.get(curr, ())):
if nxt not in visited:
visited.add(nxt)
trail.append(nxt)
dfs(nxt, visited, trail)
trail.pop()
visited.discard(nxt)
if len(paths) >= max_paths:
return
dfs(aui_start, {aui_start}, [aui_start])
return paths
def find_paths_between_auis(
aui_start: str,
aui_end: str,
*,
max_paths: int = MAX_PATH_ENUM,
max_hops: int = MAX_PATH_SEARCH_HOPS,
directed: bool = False,
) -> tuple[list[list[str]], str]:
aui_start = aui_start.strip().upper()
aui_end = aui_end.strip().upper()
mode = "directed" if directed else "undirected"
if not _aui_exists(aui_start):
return [], f"AUI '{aui_start}' not found."
if not _aui_exists(aui_end):
return [], f"AUI '{aui_end}' not found."
if aui_start == aui_end:
return [[aui_start]], f"Same start and end AUI (trivial path, {mode})."
shortest = _shortest_hop_distance(aui_start, aui_end, max_hops, directed=directed)
if shortest is None:
return (
[],
f"No {mode} path from {aui_start} to {aui_end} within {max_hops} hops "
f"(HIER_ISA + MRREL).",
)
search_depth = min(shortest + PATH_EXTRA_HOPS, max_hops)
_, adj = _build_local_adjacency({aui_start}, search_depth, directed=directed)
if aui_end not in adj:
_, adj = _build_local_adjacency(
{aui_start, aui_end}, search_depth + 2, directed=directed
)
if len(adj) > MAX_PATH_ADJ_NODES:
return (
[],
f"Path search neighborhood too large ({len(adj):,} AUIs near "
f"{aui_start}). Pick AUIs farther from high-degree hubs, use directed "
f"mode, or reduce hop depth.",
)
paths = _enumerate_simple_paths(aui_start, aui_end, adj, search_depth, max_paths)
if not paths:
return (
[],
f"Reachable in {mode} graph (shortest ≈ {shortest} hops) but no simple path "
f"found within {search_depth} hops (enumeration limit).",
)
capped = len(paths) >= max_paths
hop_lens = [len(p) - 1 for p in paths]
msg = (
f"Found {len(paths)} {mode} path{'s' if len(paths) != 1 else ''} "
f"({min(hop_lens)}{max(hop_lens)} hops, shortest={shortest})"
)
if capped:
msg += f" — showing first {max_paths} (cap)"
if search_depth > shortest:
msg += f" · searched up to {search_depth} hops"
return paths, msg
def _path_nav_neighbor_map(paths: list[list[str]]) -> dict[str, set[str]]:
nav: dict[str, set[str]] = {}
for path in paths:
for i, aui in enumerate(path):
nid = f"AUI_{aui}"
nav.setdefault(nid, set())
if i > 0:
nav[nid].add(f"AUI_{path[i - 1]}")
if i < len(path) - 1:
nav[nid].add(f"AUI_{path[i + 1]}")
return nav
def _path_edge_pairs(
paths: list[list[str]], *, directed: bool = False
) -> set[tuple[str, str]]:
pairs: set[tuple[str, str]] = set()
for path in paths:
for i in range(len(path) - 1):
a, b = path[i], path[i + 1]
if directed:
pairs.add((a, b))
else:
pairs.add((a, b) if a <= b else (b, a))
return pairs
def _fetch_hier_edges(auis: list[str]) -> list[dict]:
if not auis:
return []
rows = GRAPH.query(
f"""
MATCH (x:{NODE_LABEL})-[r:{REL_HIER}]->(y:{NODE_LABEL})
WHERE x.aui IN $auis AND y.aui IN $auis
RETURN x.aui AS u, y.aui AS v, r.sab AS sab
""",
params={"auis": auis},
).result_set
return [
{
"u": _prop_str(row[0]),
"v": _prop_str(row[1]),
"sab": _prop_str(row[2]),
"rela": "isa",
}
for row in rows
]
def _fetch_mrrel_edges(auis: list[str], limit: int | None = None) -> list[dict]:
if not auis:
return []
lim = f"LIMIT {int(limit)}" if limit else ""
rows = GRAPH.query(
f"""
MATCH (x:{NODE_LABEL})-[r:{REL_MRREL}]->(y:{NODE_LABEL})
WHERE x.aui IN $auis AND y.aui IN $auis AND x.aui <> y.aui
RETURN x.aui AS u, y.aui AS v, r.sab AS sab, r.rela AS rela, r.rel AS rel
{lim}
""",
params={"auis": auis},
).result_set
return [
{
"u": _prop_str(row[0]),
"v": _prop_str(row[1]),
"sab": _prop_str(row[2]),
"rela": _prop_str(row[3]) or _prop_str(row[4]) or "rel",
"rel": _prop_str(row[4]),
}
for row in rows
]
def _fetch_path_edges(
pairs: set[tuple[str, str]], *, directed: bool = False
) -> list[dict]:
if not pairs:
return []
pair_list = [{"u": u, "v": v} for u, v in sorted(pairs)]
if directed:
hier_q = f"""
UNWIND $pairs AS p
MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_HIER}]->(y:{NODE_LABEL} {{aui: p.v}})
RETURN 'mrhier' AS et, x.aui AS u, y.aui AS v, r.sab AS sab, 'isa' AS rela, '' AS rel
"""
mrrel_q = f"""
UNWIND $pairs AS p
MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_MRREL}]->(y:{NODE_LABEL} {{aui: p.v}})
RETURN 'mrrel' AS et, x.aui AS u, y.aui AS v, r.sab AS sab,
coalesce(r.rela, r.rel, 'rel') AS rela, coalesce(r.rel, '') AS rel
"""
else:
hier_q = f"""
UNWIND $pairs AS p
MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_HIER}]-(y:{NODE_LABEL} {{aui: p.v}})
RETURN 'mrhier' AS et, x.aui AS u, y.aui AS v, r.sab AS sab, 'isa' AS rela, '' AS rel
"""
mrrel_q = f"""
UNWIND $pairs AS p
MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_MRREL}]-(y:{NODE_LABEL} {{aui: p.v}})
RETURN 'mrrel' AS et, x.aui AS u, y.aui AS v, r.sab AS sab,
coalesce(r.rela, r.rel, 'rel') AS rela, coalesce(r.rel, '') AS rel
"""
rows = GRAPH.query(hier_q, params={"pairs": pair_list}).result_set
rows += GRAPH.query(mrrel_q, params={"pairs": pair_list}).result_set
out: list[dict] = []
for row in rows:
out.append(
{
"et": _prop_str(row[0]),
"u": _prop_str(row[1]),
"v": _prop_str(row[2]),
"sab": _prop_str(row[3]),
"rela": _prop_str(row[4]),
"rel": _prop_str(row[5]),
}
)
return out
def build_path_subgraph(
aui_start: str,
aui_end: str,
paths: list[list[str]],
*,
directed: bool = False,
) -> tuple[list[dict], str]:
if not paths:
return [], "No paths to display."
aui_start = aui_start.strip().upper()
aui_end = aui_end.strip().upper()
all_auis: set[str] = set()
for p in paths:
all_auis.update(p)
meta_rows = _fetch_concepts_by_auis(list(all_auis))
meta_by_aui = {r["AUI"]: r for r in meta_rows}
elements: list[dict] = []
aui_node_ids: set[str] = set()
path_pairs = _path_edge_pairs(paths, directed=directed)
for aui in sorted(all_auis):
r = meta_by_aui.get(aui, {})
nid = f"AUI_{aui}"
aui_node_ids.add(nid)
is_start = aui == aui_start
is_end = aui == aui_end
classes = "aui-node"
if is_start:
classes += " path-endpoint-start"
elif is_end:
classes += " path-endpoint-end"
else:
classes += " path-node"
if is_start or is_end:
classes += " seed-node"
elements.append(
{
"data": {
"id": nid,
"label": _aui_node_label(aui, r.get("STR", "") or ""),
"node_type": "AUI",
"aui": aui,
"cui": r.get("CUI", "") or "",
"str": r.get("STR", "") or "",
"sab": r.get("SAB", "") or "",
"tui": r.get("TUI", "") or "",
"tui_label": format_tui_list(r.get("TUI", "") or ""),
"tty": r.get("TTY", "") or "",
"category": _category_for_aui(aui),
"is_seed": is_start or is_end,
"path_role": "start" if is_start else ("end" if is_end else "via"),
},
"classes": classes,
}
)
nav_map = _path_nav_neighbor_map(paths)
if not path_pairs:
for el in elements:
if "source" in el["data"]:
continue
nid = el["data"].get("id")
if not nid:
continue
step_nbrs = sorted(nav_map.get(nid, ()))
if step_nbrs:
el["data"]["nav_neighbor_ids"] = step_nbrs
n_edges = 0
msg = (
f"Path view: {aui_start}{aui_end} | "
f"{len(paths)} path(s) | {len(all_auis)} AUIs | {n_edges} edges"
)
return elements, msg
edge_rows = _fetch_path_edges(path_pairs, directed=directed)
seen_edge: set[tuple[str, str, str]] = set()
for r in edge_rows:
u, v = r["u"], r["v"]
et = r["et"]
key = (et, u, v)
if key in seen_edge:
continue
seen_edge.add(key)
src, tgt = f"AUI_{u}", f"AUI_{v}"
if src not in aui_node_ids or tgt not in aui_node_ids:
continue
if et == "mrhier":
elements.append(
{
"data": {
"id": f"mrhier_{u}_{v}",
"source": src,
"target": tgt,
"edge_type": "mrhier",
"label": r["rela"],
"sab": r["sab"],
"on_path": True,
},
"classes": "mrhier-edge path-edge",
}
)
else:
rela = str(r["rela"] or "rel")
label = rela if len(rela) <= 22 else rela[:20] + "…"
elements.append(
{
"data": {
"id": f"mrrel_{u}_{v}_{r['rel']}_{rela}",
"source": src,
"target": tgt,
"edge_type": "mrrel",
"label": label,
"rela": rela,
"rel": r["rel"] or "",
"sab": r["sab"] or "",
"on_path": True,
},
"classes": "mrrel-edge path-edge",
}
)
for el in elements:
if "source" in el["data"]:
continue
nid = el["data"].get("id")
if not nid:
continue
step_nbrs = sorted(nav_map.get(nid, ()))
if step_nbrs:
el["data"]["nav_neighbor_ids"] = step_nbrs
n_edges = sum(1 for e in elements if "source" in e["data"])
msg = (
f"Path view: {aui_start}{aui_end} | "
f"{len(paths)} path(s) | {len(all_auis)} AUIs | {n_edges} edges"
)
return elements, msg
def build_subgraph(
seed: str,
depth: int,
show_mrhier: bool,
show_aui_cui: bool,
show_same_cui: bool,
show_cui_cui: bool,
show_mrrel: bool,
show_cui_layer: bool = True,
) -> tuple[list[dict], str]:
del show_cui_cui # CUI parent hierarchy not in this dataset
seed = seed.strip().upper()
depth = _clamp_depth(depth)
resolved = _resolve_seed(seed)
if resolved[0] is None:
return [], f"'{seed}' not found as AUI or CUI."
seed_auis, seed_type = resolved
all_auis, aui_hop, bfs_capped = _bfs_auis(seed_auis, depth)
if not all_auis:
return [], "No metadata found for the resolved AUIs."
# Select the display set (nearest hops) BEFORE fetching metadata/edges so we
# never exceed FalkorDB's RESULTSET_SIZE cap and keep the seed connected.
kept_auis = _select_auis_by_hop(all_auis, aui_hop, seed_auis, MAX_DISPLAY_AUIS)
prefetch_trimmed = len(all_auis) - len(kept_auis)
meta_rows = _fetch_concepts_by_auis(kept_auis)
if not meta_rows:
return [], "No metadata found for the resolved AUIs."
seed_set: set[str] = set(seed_auis)
elements: list[dict] = []
aui_node_ids: set[str] = set()
cui_node_ids: set[str] = set()
aui_to_cui: dict[str, str] = {}
visited_auis = [r["AUI"] for r in meta_rows]
for r in meta_rows:
aui = r["AUI"]
cui = r["CUI"] or ""
s = r["STR"] or ""
sab = r["SAB"] or ""
tui = r["TUI"] or ""
tty = r["TTY"] or ""
aui_to_cui[aui] = cui
nid = f"AUI_{aui}"
aui_node_ids.add(nid)
is_seed = aui in seed_set
tui_fmt = format_tui_list(tui)
elements.append(
{
"data": {
"id": nid,
"label": _aui_node_label(aui, s),
"node_type": "AUI",
"aui": aui,
"cui": cui,
"str": s,
"sab": sab,
"tui": tui,
"tui_label": tui_fmt,
"tty": tty,
"category": _category_for_aui(aui),
"is_seed": is_seed,
},
"classes": "aui-node" + (" seed-node" if is_seed else ""),
}
)
all_cuis: set[str] = {c for c in aui_to_cui.values() if c}
cui_label: dict[str, tuple[str, str]] = {}
if show_cui_layer and all_cuis:
lbl_rows = _fetch_concepts_by_cuis(list(all_cuis))
cui_buckets: dict[str, list[dict]] = {}
for row in lbl_rows:
cui_buckets.setdefault(row["CUI"], []).append(row)
for cui, bucket in cui_buckets.items():
pref = sorted(bucket, key=lambda x: x["AUI"])[0]
cui_label[cui] = (pref["STR"] or "", pref["TUI"] or "")
for cui in all_cuis:
nid = f"CUI_{cui}"
cui_node_ids.add(nid)
raw, tui_raw = cui_label.get(cui, ("", ""))
tui_fmt = format_tui_list(tui_raw)
is_seed_cui = seed_type == "CUI" and cui == seed
elements.append(
{
"data": {
"id": nid,
"label": _cui_node_label(cui),
"node_type": "CUI",
"cui": cui,
"str": raw,
"tui": tui_raw,
"tui_label": tui_fmt,
"is_seed": is_seed_cui,
},
"classes": "cui-node" + (" seed-node" if is_seed_cui else ""),
}
)
all_ids = aui_node_ids | cui_node_ids
mrrel_truncated = False
if show_mrhier:
for r in _fetch_hier_edges(visited_auis):
src = f"AUI_{r['u']}"
tgt = f"AUI_{r['v']}"
if src in all_ids and tgt in all_ids:
elements.append(
{
"data": {
"id": f"mrhier_{r['u']}_{r['v']}",
"source": src,
"target": tgt,
"edge_type": "mrhier",
"label": r["rela"],
"sab": r["sab"],
},
"classes": "mrhier-edge",
}
)
if show_cui_layer and show_aui_cui:
for aui, cui in aui_to_cui.items():
if not cui:
continue
src = f"AUI_{aui}"
tgt = f"CUI_{cui}"
if src in all_ids and tgt in all_ids:
elements.append(
{
"data": {
"id": f"aui_cui_{aui}",
"source": src,
"target": tgt,
"edge_type": "aui_cui",
"label": "belongs_to",
},
"classes": "aui-cui-edge",
}
)
if show_cui_layer and show_same_cui:
cui_buckets: dict[str, list[str]] = {}
for aui, cui in aui_to_cui.items():
if cui:
cui_buckets.setdefault(cui, []).append(aui)
for _cui, auis in cui_buckets.items():
if len(auis) > MAX_SAME_CUI_AUIS:
continue
if len(auis) > 1:
for a, b in combinations(sorted(auis), 2):
elements.append(
{
"data": {
"id": f"same_cui_{a}_{b}",
"source": f"AUI_{a}",
"target": f"AUI_{b}",
"edge_type": "same_cui",
"label": "same_cui",
},
"classes": "same-cui-edge",
}
)
if show_mrrel:
re = _fetch_mrrel_edges(visited_auis, limit=None)
if len(re) > MAX_MRREL_EDGES and len(re) >= MRREL_CAP_MIN_EDGES:
mrrel_truncated = True
re = sorted(
re,
key=lambda r: aui_hop.get(r["u"], 10**9) + aui_hop.get(r["v"], 10**9),
)[:MAX_MRREL_EDGES]
for r in re:
src = f"AUI_{r['u']}"
tgt = f"AUI_{r['v']}"
if src in all_ids and tgt in all_ids:
rela = str(r["rela"] or "rel")
label = rela if len(rela) <= 22 else rela[:20] + "…"
elements.append(
{
"data": {
"id": f"mrrel_{r['u']}_{r['v']}_{r['rel']}_{rela}",
"source": src,
"target": tgt,
"edge_type": "mrrel",
"label": label,
"rela": rela,
"rel": r["rel"] or "",
"sab": r["sab"] or "",
},
"classes": "mrrel-edge",
}
)
elements, limit_notes = _apply_graph_limits(
elements,
seed,
seed_auis,
seed_type,
show_cui_layer,
aui_hop=aui_hop,
)
# Connectivity is enforced inside _apply_graph_limits (hop-trimmed AUIs are pinned).
n_aui = sum(1 for e in elements if e["data"].get("node_type") == "AUI")
n_cui = sum(1 for e in elements if e["data"].get("node_type") == "CUI")
n_edges = sum(1 for e in elements if "source" in e["data"])
extra: list[str] = []
if bfs_capped:
extra.append(f"AUI neighborhood capped at {MAX_BFS_AUIS:,} atoms")
if prefetch_trimmed > 0:
extra.append(
f"⚠ {len(all_auis):,} AUIs in range — kept {len(kept_auis):,} nearest "
f"(hop-priority) before fetch; raise hop depth narrows scope"
)
if mrrel_truncated:
extra.append(f"MRREL capped at {MAX_MRREL_EDGES} edges (seed-adjacent first)")
extra.extend(limit_notes)
msg = (
f"Seed: {seed} ({seed_type}) | Depth: {depth} | "
f"AUI nodes: {n_aui} | CUI nodes: {n_cui} | Edges: {n_edges}"
)
if extra:
msg += " | " + " · ".join(extra)
return elements, msg
def _extract_filter_options(elements: list[dict]) -> dict:
auis: set[str] = set()
cuis: set[str] = set()
for e in elements:
d = e["data"]
if d.get("node_type") == "AUI" and d.get("aui"):
auis.add(d["aui"])
elif d.get("node_type") == "CUI" and d.get("cui"):
cuis.add(d["cui"])
sabs: set[str] = set()
tuis: set[str] = set()
ttys: set[str] = set()
categories: set[str] = set()
if auis:
for r in _fetch_concepts_by_auis(list(auis)):
if r["SAB"]:
sabs.add(r["SAB"])
if r["TTY"]:
ttys.add(r["TTY"])
tuis |= tui_codes(r["TUI"] or "")
cat = _category_for_aui(r["AUI"])
if cat:
categories.add(cat)
if cuis:
for r in _fetch_concepts_by_cuis(list(cuis)):
if r.get("SAB"):
sabs.add(r["SAB"])
tuis |= tui_codes(r.get("TUI") or "")
return {
"sab": [{"label": s, "value": s} for s in sorted(sabs)],
"tui": [{"label": tui_display(t), "value": t} for t in sorted(tuis)],
"tty": [{"label": t, "value": t} for t in sorted(ttys)],
"category": [{"label": c, "value": c} for c in sorted(categories)],
}
def _aui_passes_filters(
node_id: str,
meta: dict[str, dict],
sab_keep: set[str],
tui_keep: set[str],
tty_keep: set[str],
category_keep: set[str] | None = None,
) -> bool:
m = meta.get(node_id)
if not m:
return True
if sab_keep and m["sab"] not in sab_keep:
return False
if tty_keep and m["tty"] not in tty_keep:
return False
if tui_keep and not (m["tuis"] & tui_keep):
return False
if category_keep and m.get("category") not in category_keep:
return False
return True