Add node_label_neo4j parameter to Qdrant and Pinecone retrievers

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

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.
 
     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.
 
     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/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.
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -170,6 +171,7 @@ def get_search_results(
               collection="Jeopardy",
               id_property_external="neo4j_id",
               id_property_neo4j="id",
+              node_label_neo4j="Document",
           )
 
           biology_embedding = ...

tests/e2e/pinecone_e2e/test_pinecone_e2e.py

@@ -52,7 +52,11 @@ def populate_neo4j_db(driver: MagicMock) -> None:
 @pytest.mark.usefixtures("populate_neo4j_db")
 def test_pinecone_neo4j_vector_input(driver: MagicMock, client: MagicMock) -> None:
     retriever = PineconeNeo4jRetriever(
-        driver=driver, client=client, index_name="jeopardy", id_property_neo4j="id"
+        driver=driver,
+        client=client,
+        index_name="jeopardy",
+        id_property_neo4j="id",
+        node_label_neo4j="Question",
     )
     with mock.patch.object(retriever, "index") as mock_index:
         top_k = 2

tests/e2e/qdrant_e2e/test_qdrant_e2e.py

@@ -59,6 +59,7 @@ def test_qdrant_neo4j_vector_input(driver: Driver, qdrant_client: QdrantClient)
         collection_name="Jeopardy",
         id_property_external="neo4j_id",
         id_property_neo4j="id",
+        node_label_neo4j="Question",
     )
 
     top_k = 1

tests/e2e/weaviate_e2e/test_weaviate_e2e.py

@@ -63,6 +63,7 @@ def test_weaviate_neo4j_vector_input(
         collection="Jeopardy",
         id_property_external="neo4j_id",
         id_property_neo4j="id",
+        node_label_neo4j="Question",
     )
 
     top_k = 2

tests/unit/retrievers/external/test_weaviate.py

@@ -261,7 +261,6 @@ def test_match_query_with_both_return_properties_and_retrieval_query() -> None:
 
 
 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"
     )