From e182a3c76ce289f12e9e88c17269e296f2a4841e Mon Sep 17 00:00:00 2001 From: Brian Shi Date: Tue, 9 Sep 2025 14:25:54 +0100 Subject: [PATCH 1/2] add snowflake config --- mcp_server/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/mcp_server/pyproject.toml b/mcp_server/pyproject.toml index 3aec647..c8af051 100644 --- a/mcp_server/pyproject.toml +++ b/mcp_server/pyproject.toml @@ -16,6 +16,7 @@ classifiers = [ dependencies = [ "graphdatascience>=1.16", "mcp[cli]>=1.11.0", + "snowflake-connector-python==3.17.3", ] [project.urls] From 8dbc313bc5cc36ee9b68327dace8ca8e83a5fd38 Mon Sep 17 00:00:00 2001 From: Brian Shi Date: Wed, 10 Sep 2025 16:43:11 +0100 Subject: [PATCH 2/2] wip use snowgraph degree centrality --- mcp_server/pyproject.toml | 2 +- .../centrality_algorithm_handlers.py | 92 ++++++++++++++----- .../centrality_algorithm_specs.py | 13 +-- mcp_server/src/mcp_server_neo4j_gds/server.py | 9 +- 4 files changed, 83 insertions(+), 33 deletions(-) diff --git a/mcp_server/pyproject.toml b/mcp_server/pyproject.toml index c8af051..020274a 100644 --- a/mcp_server/pyproject.toml +++ b/mcp_server/pyproject.toml @@ -16,7 +16,7 @@ classifiers = [ dependencies = [ "graphdatascience>=1.16", "mcp[cli]>=1.11.0", - "snowflake-connector-python==3.17.3", + "snowflake-snowpark-python==1.38.0", ] [project.urls] diff --git a/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_handlers.py b/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_handlers.py index 97f7892..67e52c8 100644 --- a/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_handlers.py +++ b/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_handlers.py @@ -187,32 +187,82 @@ def execute(self, arguments: Dict[str, Any]) -> Any: class DegreeCentralityHandler(AlgorithmHandler): def degree_centrality(self, **kwargs): - with projected_graph(self.gds) as G: - params = { - k: v - for k, v in kwargs.items() - if v is not None and k not in ["nodes", "nodeIdentifierProperty"] + # with projected_graph(self.gds) as G: + # params = { + # k: v + # for k, v in kwargs.items() + # if v is not None and k not in ["nodes", "nodeIdentifierProperty"] + # } + # logger.info(f"Degree centrality parameters: {params}") + # centrality = self.gds.degree.stream(G, **params) + + # # Add node names to the results if nodeIdentifierProperty is provided + # node_identifier_property = kwargs.get("nodeIdentifierProperty") + # translate_ids_to_identifiers(self.gds, node_identifier_property, centrality) + + # # Filter results by node names if provided + # node_names = kwargs.get("nodes", None) + # centrality = filter_identifiers( + # self.gds, node_identifier_property, node_names, centrality + # ) + orientation = kwargs.get("orientation") + if orientation is None: + relationships_dict = """ + { + 'RELATIONSHIPS': { + 'sourceTable': 'NODES', + 'targetTable': 'NODES', + } } - logger.info(f"Degree centrality parameters: {params}") - centrality = self.gds.degree.stream(G, **params) - - # Add node names to the results if nodeIdentifierProperty is provided - node_identifier_property = kwargs.get("nodeIdentifierProperty") - translate_ids_to_identifiers(self.gds, node_identifier_property, centrality) - - # Filter results by node names if provided - node_names = kwargs.get("nodes", None) - centrality = filter_identifiers( - self.gds, node_identifier_property, node_names, centrality - ) - - return centrality + """ + else: + relationships_dict = f""" + {{ + 'RELATIONSHIPS': {{ + 'sourceTable': 'NODES', + 'targetTable': 'NODES', + 'orientation': '{orientation}' + }} + }} + """ + + relationshipWeightProperty = kwargs.get("relationshipWeightProperty") + if relationshipWeightProperty is None: + compute_dict = """ + { } + """ + else: + compute_dict = f""" + {{ + 'relationshipWeightProperty': '{relationshipWeightProperty}' + }} + """ + + res = self.gds.sql( + f""" + CALL Neo4j_Graph_Analytics.graph.degree('CPU_X64_XS', {{ + 'defaultTablePrefix': 'EXAMPLE_DB.PUBLIC', + 'project': {{ + 'nodeTables': [ 'NODES' ], + 'relationshipTables': {relationships_dict} + }}, + 'compute': {compute_dict}, + 'write': [{{ + 'nodeLabel': 'NODES', + 'outputTable': 'NODES_DEGREE_CENTRALITY' + }}] + }}); + """ + ).collect() + + logger.info(f"Degree centrality execution: {res}") + output_table = self.gds.table("EXAMPLE_DB.PUBLIC.NODES_DEGREE_CENTRALITY") + return output_table.to_pandas() def execute(self, arguments: Dict[str, Any]) -> Any: return self.degree_centrality( - nodes=arguments.get("nodes"), - nodeIdentifierProperty=arguments.get("nodeIdentifierProperty"), orientation=arguments.get("orientation"), + relationshipWeightProperty=arguments.get("relationshipWeightProperty"), ) diff --git a/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_specs.py b/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_specs.py index 3c71dbf..05339e9 100644 --- a/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_specs.py +++ b/mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_specs.py @@ -177,19 +177,14 @@ inputSchema={ "type": "object", "properties": { - "nodes": { - "type": "array", - "items": {"type": "string"}, - "description": "List of node names to filter degree centrality results for.", - }, - "nodeIdentifierProperty": { - "type": "string", - "description": "Property name to use for identifying nodes (e.g., 'name', 'Name', 'title'). Use get_node_properties_keys to find available properties.", - }, "orientation": { "type": "string", "description": "The orientation used to compute node degrees. Supported orientations are NATURAL (for out-degree), REVERSE (for in-degree) and UNDIRECTED (for both in-degree and out-degree) ", }, + "relationshipWeightProperty": { + "type": "string", + "description": "Property of the relationship to use for weighting. If not specified, all relationships are treated equally.", + }, }, "required": [], }, diff --git a/mcp_server/src/mcp_server_neo4j_gds/server.py b/mcp_server/src/mcp_server_neo4j_gds/server.py index 0a28114..8d51746 100644 --- a/mcp_server/src/mcp_server_neo4j_gds/server.py +++ b/mcp_server/src/mcp_server_neo4j_gds/server.py @@ -8,7 +8,7 @@ import pandas as pd import json from graphdatascience import GraphDataScience - +from snowflake.snowpark import Session from .similarity_algorithm_specs import similarity_tool_definitions from .centrality_algorithm_specs import centrality_tool_definitions from .community_algorithm_specs import community_tool_definitions @@ -56,7 +56,12 @@ async def main(db_url: str, username: str, password: str, database: str = None): db_url, auth=(username, password), aura_ds=False, database=database ) else: - gds = GraphDataScience(db_url, auth=(username, password), aura_ds=False) + if db_url: + gds = GraphDataScience(db_url, auth=(username, password), aura_ds=False) + else: + gds = Session.builder.config("connection_name", "snowflake-gds-mcp").create() + print(gds.sql("SELECT 1;").collect()) + logger.info("Successfully connected to Snowflake database") logger.info("Successfully connected to Neo4j database") except Exception as e: logger.error(f"Failed to connect to Neo4j database: {e}")