"""Query-Aware Weighted Flow Diffusion retriever. This module implements the Query-Aware Weighted Flow Diffusion algorithm for knowledge graph traversal. The algorithm uses flow diffusion to rank nodes by relevance, with query-aware edge weighting based on node embeddings and query embeddings. Algorithm (push-relabel flow diffusion): 1. Initialize mass at source node(s), warm-start x via lazy random walk 2. Iteratively push excess mass to neighbors via query-aware edges 3. Accumulate importance scores in x with configurable step size 4. Return node flow values as relevance ranking """ import random from collections import defaultdict from typing import Dict, List, Optional import networkx as nx from ..utils import logger from .base import BaseRetriever, RetrieverResult class FlowDiffusionRetriever(BaseRetriever): """Query-Aware Weighted Flow Diffusion algorithm for graph retrieval. Uses push-relabel flow diffusion to rank nodes by relevance. Each source node diffuses mass through the graph; nodes accumulate importance scores proportional to flow passing through them. Attributes: graph: NetworkX graph to traverse source: Source node ID confidence: Confidence level (0.0-1.0) epsilon: Convergence threshold step_size: Learning rate for flow accumulation node_embeddings: Optional dict of node embeddings subquery_embedding: Optional query embedding vector weight_func: Weight function type ("multiply", "add", or None) """ def __init__( self, graph, source_node: str, target_node: str, confidence: float = 0.5, epsilon: float = 0.01, node_embeddings: Optional[Dict] = None, subquery_embedding: Optional[List[float]] = None, weight_func: Optional[str] = None, step_size: float = 0.2, random_seed: int = 42 ): """Initialize the Query-Aware Weighted Flow Diffusion algorithm. Args: graph: NetworkX graph to traverse source_node: Starting node for flow diffusion target_node: Target node (unused in ranking mode) confidence: Confidence level (0.0-1.0, default 0.5) epsilon: Convergence threshold (default 0.02) node_embeddings: Optional dict mapping node IDs to embeddings subquery_embedding: Optional query embedding vector weight_func: Optional weight function ("multiply", "add", or None) step_size: Learning rate for flow accumulation (default 0.2) random_seed: Random seed for reproducibility (default 42) """ super().__init__(graph) self.source = source_node self.target = target_node self.confidence = max(0.0, min(1.0, confidence)) self.epsilon = epsilon self.step_size = step_size self.mass = defaultdict(float) self.x = defaultdict(float) self.sink_capacity = defaultdict(float) # Query-aware components self.node_embeddings = node_embeddings or {} embedding_dim = len(subquery_embedding) if subquery_embedding else 1536 self.subquery_embedding = subquery_embedding or [0.0] * embedding_dim self.weight_func = weight_func self.edge_weights_cache = {} # Deterministic seeding for reproducibility random.seed(random_seed) def cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float: """Cosine similarity normalized to [0, 1].""" if not vec1 or not vec2: return 0.0 try: dot_product = sum(a * b for a, b in zip(vec1, vec2)) mag1 = sum(a * a for a in vec1) ** 0.5 mag2 = sum(b * b for b in vec2) ** 0.5 if mag1 == 0 or mag2 == 0: return 0.0 similarity = dot_product / (mag1 * mag2) return max(0.0, (similarity + 1.0) / 2.0) except Exception: return 0.0 def get_edge_weight(self, node1: str, node2: str) -> float: """Get query-aware edge weight: w'(u,v) = w(u,v) * f(sim(u,q), sim(v,q)).""" cache_key = (node1, node2) if cache_key in self.edge_weights_cache: return self.edge_weights_cache[cache_key] edge_data = self.graph[node1][node2] original_weight = edge_data.get('weight', 1.0) if not self.node_embeddings or not self.subquery_embedding: self.edge_weights_cache[cache_key] = original_weight return original_weight if original_weight <= 0: self.edge_weights_cache[cache_key] = 0.0 return 0.0 node1_emb = self.node_embeddings.get(node1, [0.0] * len(self.subquery_embedding)) node2_emb = self.node_embeddings.get(node2, [0.0] * len(self.subquery_embedding)) node1_query_sim = self.cosine_similarity(node1_emb, self.subquery_embedding) node2_query_sim = self.cosine_similarity(node2_emb, self.subquery_embedding) if self.weight_func == "multiply": smart_weight = original_weight * node1_query_sim * node2_query_sim elif self.weight_func == "add": smart_weight = (original_weight + node1_query_sim + node2_query_sim) / 3.0 else: query_factor = (node1_query_sim + node2_query_sim) / 2.0 smart_weight = original_weight * (1.0 + query_factor * 0.5) self.edge_weights_cache[cache_key] = smart_weight return smart_weight def initialize(self, alpha: float = 50, use_node_degree: bool = True): """Initialize sink capacities, source mass, and warm-start x. Args: alpha: Mass initialization factor (default 50) use_node_degree: Whether to use node degree for sink capacity (default True) """ # Set sink capacity proportional to node degree for node in self.graph.nodes(): if use_node_degree: self.sink_capacity[node] = max(self.graph.degree(node), 1) else: self.sink_capacity[node] = 1 # Normalize sink capacity to sum to 10 total_sink = sum(self.sink_capacity.values()) if total_sink > 0: for node in self.sink_capacity: self.sink_capacity[node] = 10.0 * self.sink_capacity[node] / total_sink total_sink = sum(self.sink_capacity.values()) # Set all masses to 0 for node in self.graph.nodes(): self.mass[node] = 0 # Inject mass at source: alpha * total_sink * confidence_boost confidence_boost = 1.0 + self.confidence self.mass[self.source] = alpha * total_sink * confidence_boost # Warm-start x: 2-step lazy random walk from source self.x = defaultdict(float) self.x[self.source] = 1.0 for _ in range(2): x_new = defaultdict(float) for node, val in self.x.items(): if val > 0: neighbors = list(self.graph.neighbors(node)) if neighbors: for neighbor in neighbors: x_new[neighbor] += val / len(neighbors) # Lazy update: average diffused values with source-anchored distribution x_combined = defaultdict(float) x_combined[self.source] = 1.0 for node in set(list(x_new.keys()) + [self.source]): x_combined[node] = (x_combined.get(node, 0.0) + x_new.get(node, 0.0)) / 2.0 self.x = x_combined def push(self, node: str) -> bool: """Push excess mass from node to neighbors. Decoupled accumulation/routing: x accumulates by structural degree, mass routes by query-aware edge weights. """ neighbors = list(self.graph.neighbors(node)) if not neighbors: return False # Query-aware weights (for routing) w_qa = 0 for neighbor in neighbors: w_qa += self.get_edge_weight(node, neighbor) if w_qa == 0: return False excess = self.mass[node] - self.sink_capacity[node] if excess <= 0: return False # Structural weights (for accumulation) — decoupled from QA w_struct = 0 for neighbor in neighbors: edge_data = self.graph[node][neighbor] w_struct += edge_data.get('weight', 1.0) if w_struct == 0: w_struct = w_qa # Accumulate importance based on structural degree self.x[node] += self.step_size * excess / (w_struct + 1e-8) # Absorb what sink can hold self.mass[node] = self.sink_capacity[node] # Route mass using query-aware weights for neighbor in neighbors: w_ij = self.get_edge_weight(node, neighbor) if w_ij > 0: self.mass[neighbor] += excess * w_ij / (w_qa + 1e-8) return True def flow_diffusion(self, max_iterations: int = 500) -> Dict[str, float]: """Run push-relabel flow diffusion until convergence. Args: max_iterations: Maximum iterations to run (default 500) Returns: Dictionary of nodes with positive flow values """ iterations = 0 pushes = 0 while iterations < max_iterations: iterations += 1 # Find nodes with excess mass excess_nodes = [node for node in self.graph.nodes() if self.mass[node] > self.sink_capacity[node] + self.epsilon] if not excess_nodes: logger.debug(f"QAFD converged in {iterations} iterations ({pushes} pushes)") break node = random.choice(excess_nodes) if self.push(node): pushes += 1 # Check convergence every 10 iterations if iterations % 10 == 0: remaining_excess = sum(max(0, self.mass[node] - self.sink_capacity[node]) for node in self.graph.nodes()) if remaining_excess < self.epsilon: logger.debug(f"QAFD converged in {iterations} iterations ({pushes} pushes)") break if iterations >= max_iterations: logger.warning(f"QAFD did not converge after {max_iterations} iterations") return {node: val for node, val in self.x.items() if val > 0} def retrieve( self, source_node: Optional[str] = None, target_node: Optional[str] = None, **kwargs ) -> RetrieverResult: """Retrieve nodes using flow diffusion. Args: source_node: Optional override for source node target_node: Optional override for target node **kwargs: Additional parameters: - alpha: Mass initialization factor (default 10) - max_iterations: Max diffusion iterations (default 500) Returns: RetrieverResult with diffused nodes and scores """ alpha = kwargs.get('alpha', 50) max_iterations = kwargs.get('max_iterations', 500) self.initialize(alpha=alpha) diffused_nodes = self.flow_diffusion(max_iterations=max_iterations) return RetrieverResult( nodes=diffused_nodes, path=None, score=0.0, metadata={ 'source': self.source, 'target': self.target, 'confidence': self.confidence, 'weight_func': self.weight_func } ) def get_node_scores(self) -> Dict[str, float]: """Get the flow values for all processed nodes.""" return dict(self.x) # Alias for backward compatibility QueryAwareWeightedFlowDiffusion = FlowDiffusionRetriever