gds-agent/mcp_server/src/mcp_server_neo4j_gds/registry.py at b28fc2aa640609bab209b27bdf33e0aba0bca478 · brs96/gds-agent · GitHub

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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
from typing import Dict, Type
from graphdatascience import GraphDataScience
from .algorithm_handler import AlgorithmHandler
from .centrality_algorithm_handlers import (
    PageRankHandler,
    ArticleRankHandler,
    DegreeCentralityHandler,
    ArticulationPointsHandler,
    BetweennessCentralityHandler,
    BridgesHandler,
    CELFHandler,
    ClosenessCentralityHandler,
    EigenvectorCentralityHandler,
    HarmonicCentralityHandler,
    HITSHandler,
)
from .community_algorithm_handlers import (
    ConductanceHandler,
    HDBSCANHandler,
    KCoreDecompositionHandler,
    K1ColoringHandler,
    KMeansClusteringHandler,
    LabelPropagationHandler,
    LeidenHandler,
    LocalClusteringCoefficientHandler,
    LouvainHandler,
    ModularityMetricHandler,
    ModularityOptimizationHandler,
    StronglyConnectedComponentsHandler,
    TriangleCountHandler,
    WeaklyConnectedComponentsHandler,
    ApproximateMaximumKCutHandler,
    SpeakerListenerLabelPropagationHandler,
)
from .similarity_algorithm_handlers import (
    NodeSimilarityHandler,
    KNearestNeighborsHandler,
)
from .path_algorithm_handlers import (
    DijkstraShortestPathHandler,
    DeltaSteppingShortestPathHandler,
    DijkstraSingleSourceShortestPathHandler,
    AStarShortestPathHandler,
    YensShortestPathsHandler,
    MinimumWeightSpanningTreeHandler,
    MinimumDirectedSteinerTreeHandler,
    PrizeCollectingSteinerTreeHandler,
    AllPairsShortestPathsHandler,
    RandomWalkHandler,
    BreadthFirstSearchHandler,
    DepthFirstSearchHandler,
    BellmanFordSingleSourceShortestPathHandler,
    LongestPathHandler,
    MaxFlowHandler,
)


class AlgorithmRegistry:
    _handlers: Dict[str, Type[AlgorithmHandler]] = {
        # Centrality algorithms
        "article_rank": ArticleRankHandler,
        "articulation_points": ArticulationPointsHandler,
        "betweenness_centrality": BetweennessCentralityHandler,
        "bridges": BridgesHandler,
        "CELF": CELFHandler,
        "closeness_centrality": ClosenessCentralityHandler,
        "degree_centrality": DegreeCentralityHandler,
        "eigenvector_centrality": EigenvectorCentralityHandler,
        "pagerank": PageRankHandler,
        "harmonic_centrality": HarmonicCentralityHandler,
        "HITS": HITSHandler,
        # Community detection algorithms
        "conductance": ConductanceHandler,
        "hdbscan": HDBSCANHandler,
        "k_core_decomposition": KCoreDecompositionHandler,
        "k_1_coloring": K1ColoringHandler,
        "k_means_clustering": KMeansClusteringHandler,
        "label_propagation": LabelPropagationHandler,
        "leiden": LeidenHandler,
        "local_clustering_coefficient": LocalClusteringCoefficientHandler,
        "louvain": LouvainHandler,
        "modularity_metric": ModularityMetricHandler,
        "modularity_optimization": ModularityOptimizationHandler,
        "strongly_connected_components": StronglyConnectedComponentsHandler,
        "triangle_count": TriangleCountHandler,
        "weakly_connected_components": WeaklyConnectedComponentsHandler,
        "approximate_maximum_k_cut": ApproximateMaximumKCutHandler,
        "speaker_listener_label_propagation": SpeakerListenerLabelPropagationHandler,
        # Similarity algorithms
        "node_similarity": NodeSimilarityHandler,
        "k_nearest_neighbors": KNearestNeighborsHandler,
        # Path finding algorithms
        "find_shortest_path": DijkstraShortestPathHandler,
        "delta_stepping_shortest_path": DeltaSteppingShortestPathHandler,
        "dijkstra_single_source_shortest_path": DijkstraSingleSourceShortestPathHandler,
        "a_star_shortest_path": AStarShortestPathHandler,
        "yens_shortest_paths": YensShortestPathsHandler,
        "minimum_weight_spanning_tree": MinimumWeightSpanningTreeHandler,
        "minimum_directed_steiner_tree": MinimumDirectedSteinerTreeHandler,
        "prize_collecting_steiner_tree": PrizeCollectingSteinerTreeHandler,
        "all_pairs_shortest_paths": AllPairsShortestPathsHandler,
        "random_walk": RandomWalkHandler,
        "breadth_first_search": BreadthFirstSearchHandler,
        "depth_first_search": DepthFirstSearchHandler,
        "bellman_ford_single_source_shortest_path": BellmanFordSingleSourceShortestPathHandler,
        "longest_path": LongestPathHandler,
        "max_flow": MaxFlowHandler,
    }

    @classmethod
    def get_handler(cls, name: str, gds: GraphDataScience) -> AlgorithmHandler:
        handler_class = cls._handlers.get(name)
        if handler_class is None:
            raise ValueError(f"Unknown tool: {name}.")
        return handler_class(gds)