Add an node label to filter on for external retriever base class and WeaviateRetriever

Author: stellasia

src/neo4j_graphrag/retrievers/base.py

@@ -454,10 +454,12 @@ def __init__(
         id_property_external: str,
         id_property_neo4j: str,
         neo4j_database: Optional[str] = None,
+        node_label_neo4j: Optional[str] = None,
     ):
         super().__init__(driver)
         self.id_property_external = id_property_external
         self.id_property_neo4j = id_property_neo4j
+        self.node_label_neo4j = node_label_neo4j
         self.neo4j_database = neo4j_database
 
     @abstractmethod

src/neo4j_graphrag/retrievers/external/utils.py

@@ -20,12 +20,15 @@
 
 
 def get_match_query(
-    return_properties: Optional[list[str]] = None, retrieval_query: Optional[str] = None
+    return_properties: Optional[list[str]] = None,
+    retrieval_query: Optional[str] = None,
+    node_label: Optional[str] = None,
 ) -> str:
+    node_label_filter = f":`{node_label}`" if node_label else ""
     match_query = (
         "UNWIND $match_params AS match_param "
         "WITH match_param[0] AS match_id_value, match_param[1] AS score "
-        "MATCH (node) "
+        f"MATCH (node{node_label_filter}) "
         "WHERE node[$id_property] = match_id_value "
     )
     return match_query + get_query_tail(

src/neo4j_graphrag/retrievers/external/weaviate/types.py

@@ -57,6 +57,7 @@ class WeaviateNeo4jRetrieverModel(BaseModel):
     retrieval_query: Optional[str] = None
     result_formatter: Optional[Callable[[neo4j.Record], RetrieverResultItem]] = None
     neo4j_database: Optional[str] = None
+    node_label_neo4j: Optional[str] = None
 
 
 class WeaviateNeo4jSearchModel(VectorSearchModel):

src/neo4j_graphrag/retrievers/external/weaviate/weaviate.py

@@ -100,6 +100,7 @@ def __init__(
             Callable[[neo4j.Record], RetrieverResultItem]
         ] = None,
         neo4j_database: Optional[str] = None,
+        node_label_neo4j: Optional[str] = None,
     ):
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -116,12 +117,17 @@ def __init__(
                 retrieval_query=retrieval_query,
                 result_formatter=result_formatter,
                 neo4j_database=neo4j_database,
+                node_label_neo4j=node_label_neo4j,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors()) from e
 
         super().__init__(
-            driver, id_property_external, id_property_neo4j, neo4j_database
+            driver,
+            id_property_external,
+            id_property_neo4j,
+            neo4j_database,
+            node_label_neo4j,
         )
         self.client = validated_data.client_model.client
         collection = validated_data.collection
@@ -234,6 +240,7 @@ def get_search_results(
         search_query = get_match_query(
             return_properties=self.return_properties,
             retrieval_query=self.retrieval_query,
+            node_label=self.node_label_neo4j,
         )
 
         parameters = {

tests/unit/retrievers/external/test_weaviate.py

@@ -260,6 +260,21 @@ def test_match_query_with_both_return_properties_and_retrieval_query() -> None:
     assert match_query.strip() == expected.strip()
 
 
+def test_match_query_with_custom_node_label() -> None:
+    # Should ignore return_properties
+    match_query = get_match_query(
+        return_properties=["name", "age"], node_label="MyNodeLabel"
+    )
+    expected = (
+        "UNWIND $match_params AS match_param "
+        "WITH match_param[0] AS match_id_value, match_param[1] AS score "
+        "MATCH (node:`MyNodeLabel`) "
+        "WHERE node[$id_property] = match_id_value "
+        "RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score "
+    )
+    assert match_query.strip() == expected.strip()
+
+
 def test_weaviate_retriever_with_result_format_function(
     driver: MagicMock, neo4j_record: MagicMock, result_formatter: MagicMock
 ) -> None: