Spaces:
Sleeping
Sleeping
| 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, | |
| )) | |