Spaces:
Sleeping
Sleeping
| """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 | |