Spaces:
Sleeping
Sleeping
| from collections import deque | |
| from pydantic import BaseModel | |
| class NodeSchema(BaseModel): | |
| label: str # PascalCase Neo4j label | |
| source_entity: str | None = None # folder name, e.g. "sales_order_headers" | |
| properties: dict[str, str] # column_name -> pandas dtype string | |
| id_field: str | list[str] # primary identifier column(s) — list for composite keys | |
| def id_fields(self) -> list[str]: | |
| """Always return id_field as a list.""" | |
| if isinstance(self.id_field, list): | |
| return self.id_field | |
| return [self.id_field] | |
| class EdgeSchema(BaseModel): | |
| type: str # UPPER_SNAKE relationship type | |
| from_node: str # source node label | |
| to_node: str # target node label | |
| join_on: tuple[str, str] # (source_field, target_field) | |
| source_entity: str | None = None # bridge table folder name (reads data from this DF) | |
| rel_properties: list[str] | None = None # columns to copy as relationship properties | |
| bridge_target_field: str | None = None # bridge column for target match (when != target id) | |
| class GraphSchema(BaseModel): | |
| nodes: list[NodeSchema] | |
| edges: list[EdgeSchema] | |
| def get_node(self, label: str) -> NodeSchema | None: | |
| for n in self.nodes: | |
| if n.label == label: | |
| return n | |
| return None | |
| def node_labels(self) -> list[str]: | |
| return [n.label for n in self.nodes] | |
| def edge_types(self) -> list[str]: | |
| return [e.type for e in self.edges] | |
| def edge_descriptions(self) -> list[str]: | |
| """Return edges as 'FromNode -[TYPE]-> ToNode' strings, with rel props if present.""" | |
| descs: list[str] = [] | |
| for e in self.edges: | |
| desc = f"{e.from_node} -[{e.type}]-> {e.to_node}" | |
| if e.rel_properties: | |
| desc += f" {{{', '.join(e.rel_properties)}}}" | |
| descs.append(desc) | |
| return descs | |
| def node_properties_summary(self, max_properties_per_node: int | None = None) -> list[str]: | |
| """Return per-node property summaries for LLM context.""" | |
| lines = [] | |
| for n in self.nodes: | |
| id_str = ", ".join(n.id_fields) if isinstance(n.id_field, list) else n.id_field | |
| props = [p for p in n.properties.keys() if p not in n.id_fields] | |
| if max_properties_per_node is not None and len(props) > max_properties_per_node: | |
| props = props[:max_properties_per_node] + ["..."] | |
| lines.append(f"{n.label} (PRIMARY KEY: `{id_str}`): {', '.join(props)}") | |
| return lines | |
| def extract_relevant_subschema(self, seed_labels: list[str]) -> "GraphSchema": | |
| """Extract a relevant schema slice using BFS pathfinding between seed labels.""" | |
| existing_labels = set(self.node_labels()) | |
| seeds = [label for label in dict.fromkeys(seed_labels) if label in existing_labels] | |
| if not seeds: | |
| return self | |
| selected_nodes = set(seeds) | |
| selected_edges: set[int] = set() | |
| if len(seeds) == 1: | |
| self._add_two_hop_edges(seeds[0], selected_nodes, selected_edges) | |
| return self._build_subschema(selected_nodes, selected_edges) | |
| found_path = False | |
| for i, start in enumerate(seeds): | |
| for end in seeds[i + 1:]: | |
| path_edges = self._shortest_path_edges(start, end) | |
| if not path_edges: | |
| continue | |
| found_path = True | |
| for edge_idx in path_edges: | |
| edge = self.edges[edge_idx] | |
| selected_edges.add(edge_idx) | |
| selected_nodes.add(edge.from_node) | |
| selected_nodes.add(edge.to_node) | |
| if not found_path: | |
| for seed in seeds: | |
| self._add_incident_edges(seed, selected_nodes, selected_edges) | |
| return self._build_subschema(selected_nodes, selected_edges) | |
| def _shortest_path_edges(self, start: str, end: str) -> list[int]: | |
| adjacency: dict[str, list[tuple[str, int]]] = {} | |
| for idx, edge in enumerate(self.edges): | |
| adjacency.setdefault(edge.from_node, []).append((edge.to_node, idx)) | |
| adjacency.setdefault(edge.to_node, []).append((edge.from_node, idx)) | |
| queue = deque([start]) | |
| parents: dict[str, tuple[str, int] | None] = {start: None} | |
| while queue: | |
| current = queue.popleft() | |
| if current == end: | |
| break | |
| for neighbor, edge_idx in adjacency.get(current, []): | |
| if neighbor in parents: | |
| continue | |
| parents[neighbor] = (current, edge_idx) | |
| queue.append(neighbor) | |
| if end not in parents: | |
| return [] | |
| path_edges: list[int] = [] | |
| cursor = end | |
| while parents[cursor] is not None: | |
| parent, edge_idx = parents[cursor] | |
| path_edges.append(edge_idx) | |
| cursor = parent | |
| path_edges.reverse() | |
| return path_edges | |
| def _add_incident_edges( | |
| self, | |
| label: str, | |
| selected_nodes: set[str], | |
| selected_edges: set[int], | |
| ) -> None: | |
| for idx, edge in enumerate(self.edges): | |
| if edge.from_node == label or edge.to_node == label: | |
| selected_edges.add(idx) | |
| selected_nodes.add(edge.from_node) | |
| selected_nodes.add(edge.to_node) | |
| def _add_two_hop_edges( | |
| self, | |
| label: str, | |
| selected_nodes: set[str], | |
| selected_edges: set[int], | |
| ) -> None: | |
| """Include incident edges for seed node and its immediate neighbors.""" | |
| self._add_incident_edges(label, selected_nodes, selected_edges) | |
| neighbors = {n for n in selected_nodes if n != label} | |
| for neighbor in neighbors: | |
| self._add_incident_edges(neighbor, selected_nodes, selected_edges) | |
| def _build_subschema(self, selected_nodes: set[str], selected_edges: set[int]) -> "GraphSchema": | |
| nodes = [node for node in self.nodes if node.label in selected_nodes] | |
| edges = [edge for idx, edge in enumerate(self.edges) if idx in selected_edges] | |
| return GraphSchema(nodes=nodes, edges=edges) | |
| def find_candidate_path_patterns(self, seed_labels: list[str]) -> list[str]: | |
| """Find candidate traversal patterns between seed labels using BFS shortest paths.""" | |
| existing_labels = set(self.node_labels()) | |
| seeds = [label for label in dict.fromkeys(seed_labels) if label in existing_labels] | |
| if not seeds: | |
| return [] | |
| patterns: list[str] = [] | |
| if len(seeds) == 1: | |
| return self._two_hop_path_patterns(seeds[0]) | |
| for i, start in enumerate(seeds): | |
| for end in seeds[i + 1:]: | |
| steps = self._shortest_path_steps(start, end) | |
| if not steps: | |
| continue | |
| patterns.append(self._steps_to_pattern(start, steps)) | |
| if patterns: | |
| return list(dict.fromkeys(patterns)) | |
| for seed in seeds: | |
| patterns.extend(self._incident_path_patterns(seed)) | |
| return list(dict.fromkeys(patterns)) | |
| def _shortest_path_steps(self, start: str, end: str) -> list[tuple[str, bool, str]]: | |
| """Return shortest path as (relationship_type, forward_direction, next_label) steps.""" | |
| adjacency: dict[str, list[tuple[str, int]]] = {} | |
| for idx, edge in enumerate(self.edges): | |
| adjacency.setdefault(edge.from_node, []).append((edge.to_node, idx)) | |
| adjacency.setdefault(edge.to_node, []).append((edge.from_node, idx)) | |
| queue = deque([start]) | |
| parents: dict[str, tuple[str, int] | None] = {start: None} | |
| while queue: | |
| current = queue.popleft() | |
| if current == end: | |
| break | |
| for neighbor, edge_idx in adjacency.get(current, []): | |
| if neighbor in parents: | |
| continue | |
| parents[neighbor] = (current, edge_idx) | |
| queue.append(neighbor) | |
| if end not in parents: | |
| return [] | |
| node_path: list[str] = [end] | |
| edge_path: list[int] = [] | |
| cursor = end | |
| while parents[cursor] is not None: | |
| parent, edge_idx = parents[cursor] | |
| edge_path.append(edge_idx) | |
| node_path.append(parent) | |
| cursor = parent | |
| node_path.reverse() | |
| edge_path.reverse() | |
| steps: list[tuple[str, bool, str]] = [] | |
| for i, edge_idx in enumerate(edge_path): | |
| edge = self.edges[edge_idx] | |
| current_label = node_path[i] | |
| next_label = node_path[i + 1] | |
| is_forward = edge.from_node == current_label and edge.to_node == next_label | |
| steps.append((edge.type, is_forward, next_label)) | |
| return steps | |
| def _steps_to_pattern(start_label: str, steps: list[tuple[str, bool, str]]) -> str: | |
| pattern = f"(:`{start_label}`)" | |
| for rel_type, is_forward, next_label in steps: | |
| if is_forward: | |
| pattern += f"-[:{rel_type}]->(:`{next_label}`)" | |
| else: | |
| pattern += f"<-[:{rel_type}]-(:`{next_label}`)" | |
| return pattern | |
| def _incident_path_patterns(self, label: str) -> list[str]: | |
| patterns: list[str] = [] | |
| for edge in self.edges: | |
| if edge.from_node == label: | |
| patterns.append(f"(:`{label}`)-[:{edge.type}]->(:`{edge.to_node}`)") | |
| elif edge.to_node == label: | |
| patterns.append(f"(:`{label}`)<-[:{edge.type}]-(:`{edge.from_node}`)") | |
| return patterns | |
| def _two_hop_path_patterns(self, label: str) -> list[str]: | |
| """Build 1-hop and 2-hop directional candidate patterns from a single seed.""" | |
| one_hop = self._incident_path_patterns(label) | |
| patterns: list[str] = list(one_hop) | |
| for idx, edge in enumerate(self.edges): | |
| if edge.from_node == label: | |
| neighbor = edge.to_node | |
| first = f"(:`{label}`)-[:{edge.type}]->(:`{neighbor}`)" | |
| elif edge.to_node == label: | |
| neighbor = edge.from_node | |
| first = f"(:`{label}`)<-[:{edge.type}]-(:`{neighbor}`)" | |
| else: | |
| continue | |
| for j, second_edge in enumerate(self.edges): | |
| if j == idx: | |
| continue | |
| if second_edge.from_node == neighbor: | |
| second = f"-[:{second_edge.type}]->(:`{second_edge.to_node}`)" | |
| elif second_edge.to_node == neighbor: | |
| second = f"<-[:{second_edge.type}]-(:`{second_edge.from_node}`)" | |
| else: | |
| continue | |
| patterns.append(first + second) | |
| return list(dict.fromkeys(patterns)) | |