Fix projection and missing functionalities

CHANGELOG.md

@@ -1,9 +1,10 @@
 # Changelog - [0.3.0] - yyyy-mm-dd
 
 ### New Features
-
+1. Add a new get_relationship_properties_keys tool.
 
 ### Bug Fixes
+1. Return node names in several path algorithms that only returned node ids.
 
 * Fix a bug with loading node properties correctly.
 

mcp_server/src/mcp_server_neo4j_gds/gds.py

@@ -158,8 +158,7 @@ def projected_graph(gds, undirected=False):
 
         # Use separate data and additional configuration parameters
         if additional_config:
-            G, _ = gds.graph.cypher.project(
-                f"""
+            project_query = f"""
                        MATCH (n)-[r]->(m)
                        WITH n, r, m
                        RETURN gds.graph.project(
@@ -169,12 +168,14 @@ def projected_graph(gds, undirected=False):
                            {{{data_config}}},
                            {{{additional_config}}}
                        )
-                       """,
+                       """
+            logger.info(f"Project query: '{project_query}'")
+            G, _ = gds.graph.cypher.project(
+                project_query,
                 graph_name=graph_name,
             )
         else:
-            G, _ = gds.graph.cypher.project(
-                f"""
+            projection_query = f"""
                        MATCH (n)-[r]->(m)
                        WITH n, r, m
                        RETURN gds.graph.project(
@@ -183,7 +184,10 @@ def projected_graph(gds, undirected=False):
                            m,
                            {{{data_config}}}
                        )
-                       """,
+                       """
+            logger.info(f"Projection query: '{projection_query}'")
+            G, _ = gds.graph.cypher.project(
+                projection_query,
                 graph_name=graph_name,
             )
         yield G
@@ -206,3 +210,15 @@ def get_node_properties_keys(gds: GraphDataScience):
         if df.empty:
             return []
         return df["properties_keys"].iloc[0]
+
+
+def get_relationship_properties_keys(gds: GraphDataScience):
+    with projected_graph(gds):
+        query = """
+        MATCH (n)-[r]->(m)
+        RETURN DISTINCT keys(properties(r)) AS properties_keys
+        """
+        df = gds.run_cypher(query)
+        if df.empty:
+            return []
+        return df["properties_keys"].iloc[0]

mcp_server/src/mcp_server_neo4j_gds/path_algorithm_handlers.py

@@ -562,10 +562,16 @@ def minimum_directed_steiner_tree(
 
                 total_weight += weight
 
+                # Get node names using GDS utility function
+                node_name = self.gds.util.asNode(node_id)
+                parent_name = self.gds.util.asNode(parent_id)
+
                 edges.append(
                     {
                         "nodeId": node_id,
                         "parentId": parent_id,
+                        "nodeName": node_name,
+                        "parentName": parent_name,
                         "weight": weight,
                     }
                 )
@@ -590,7 +596,7 @@ def execute(self, arguments: Dict[str, Any]) -> Any:
 class PrizeCollectingSteinerTreeHandler(AlgorithmHandler):
     def prize_collecting_steiner_tree(self, **kwargs):
         with projected_graph(self.gds, undirected=True) as G:
-            # If any optional parameter is not None, use that parameter
+            # Prepare parameters for the algorithm
             params = {k: v for k, v in kwargs.items() if v is not None}
             logger.info(f"Prize-Collecting Steiner Tree parameters: {params}")
 
@@ -618,10 +624,16 @@ def prize_collecting_steiner_tree(self, **kwargs):
 
                 total_weight += weight
 
+                # Get node names using GDS utility function if available
+                node_name = self.gds.util.asNode(node_id)
+                parent_name = self.gds.util.asNode(parent_id)
+
                 edges.append(
                     {
                         "nodeId": node_id,
                         "parentId": parent_id,
+                        "nodeName": node_name,
+                        "parentName": parent_name,
                         "weight": weight,
                     }
                 )
@@ -660,10 +672,16 @@ def all_pairs_shortest_paths(self, **kwargs):
                 target_id = int(row["targetNodeId"])
                 distance = float(row["distance"])
 
+                # Get node names using GDS utility function
+                source_name = self.gds.util.asNode(source_id)
+                target_name = self.gds.util.asNode(target_id)
+
                 paths.append(
                     {
                         "sourceNodeId": source_id,
                         "targetNodeId": target_id,
+                        "sourceNodeName": source_name,
+                        "targetNodeName": target_name,
                         "distance": distance,
                     }
                 )

mcp_server/src/mcp_server_neo4j_gds/server.py

@@ -12,7 +12,7 @@
 from .community_algorithm_specs import community_tool_definitions
 from .path_algorithm_specs import path_tool_definitions
 from .registry import AlgorithmRegistry
-from .gds import count_nodes, get_node_properties_keys
+from .gds import count_nodes, get_node_properties_keys, get_relationship_properties_keys
 
 logger = logging.getLogger("mcp_server_neo4j_gds")
 
@@ -80,6 +80,13 @@ async def handle_list_tools() -> list[types.Tool]:
                             "type": "object",
                         },
                     ),
+                    types.Tool(
+                        name="get_relationship_properties_keys",
+                        description="""Get all relationship properties keys in the database""",
+                        inputSchema={
+                            "type": "object",
+                        },
+                    ),
                 ]
                 + centrality_tool_definitions
                 + community_tool_definitions
@@ -105,6 +112,10 @@ async def handle_call_tool(
                 result = get_node_properties_keys(gds)
                 return [types.TextContent(type="text", text=serialize_result(result))]
 
+            elif name == "get_relationship_properties_keys":
+                result = get_relationship_properties_keys(gds)
+                return [types.TextContent(type="text", text=serialize_result(result))]
+
             else:
                 handler = AlgorithmRegistry.get_handler(name, gds)
                 result = handler.execute(arguments or {})

mcp_server/tests/test_basic_tools.py

@@ -25,6 +25,7 @@ async def test_list_tools(mcp_client):
         # Basic tools
         "count_nodes",
         "get_node_properties_keys",
+        "get_relationship_properties_keys",
         # Centrality algorithms
         "article_rank",
         "articulation_points",
@@ -97,3 +98,14 @@ async def test_get_node_properties_keys(mcp_client):
         "display_name",
         "longitude",
     ]
+
+
+@pytest.mark.asyncio
+async def test_get_relationship_properties_keys(mcp_client):
+    result = await mcp_client.call_tool("get_relationship_properties_keys")
+
+    assert len(result) == 1
+    result_text = result[0]["text"]
+    properties_keys = json.loads(result_text)
+
+    assert properties_keys == ["distance", "line", "time"]

mcp_server/tests/test_path_algorithms.py

@@ -317,6 +317,8 @@ async def test_minimum_directed_steiner_tree(mcp_client):
     first_edge = edges[0]
     assert "nodeId" in first_edge
     assert "parentId" in first_edge
+    assert "nodeName" in first_edge
+    assert "parentName" in first_edge
     assert "weight" in first_edge
     assert first_edge["weight"] > 0
 
@@ -360,6 +362,8 @@ async def test_prize_collecting_steiner_tree(mcp_client):
     first_edge = edges[0]
     assert "nodeId" in first_edge
     assert "parentId" in first_edge
+    assert "nodeName" in first_edge
+    assert "parentName" in first_edge
     assert "weight" in first_edge
     assert first_edge["weight"] > 0
 
@@ -388,6 +392,8 @@ async def test_all_pairs_shortest_paths(mcp_client):
     first_path = paths[0]
     assert "sourceNodeId" in first_path
     assert "targetNodeId" in first_path
+    assert "sourceNodeName" in first_path
+    assert "targetNodeName" in first_path
     assert "distance" in first_path
 
     distance = first_path["distance"]