External retriever node label (#451)

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

* Add node_label_neo4j parameter to Qdrant and Pinecone retrievers

* Update doc + do not force escaping the node label to allow node label expression

* Update unit test

Author: stellasia

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## Next
 
+### Added
+
+- Added an optional `node_label_neo4j` parameter in the external retrievers to speed up the search query in Neo4j.
+
+
 ## 1.10.1
 
 ### Added

docs/source/user_guide_rag.rst

@@ -762,6 +762,7 @@ Weaviate Retrievers
         collection="Movies",
         id_property_external="neo4j_id",
         id_property_neo4j="id",
+        node_label_neo4j="Document",  # optional
     )
 
 Internally, this retriever performs the vector search in Weaviate, finds the corresponding node by matching
@@ -795,6 +796,7 @@ Pinecone Retrievers
         index_name="Movies",
         id_property_neo4j="id",
         embedder=embedder,
+        node_label_neo4j="Document",  # optional
     )
 
 Also see :ref:`pineconeneo4jretriever`.
@@ -825,6 +827,7 @@ Qdrant Retrievers
         id_property_external="neo4j_id",    # The payload field that contains identifier to a corresponding Neo4j node id property
         id_property_neo4j="id",
         embedder=embedder,
+        node_label_neo4j="Document",  # optional
     )
 
 See :ref:`qdrantneo4jretriever`.

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/pinecone/pinecone.py

@@ -83,6 +83,7 @@ class PineconeNeo4jRetriever(ExternalRetriever):
         retrieval_query (str): Cypher query that gets appended.
         result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Function to transform a neo4j.Record to a RetrieverResultItem.
         neo4j_database (Optional[str]): The name of the Neo4j database. If not provided, this defaults to the server's default database ("neo4j" by default) (`see reference to documentation <https://neo4j.com/docs/operations-manual/current/database-administration/#manage-databases-default>`_).
+        node_label_neo4j (Optional[str]): The label of the Neo4j node to retrieve. This label must be properly escaped if needed, eg "`Label with spaces`".
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -101,6 +102,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,6 +118,7 @@ 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
@@ -125,6 +128,7 @@ def __init__(
             id_property_external="id",
             id_property_neo4j=validated_data.id_property_neo4j,
             neo4j_database=neo4j_database,
+            node_label_neo4j=node_label_neo4j,
         )
         self.driver = validated_data.driver_model.driver
         self.client = validated_data.client_model.client
@@ -172,7 +176,8 @@ def get_search_results(
                   driver=neo4j_driver,
                   client=pc_client,
                   index_name="jeopardy",
-                  id_property_neo4j="id"
+                  id_property_neo4j="id",
+                  node_label_neo4j="Document",
               )
               biology_embedding = ...
               retriever.search(query_vector=biology_embedding, top_k=2)
@@ -223,6 +228,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 = {

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

@@ -59,3 +59,4 @@ class PineconeNeo4jRetrieverModel(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

src/neo4j_graphrag/retrievers/external/qdrant/qdrant.py

@@ -79,6 +79,7 @@ class QdrantNeo4jRetriever(ExternalRetriever):
         return_properties (Optional[list[str]]): List of node properties to return.
         result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Function to transform a neo4j.Record to a RetrieverResultItem.
         neo4j_database (Optional[str]): The name of the Neo4j database. If not provided, this defaults to the server's default database ("neo4j" by default) (`see reference to documentation <https://neo4j.com/docs/operations-manual/current/database-administration/#manage-databases-default>`_).
+        node_label_neo4j (Optional[str]): The label of the Neo4j node to retrieve. This label must be properly escaped if needed, eg "`Label with spaces`".
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -99,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,6 +118,7 @@ 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
@@ -125,6 +128,7 @@ def __init__(
             id_property_external=validated_data.id_property_external,
             id_property_neo4j=validated_data.id_property_neo4j,
             neo4j_database=neo4j_database,
+            node_label_neo4j=node_label_neo4j,
         )
         self.driver = validated_data.driver_model.driver
         self.client = validated_data.client_model.client
@@ -169,7 +173,8 @@ def get_search_results(
                     driver=neo4j_driver,
                     client=client,
                     collection_name="my_collection",
-                    id_property_external="neo4j_id"
+                    id_property_external="neo4j_id",
+                    node_label_neo4j="Document",
                 )
                 embedding = ...
                 retriever.search(query_vector=embedding, top_k=2)
@@ -223,6 +228,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 = {

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

@@ -54,3 +54,4 @@ class QdrantNeo4jRetrieverModel(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

src/neo4j_graphrag/retrievers/external/utils.py

@@ -20,12 +20,18 @@
 
 
 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 is not escaped on purpose, allowing users to use any valid
+    # node label expression, e.g. "Actor|Director". It's up to the user to ensure
+    # labels are properly escaped, i.e. "`My label with space`".
+    node_label_expression = 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_expression}) "
         "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

@@ -81,6 +81,7 @@ class WeaviateNeo4jRetriever(ExternalRetriever):
         return_properties (Optional[list[str]]): List of node properties to return.
         result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Function to transform a neo4j.Record to a RetrieverResultItem.
         neo4j_database (Optional[str]): The name of the Neo4j database. If not provided, this defaults to the server's default database ("neo4j" by default) (`see reference to documentation <https://neo4j.com/docs/operations-manual/current/database-administration/#manage-databases-default>`_).
+        node_label_neo4j (Optional[str]): The label of the Neo4j node to retrieve. This label must be properly escaped if needed, eg "`Label with spaces`".
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -100,6 +101,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 +118,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
@@ -164,6 +171,7 @@ def get_search_results(
               collection="Jeopardy",
               id_property_external="neo4j_id",
               id_property_neo4j="id",
+              node_label_neo4j="Document",
           )
 
           biology_embedding = ...
@@ -234,6 +242,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 = {