File size: 2,738 Bytes
8e874f5 | 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 | """Abstract base class for graph retrievers.
This module defines the interface that all retriever algorithms must implement.
Retrievers are responsible for traversing knowledge graphs to find relevant
nodes and paths based on various algorithms (flow diffusion, PageRank, BFS, etc.).
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Optional, Any
@dataclass
class RetrieverResult:
"""Result from a retriever operation.
Attributes:
nodes: Dictionary mapping node IDs to their scores/flow values
path: Optional list of node IDs representing a path
score: Overall score for the retrieval result
metadata: Optional additional metadata from the retrieval
"""
nodes: Dict[str, float] = field(default_factory=dict)
path: Optional[List[str]] = None
score: float = 0.0
metadata: Dict[str, Any] = field(default_factory=dict)
class BaseRetriever(ABC):
"""Abstract base class for graph traversal/retrieval algorithms.
All retriever implementations should inherit from this class and
implement the required abstract methods.
Attributes:
graph: The graph object to traverse (typically NetworkX graph)
"""
def __init__(self, graph, **kwargs):
"""Initialize the retriever with a graph.
Args:
graph: The graph to traverse
**kwargs: Additional algorithm-specific parameters
"""
self.graph = graph
@abstractmethod
def retrieve(
self,
source_node: str,
target_node: Optional[str] = None,
**kwargs
) -> RetrieverResult:
"""Retrieve relevant nodes from the graph.
Args:
source_node: Starting node for retrieval
target_node: Optional target node for path finding
**kwargs: Algorithm-specific parameters
Returns:
RetrieverResult containing nodes, optional path, and score
"""
pass
def find_path(
self,
source: str,
target: str
) -> Tuple[Optional[List[str]], float]:
"""Find path between source and target nodes.
Optional — not all retrievers support pathfinding (e.g. flow diffusion
returns ranked clusters, not paths).
Args:
source: Source node ID
target: Target node ID
Returns:
Tuple of (path as list of node IDs or None, score)
"""
return None, 0.0
def get_node_scores(self) -> Dict[str, float]:
"""Get the scores/values for all processed nodes.
Returns:
Dictionary mapping node IDs to their scores
"""
return {}
|