Support filtered knn

Author: IoannisPanagiotas

mcp_server/src/mcp_server_neo4j_gds/similarity_algorithm_handlers.py

@@ -1,8 +1,6 @@
 import logging
 from typing import Dict, Any
 
-from graphdatascience import GraphDataScience
-
 from .algorithm_handler import AlgorithmHandler
 from .gds import projected_graph
 
@@ -201,21 +199,97 @@ def execute(self, arguments: Dict[str, Any]) -> Any:
 
 
 class FilteredKNearestNeighborsHandler(AlgorithmHandler):
-    def filtered_k_nearest_neighbors(
-        self, db_url: str, username: str, password: str, **kwargs
+    def handle_input_nodes(
+        self,
+        input_nodes,
+        input_nodes_variable_name,
+        node_identifier_property,
+        call_params,
     ):
-        gds = GraphDataScience(db_url, auth=(username, password), aura_ds=False)
-        with projected_graph(gds) as G:
+        # Handle input nodes - convert names to IDs if nodeIdentifierProperty is provided
+        if input_nodes is not None and node_identifier_property is not None:
+            if isinstance(input_nodes, list):
+                # Handle list of node names
+                query = f"""
+                    UNWIND $names AS name
+                    MATCH (s)
+                    WHERE toLower(s.{node_identifier_property}) CONTAINS toLower(name)
+                    RETURN id(s) as node_id
+                    """
+                df = self.gds.run_cypher(
+                    query,
+                    params={
+                        "names": input_nodes,
+                    },
+                )
+                input_node_ids = df["node_id"].tolist()
+                call_params[input_nodes_variable_name] = input_node_ids
+            else:
+                # Handle single  node name
+                query = f"""
+                    MATCH (s)
+                    WHERE toLower(s.{node_identifier_property}) CONTAINS toLower($name)
+                    RETURN id(s) as node_id
+                    """
+                df = self.gds.run_cypher(
+                    query,
+                    params={
+                        "name": input_nodes,
+                    },
+                )
+                if not df.empty:
+                    call_params[input_nodes_variable_name] = int(df["node_id"].iloc[0])
+        elif input_nodes is not None:
+            # If input_nodes provided but no nodeIdentifierProperty, pass through as-is
+            call_params[input_nodes_variable_name] = input_nodes
+
+    def filtered_k_nearest_neighbors(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 [
+                    "nodeIdentifierProperty",
+                    "sourceNodeFilter",
+                    "targetNodeFilter",
+                ]
+            }
+            node_identifier_property = kwargs.get("nodeIdentifierProperty")
+            source_nodes = kwargs.get("sourceNodeFilter", None)
+            target_nodes = kwargs.get("targetNodeFilter", None)
+            self.handle_input_nodes(
+                source_nodes, "sourceNodeFilter", node_identifier_property, params
+            )
+            self.handle_input_nodes(
+                target_nodes, "targetNodeFilter", node_identifier_property, params
+            )
+
             logger.info(f"Filtered K-Nearest Neighbors parameters: {kwargs}")
-            filtered_k_nearest_neighbors_result = gds.knn.filtered.stream(G, **kwargs)
+            filtered_k_nearest_neighbors_result = self.gds.knn.filtered.stream(
+                G, **params
+            )
+
+        # Add node names to the results if nodeIdentifierProperty is provided
+        node_identifier_property = kwargs.get("nodeIdentifierProperty")
+        if node_identifier_property is not None:
+            node1_name_values = [
+                self.gds.util.asNode(node_id).get(node_identifier_property)
+                for node_id in filtered_k_nearest_neighbors_result["node1"]
+            ]
+            node2_name_values = [
+                self.gds.util.asNode(node_id).get(node_identifier_property)
+                for node_id in filtered_k_nearest_neighbors_result["node2"]
+            ]
+            filtered_k_nearest_neighbors_result["node1Name"] = node1_name_values
+            filtered_k_nearest_neighbors_result["node2Name"] = node2_name_values
 
         return filtered_k_nearest_neighbors_result
 
     def execute(self, arguments: Dict[str, Any]) -> Any:
         return self.filtered_k_nearest_neighbors(
-            self.db_url,
-            self.username,
-            self.password,
+            nodeIdentifierProperty=arguments.get("nodeIdentifierProperty"),
             sourceNodeFilter=arguments.get("sourceNodeFilter"),
             targetNodeFilter=arguments.get("targetNodeFilter"),
             nodeProperties=arguments.get("nodeProperties"),

mcp_server/src/mcp_server_neo4j_gds/similarity_algorithm_specs.py

@@ -117,7 +117,7 @@
                 "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.",
-                }
+                },
             },
         },
     ),
@@ -185,7 +185,7 @@
                 "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.",
-                }
+                },
             },
             "required": ["nodeProperties"],
         },
@@ -243,12 +243,12 @@
                 },
                 "seedTargetNodes": {
                     "type": "boolean",
-                    "description": "Enable seeding of target nodes.",
+                    "description": "Enable seeding of target nodes. If seeded, every node picks some of the target nodes initially. This guarantees that for every node we can avoid empty result (when the algorithm did not find for it any similar neighbors from the target set). Can only be used if targetNodeFilter is set.",
                 },
                 "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.",
-                }
+                },
             },
             "required": ["sourceNodeFilter", "targetNodeFilter", "nodeProperties"],
         },

mcp_server/tests/test_similarity_algorithms.py

@@ -117,3 +117,80 @@ async def test_k_nearest_neighbors(mcp_client):
     lines = result_text.strip().split("\n")
     data_lines = [line for line in lines[1:] if line.strip()]
     assert len(data_lines) == 302 * 3
+
+
+@pytest.mark.asyncio
+async def test_filtered_knn(mcp_client):
+    # test source-filter only
+    result = await mcp_client.call_tool(
+        "filtered_k_nearest_neighbors",
+        {
+            "nodeIdentifierProperty": "name",
+            "topK": 3,
+            "sourceNodeFilter": ["Acton Town"],
+            "nodeProperties": "rail",
+        },
+    )
+
+    assert len(result) == 1
+    result_text = result[0]["text"]
+    # Verify structure of a result entry
+    assert "node1" in result_text
+    assert "node2" in result_text
+    assert "node1Name" in result_text
+    assert "node2Name" in result_text
+    assert "similarity" in result_text
+    lines = result_text.strip().split("\n")
+    data_lines = [line for line in lines[1:] if line.strip()]
+    assert len(data_lines) == 3
+    assert "Acton Town" in data_lines[0]
+
+    # test target-filter alone
+
+    result = await mcp_client.call_tool(
+        "filtered_k_nearest_neighbors",
+        {
+            "nodeIdentifierProperty": "name",
+            "topK": 3,
+            "targetNodeFilter": "Stamford Brook",
+            "nodeProperties": "rail",
+        },
+    )
+    assert len(result) == 1
+    result_text = result[0]["text"]
+    # Verify structure of a result entry
+    assert "node1" in result_text
+    assert "node2" in result_text
+    assert "node1Name" in result_text
+    assert "node2Name" in result_text
+    assert "similarity" in result_text
+    lines = result_text.strip().split("\n")
+    data_lines = [line for line in lines[1:] if line.strip()]
+    assert len(data_lines) > 0
+    assert "Stamford Brook" in data_lines[0]
+
+    # test combination of filters
+    result = await mcp_client.call_tool(
+        "filtered_node_similarity",
+        {
+            "nodeIdentifierProperty": "name",
+            "topK": 3,
+            "sourceNodeFilter": ["Acton Town"],
+            "targetNodeFilter": ["Stamford Brook"],
+            "seedTargetNodes": True,  # k-nn filtering is a bit special, it might not necessarily find answer if this is not specified (at least for this small example graph)
+        },
+    )
+
+    assert len(result) == 1
+    result_text = result[0]["text"]
+    # Verify structure of a result entry
+    assert "node1" in result_text
+    assert "node2" in result_text
+    assert "node1Name" in result_text
+    assert "node2Name" in result_text
+    assert "similarity" in result_text
+    lines = result_text.strip().split("\n")
+    data_lines = [line for line in lines[1:] if line.strip()]
+    assert len(data_lines) == 1
+    assert "Acton Town" in data_lines[0]
+    assert "Stamford Brook" in data_lines[0]