Spaces:
Sleeping
Sleeping
File size: 5,989 Bytes
39707cc | 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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | from pydantic import BaseModel
from app.utils.logger import get_logger
logger = get_logger("query.path_extractor")
class GraphNode(BaseModel):
id: str
label: str
properties: dict
class GraphEdge(BaseModel):
source: str
target: str
type: str
properties: dict
class GraphPath(BaseModel):
nodes: list[GraphNode] = []
edges: list[GraphEdge] = []
class PathExtractor:
"""Extracts structured node/edge data from raw Neo4j result records for graph highlighting."""
def extract(self, raw_records: list[dict]) -> GraphPath:
"""Parse Neo4j result records and extract unique nodes and edges."""
nodes_map: dict[str, GraphNode] = {}
edges_list: list[GraphEdge] = []
seen_edges: set[tuple[str, str, str]] = set()
for record in raw_records:
for value in record.values():
self._process_value(value, nodes_map, edges_list, seen_edges)
return GraphPath(
nodes=list(nodes_map.values()),
edges=edges_list,
)
def _process_value(
self,
value,
nodes_map: dict[str, GraphNode],
edges_list: list[GraphEdge],
seen_edges: set[tuple[str, str, str]],
) -> None:
"""Process a single value from a Neo4j record."""
if value is None:
return
# Neo4j driver Node object
if hasattr(value, "labels") and hasattr(value, "element_id"):
self._extract_node(value, nodes_map)
return
# Neo4j driver Relationship object
if hasattr(value, "type") and hasattr(value, "start_node") and hasattr(value, "end_node"):
self._extract_edge(value, nodes_map, edges_list, seen_edges)
return
# Neo4j driver Path object
if hasattr(value, "nodes") and hasattr(value, "relationships"):
for n in value.nodes:
self._extract_node(n, nodes_map)
for r in value.relationships:
self._extract_edge(r, nodes_map, edges_list, seen_edges)
return
# Neo4j Node (dict with _id, _labels, _properties or element_id)
if isinstance(value, dict):
if "_labels" in value or "labels" in value:
self._extract_node(value, nodes_map)
elif "_type" in value or "type" in value:
self._extract_edge(value, nodes_map, edges_list, seen_edges)
# Path objects from Neo4j driver come as dicts with nodes/relationships
elif "nodes" in value and "relationships" in value:
for n in value["nodes"]:
self._extract_node(n, nodes_map)
for r in value["relationships"]:
self._extract_edge(r, nodes_map, edges_list, seen_edges)
# List of values (e.g., from COLLECT or path results)
elif isinstance(value, list):
for item in value:
self._process_value(item, nodes_map, edges_list, seen_edges)
def _extract_node(self, data, nodes_map: dict[str, GraphNode]) -> None:
"""Extract a GraphNode from a Neo4j node dict."""
if isinstance(data, dict):
node_id = str(
data.get("element_id")
or data.get("_id")
or data.get("id")
or id(data)
)
else:
node_id = str(getattr(data, "element_id", None) or id(data))
if node_id in nodes_map:
return
if isinstance(data, dict):
labels = data.get("_labels") or data.get("labels") or []
props = data.get("_properties") or data.get("properties") or {}
else:
labels = list(getattr(data, "labels", []) or [])
props = dict(getattr(data, "_properties", {}) or {})
label = labels[0] if labels else "Unknown"
# If neither _properties nor properties exists, treat remaining keys as properties
if not props and isinstance(data, dict):
skip_keys = {"_id", "_labels", "labels", "element_id", "id", "_properties", "properties"}
props = {k: v for k, v in data.items() if k not in skip_keys}
nodes_map[node_id] = GraphNode(id=node_id, label=label, properties=props)
def _extract_edge(
self,
data,
nodes_map: dict[str, GraphNode],
edges_list: list[GraphEdge],
seen_edges: set[tuple[str, str, str]],
) -> None:
"""Extract a GraphEdge from a Neo4j relationship dict."""
if isinstance(data, dict):
rel_type = data.get("_type") or data.get("type") or "UNKNOWN"
start_id = str(
data.get("_start_node_element_id")
or data.get("start_node_element_id")
or data.get("source")
or ""
)
end_id = str(
data.get("_end_node_element_id")
or data.get("end_node_element_id")
or data.get("target")
or ""
)
props = data.get("_properties") or data.get("properties") or {}
else:
rel_type = str(getattr(data, "type", "UNKNOWN"))
start_node = getattr(data, "start_node", None)
end_node = getattr(data, "end_node", None)
start_id = str(getattr(start_node, "element_id", "") or "")
end_id = str(getattr(end_node, "element_id", "") or "")
props = dict(getattr(data, "_properties", {}) or {})
if start_node is not None:
self._extract_node(start_node, nodes_map)
if end_node is not None:
self._extract_node(end_node, nodes_map)
key = (start_id, end_id, rel_type)
if key in seen_edges:
return
seen_edges.add(key)
edges_list.append(GraphEdge(
source=start_id,
target=end_id,
type=rel_type,
properties=props,
))
|