# Enterprise Graph Traversal Module """ Advanced graph traversal algorithms for GraphRAG. Features: - Multi-hop reasoning with path tracking - PageRank for node importance - BFS/DFS traversal strategies - Causal chain detection - Path explanation generation """ from dataclasses import dataclass, field from typing import List, Dict, Optional, Set, Tuple, Any from enum import Enum import networkx as nx from collections import defaultdict import heapq class TraversalStrategy(Enum): """Graph traversal strategies""" BFS = "bfs" # Breadth-first for shortest paths DFS = "dfs" # Depth-first for exhaustive search WEIGHTED = "weighted" # Dijkstra for weighted paths PAGERANK = "pagerank" # Importance-based traversal @dataclass class GraphPath: """A path through the knowledge graph""" nodes: List[str] edges: List[str] total_weight: float hop_count: int path_type: str # "direct", "multi-hop", "causal" explanation: str @dataclass class TraversalResult: """Complete traversal result""" paths: List[GraphPath] visited_nodes: Set[str] node_importance: Dict[str, float] total_hops: int reasoning_chain: List[str] summary: str @dataclass class NodeInfo: """Information about a graph node""" node_id: str node_type: str label: str attributes: Dict[str, Any] importance: float neighbors: List[str] class EnterpriseGraphTraversal: """ Enterprise-grade graph traversal for knowledge graph reasoning. Supports: - Multi-hop path finding - Importance-weighted traversal - Causal chain detection - Path explanation generation """ def __init__(self, graph: nx.Graph): self.graph = graph self._pagerank_cache = None self._node_types_cache = None @property def pagerank(self) -> Dict[str, float]: """Lazy-computed PageRank scores""" if self._pagerank_cache is None: if self.graph.number_of_nodes() > 0: self._pagerank_cache = nx.pagerank(self.graph, weight='weight') else: self._pagerank_cache = {} return self._pagerank_cache def invalidate_cache(self): """Invalidate cached computations""" self._pagerank_cache = None self._node_types_cache = None def get_node_info(self, node_id: str) -> Optional[NodeInfo]: """Get detailed information about a node""" if not self.graph.has_node(node_id): return None attrs = dict(self.graph.nodes[node_id]) node_type = attrs.get('type', 'unknown') label = attrs.get('label', node_id) importance = self.pagerank.get(node_id, 0.0) neighbors = list(self.graph.neighbors(node_id)) return NodeInfo( node_id=node_id, node_type=node_type, label=label, attributes=attrs, importance=importance, neighbors=neighbors ) def find_paths( self, source_entities: List[str], target_entities: Optional[List[str]] = None, max_hops: int = 3, strategy: TraversalStrategy = TraversalStrategy.BFS, max_paths: int = 10 ) -> TraversalResult: """ Find paths between entities in the graph. Args: source_entities: Starting entity IDs or patterns target_entities: Optional target entity IDs max_hops: Maximum path length strategy: Traversal strategy to use max_paths: Maximum number of paths to return Returns: TraversalResult with found paths and metadata """ # Find matching source nodes source_nodes = self._find_matching_nodes(source_entities) if not source_nodes: return TraversalResult( paths=[], visited_nodes=set(), node_importance={}, total_hops=0, reasoning_chain=["No matching source nodes found"], summary="Could not find relevant entities in the knowledge graph" ) # Find matching target nodes if specified target_nodes = None if target_entities: target_nodes = self._find_matching_nodes(target_entities) # Execute traversal based on strategy if strategy == TraversalStrategy.BFS: paths, visited = self._bfs_traversal(source_nodes, target_nodes, max_hops, max_paths) elif strategy == TraversalStrategy.DFS: paths, visited = self._dfs_traversal(source_nodes, target_nodes, max_hops, max_paths) elif strategy == TraversalStrategy.WEIGHTED: paths, visited = self._weighted_traversal(source_nodes, target_nodes, max_hops, max_paths) elif strategy == TraversalStrategy.PAGERANK: paths, visited = self._pagerank_traversal(source_nodes, max_hops, max_paths) else: paths, visited = self._bfs_traversal(source_nodes, target_nodes, max_hops, max_paths) # Compute node importance for visited nodes node_importance = {n: self.pagerank.get(n, 0.0) for n in visited} # Generate reasoning chain reasoning_chain = self._generate_reasoning_chain(paths) # Generate summary summary = self._generate_summary(paths, visited) return TraversalResult( paths=paths, visited_nodes=visited, node_importance=node_importance, total_hops=max(p.hop_count for p in paths) if paths else 0, reasoning_chain=reasoning_chain, summary=summary ) def _find_matching_nodes(self, patterns: List[str]) -> Set[str]: """Find nodes matching given patterns""" matches = set() for pattern in patterns: pattern_lower = pattern.lower() for node in self.graph.nodes(): node_lower = str(node).lower() attrs = self.graph.nodes[node] label = str(attrs.get('label', '')).lower() # Match by ID, label, or type if (pattern_lower in node_lower or pattern_lower in label or pattern_lower == attrs.get('type', '').lower()): matches.add(node) return matches def _bfs_traversal( self, sources: Set[str], targets: Optional[Set[str]], max_hops: int, max_paths: int ) -> Tuple[List[GraphPath], Set[str]]: """Breadth-first traversal for shortest paths""" paths = [] visited = set() from collections import deque for source in sources: if len(paths) >= max_paths: break queue = deque([(source, [source], [], 0)]) local_visited = {source} while queue and len(paths) < max_paths: current, path, edges, hops = queue.popleft() visited.add(current) if hops >= max_hops: continue for neighbor in self.graph.neighbors(current): if neighbor not in local_visited: local_visited.add(neighbor) edge_data = self.graph.edges[current, neighbor] edge_label = edge_data.get('relation', 'connected') new_path = path + [neighbor] new_edges = edges + [edge_label] # Check if target reached if targets and neighbor in targets: paths.append(self._create_path( new_path, new_edges, hops + 1, "direct" )) elif not targets and hops + 1 <= max_hops: # Explore further for discovery paths.append(self._create_path( new_path, new_edges, hops + 1, "exploration" )) queue.append((neighbor, new_path, new_edges, hops + 1)) return paths[:max_paths], visited def _dfs_traversal( self, sources: Set[str], targets: Optional[Set[str]], max_hops: int, max_paths: int ) -> Tuple[List[GraphPath], Set[str]]: """Depth-first traversal for exhaustive search""" paths = [] visited = set() def dfs(node: str, path: List[str], edges: List[str], depth: int): if len(paths) >= max_paths: return visited.add(node) if depth >= max_hops: return for neighbor in self.graph.neighbors(node): if neighbor not in path: # Avoid cycles edge_data = self.graph.edges[node, neighbor] edge_label = edge_data.get('relation', 'connected') new_path = path + [neighbor] new_edges = edges + [edge_label] if targets and neighbor in targets: paths.append(self._create_path( new_path, new_edges, depth + 1, "multi-hop" )) dfs(neighbor, new_path, new_edges, depth + 1) for source in sources: if len(paths) >= max_paths: break dfs(source, [source], [], 0) return paths[:max_paths], visited def _weighted_traversal( self, sources: Set[str], targets: Optional[Set[str]], max_hops: int, max_paths: int ) -> Tuple[List[GraphPath], Set[str]]: """Weighted path finding using Dijkstra-like approach""" paths = [] visited = set() for source in sources: if len(paths) >= max_paths: break if not targets: continue for target in targets: if len(paths) >= max_paths: break try: path = nx.shortest_path( self.graph, source, target, weight='weight' ) if len(path) - 1 <= max_hops: edges = [] total_weight = 0 for i in range(len(path) - 1): edge_data = self.graph.edges[path[i], path[i+1]] edges.append(edge_data.get('relation', 'connected')) total_weight += edge_data.get('weight', 1.0) gpath = GraphPath( nodes=path, edges=edges, total_weight=total_weight, hop_count=len(path) - 1, path_type="weighted", explanation=self._explain_path(path, edges) ) paths.append(gpath) visited.update(path) except nx.NetworkXNoPath: pass return paths, visited def _pagerank_traversal( self, sources: Set[str], max_hops: int, max_paths: int ) -> Tuple[List[GraphPath], Set[str]]: """Traverse following high-importance nodes""" paths = [] visited = set() for source in sources: if len(paths) >= max_paths: break current = source path = [current] edges = [] for _ in range(max_hops): visited.add(current) # Get neighbors sorted by PageRank neighbors = list(self.graph.neighbors(current)) if not neighbors: break # Choose highest PageRank neighbor not in path neighbors = [n for n in neighbors if n not in path] if not neighbors: break neighbors.sort(key=lambda n: self.pagerank.get(n, 0), reverse=True) next_node = neighbors[0] edge_data = self.graph.edges[current, next_node] edges.append(edge_data.get('relation', 'connected')) path.append(next_node) current = next_node if len(path) > 1: paths.append(self._create_path( path, edges, len(path) - 1, "importance" )) return paths, visited def _create_path( self, nodes: List[str], edges: List[str], hop_count: int, path_type: str ) -> GraphPath: """Create a GraphPath with explanation""" # Calculate weight total_weight = 0 for i in range(len(nodes) - 1): if self.graph.has_edge(nodes[i], nodes[i+1]): edge_data = self.graph.edges[nodes[i], nodes[i+1]] total_weight += edge_data.get('weight', 1.0) return GraphPath( nodes=nodes, edges=edges, total_weight=total_weight, hop_count=hop_count, path_type=path_type, explanation=self._explain_path(nodes, edges) ) def _explain_path(self, nodes: List[str], edges: List[str]) -> str: """Generate human-readable explanation of a path""" if len(nodes) < 2: return f"Single entity: {self._get_node_label(nodes[0])}" parts = [] for i in range(len(nodes) - 1): src_label = self._get_node_label(nodes[i]) edge_label = edges[i] if i < len(edges) else "connected to" dst_label = self._get_node_label(nodes[i + 1]) parts.append(f"{src_label} → [{edge_label}] → {dst_label}") return " | ".join(parts) def _get_node_label(self, node_id: str) -> str: """Get human-readable label for node""" if self.graph.has_node(node_id): attrs = self.graph.nodes[node_id] label = attrs.get('label', node_id) node_type = attrs.get('type', '') if node_type: return f"{label} ({node_type})" return label return node_id def _generate_reasoning_chain(self, paths: List[GraphPath]) -> List[str]: """Generate step-by-step reasoning from paths""" if not paths: return ["No paths found in knowledge graph"] chain = [] for i, path in enumerate(paths[:5]): # Top 5 paths chain.append(f"Path {i+1}: {path.explanation}") return chain def _generate_summary(self, paths: List[GraphPath], visited: Set[str]) -> str: """Generate summary of traversal results""" if not paths: return "No relevant paths found in the knowledge graph." unique_entities = len(visited) total_paths = len(paths) avg_hops = sum(p.hop_count for p in paths) / len(paths) return ( f"Found {total_paths} paths connecting {unique_entities} entities. " f"Average path length: {avg_hops:.1f} hops." ) def get_entity_neighbors( self, entity: str, max_neighbors: int = 10 ) -> List[NodeInfo]: """Get neighbors of an entity sorted by importance""" nodes = self._find_matching_nodes([entity]) if not nodes: return [] neighbors = set() for node in nodes: neighbors.update(self.graph.neighbors(node)) # Sort by PageRank neighbor_info = [] for n in neighbors: info = self.get_node_info(n) if info: neighbor_info.append(info) neighbor_info.sort(key=lambda x: x.importance, reverse=True) return neighbor_info[:max_neighbors] def detect_causal_chain( self, effect_entity: str, max_causes: int = 5 ) -> TraversalResult: """Detect potential causal relationships leading to an effect""" # Find the effect node effect_nodes = self._find_matching_nodes([effect_entity]) if not effect_nodes: return TraversalResult( paths=[], visited_nodes=set(), node_importance={}, total_hops=0, reasoning_chain=["Effect entity not found"], summary="Could not identify the effect entity" ) # Traverse backwards to find causes causal_paths = [] visited = set() for effect in effect_nodes: # Use reverse BFS to find predecessors # In undirected graph, just find connected high-importance nodes neighbors = self.get_entity_neighbors(effect, max_neighbors=max_causes * 2) for neighbor in neighbors[:max_causes]: path = GraphPath( nodes=[neighbor.node_id, effect], edges=["contributes_to"], total_weight=neighbor.importance, hop_count=1, path_type="causal", explanation=f"{neighbor.label} may influence {self._get_node_label(effect)}" ) causal_paths.append(path) visited.add(neighbor.node_id) visited.add(effect) return TraversalResult( paths=causal_paths, visited_nodes=visited, node_importance={n: self.pagerank.get(n, 0) for n in visited}, total_hops=1, reasoning_chain=[p.explanation for p in causal_paths], summary=f"Identified {len(causal_paths)} potential causal factors" ) def create_traversal(graph: nx.Graph) -> EnterpriseGraphTraversal: """Create a graph traversal instance""" return EnterpriseGraphTraversal(graph)