Refactors upsert vector function (#272)

* Added upsert_embeddings function

* Added tests for upsert_embeddings

* Updated docs and README

* Added deprecation warnings to docstrings

* Renamed function

* Small variable name changes

* Formatting fixes

* Added deprecation comments

* Updated CHANGELOG

* Updated upsert_vectors docstring

* Fixed README issue

* Updated docs

Author: alexthomas93

CHANGELOG.md

@@ -6,10 +6,13 @@
 
 - Utility functions to retrieve metadata for vector and full-text indexes.
 - Support for effective_search_ratio parameter in vector and hybrid searches.
+- Introduced upsert_vectors utility function for batch upserting embeddings to vector indexes.
 
 ### Changed
 
 - Refactored index-related functions for improved compatibility and functionality.
+- Added deprecation warnings to upsert_vector, upsert_vector_on_relationship functions in favor of upsert_vectors.
+- Added deprecation warnings to async_upsert_vector, async_upsert_vector_on_relationship functions notifying developers that they will be removed in a future release.
 
 ## 1.4.3
 

README.md

@@ -194,7 +194,8 @@ Ensure that your vector index is created prior to executing this example.
 ```python
 from neo4j import GraphDatabase
 from neo4j_graphrag.embeddings import OpenAIEmbeddings
-from neo4j_graphrag.indexes import upsert_vector
+from neo4j_graphrag.indexes import upsert_vectors
+from neo4j_graphrag.types import EntityType
 
 NEO4J_URI = "neo4j://localhost:7687"
 NEO4J_USERNAME = "neo4j"
@@ -214,11 +215,12 @@ text = (
 vector = embedder.embed_query(text)
 
 # Upsert the vector
-upsert_vector(
+upsert_vectors(
     driver,
-    node_id=0,
-    embedding_property="embedding",
-    vector=vector,
+    ids=["1234"],
+    embedding_property="vectorProperty",
+    embeddings=[vector],
+    entity_type=EntityType.NODE,
 )
 driver.close()
 ```

docs/source/api.rst

@@ -382,6 +382,8 @@ Database Interaction
 
 .. autofunction:: neo4j_graphrag.indexes.drop_index_if_exists
 
+.. autofunction:: neo4j_graphrag.indexes.upsert_vectors
+
 .. autofunction:: neo4j_graphrag.indexes.upsert_vector
 
 .. autofunction:: neo4j_graphrag.indexes.upsert_vector_on_relationship

docs/source/index.rst

@@ -148,9 +148,9 @@ Note that the below example is not the only way you can upsert data into your Ne
 
 
 .. code:: python
-
     from neo4j import GraphDatabase
-    from neo4j_graphrag.indexes import upsert_vector
+    from neo4j_graphrag.indexes import upsert_vectors
+    from neo4j_graphrag.types import EntityType
 
     URI = "neo4j://localhost:7687"
     AUTH = ("neo4j", "password")
@@ -160,11 +160,12 @@ Note that the below example is not the only way you can upsert data into your Ne
 
     # Upsert the vector
     vector = ...
-    upsert_vector(
+    upsert_vectors(
         driver,
-        node_id=1,
+        ids=["1234"],
         embedding_property="vectorProperty",
-        vector=vector,
+        embeddings=[vector],
+        entity_type=EntityType.NODE,
     )
 
 

docs/source/user_guide_rag.rst

@@ -917,9 +917,11 @@ Populate a Vector Index
 ==========================
 
 .. code:: python
+    from random import random
 
     from neo4j import GraphDatabase
-    from random import random
+    from neo4j_graphrag.indexes import upsert_vectors
+    from neo4j_graphrag.types import EntityType
 
     URI = "neo4j://localhost:7687"
     AUTH = ("neo4j", "password")
@@ -928,10 +930,17 @@ Populate a Vector Index
     driver = GraphDatabase.driver(URI, auth=AUTH)
 
     # Upsert the vector
+    DIMENSION = 1536
     vector = [random() for _ in range(DIMENSION)]
-    upsert_vector(driver, node_id="1234", embedding_property="embedding", vector=vector)
+    upsert_vectors(
+        driver,
+        ids=["1234"],
+        embedding_property="vectorProperty",
+        embeddings=[vector],
+        entity_type=EntityType.NODE,
+    )
 
-This will update the node with `id(node)=1234` to add (or update) a `node.embedding` property.
+This will update the node with `id(node)=1234` to add (or update) a `node.vectorProperty` property.
 This property will also be added to the vector index.
 
 

src/neo4j_graphrag/indexes.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import logging
+import warnings
 from typing import List, Literal, Optional
 
 import neo4j
@@ -23,10 +24,12 @@
 from neo4j_graphrag.neo4j_queries import (
     UPSERT_VECTOR_ON_NODE_QUERY,
     UPSERT_VECTOR_ON_RELATIONSHIP_QUERY,
+    UPSERT_VECTORS_ON_NODE_QUERY,
+    UPSERT_VECTORS_ON_RELATIONSHIP_QUERY,
 )
 
 from .exceptions import Neo4jIndexError, Neo4jInsertionError
-from .types import FulltextIndexModel, VectorIndexModel
+from .types import EntityType, FulltextIndexModel, VectorIndexModel
 
 logger = logging.getLogger(__name__)
 
@@ -245,6 +248,86 @@ def drop_index_if_exists(
         raise Neo4jIndexError(f"Dropping Neo4j index failed: {e.message}") from e
 
 
+def upsert_vectors(
+    driver: neo4j.Driver,
+    ids: List[str],
+    embedding_property: str,
+    embeddings: List[List[float]],
+    neo4j_database: Optional[str] = None,
+    entity_type: EntityType = EntityType.NODE,
+) -> None:
+    """
+    This method constructs a Cypher query and executes it to upsert
+    (insert or update) embeddings on a set of nodes or relationships.
+
+    Example:
+
+    .. code-block:: python
+
+        from neo4j import GraphDatabase
+        from neo4j_graphrag.indexes import upsert_vectors
+
+        URI = "neo4j://localhost:7687"
+        AUTH = ("neo4j", "password")
+
+        # Connect to Neo4j database
+        driver = GraphDatabase.driver(URI, auth=AUTH)
+
+        # Upsert embeddings data for several nodes
+        upsert_vectors(
+            driver,
+            ids=['123', '456', '789'],
+            embedding_property="vectorProperty",
+            embeddings=[
+                [0.12, 0.34, 0.56],
+                [0.78, 0.90, 0.12],
+                [0.34, 0.56, 0.78],
+            ],
+            neo4j_database="neo4j",
+            entity_type='NODE',
+        )
+
+    Args:
+        driver (neo4j.Driver): Neo4j Python driver instance.
+        ids (List[int]): The element IDs of the nodes or relationships.
+        embedding_property (str): The name of the property to store the vectors in.
+        embeddings (List[List[float]]): The list of vectors to store, one per ID.
+        neo4j_database (Optional[str]): The name of the Neo4j database.
+            If not provided, defaults to the server's default database. 'neo4j' by default.
+        entity_type (EntityType): Specifies whether to upsert to nodes ('NODE') or relationships ('RELATIONSHIP').
+            Defaults to 'NODE'.
+
+    Raises:
+        ValueError: If the lengths of IDs and embeddings do not match, or if embeddings are not of uniform dimension.
+        Neo4jInsertionError: If an error occurs while attempting to upsert the vectors in Neo4j.
+    """
+    if entity_type == EntityType.NODE:
+        query = UPSERT_VECTORS_ON_NODE_QUERY
+    elif entity_type == EntityType.RELATIONSHIP:
+        query = UPSERT_VECTORS_ON_RELATIONSHIP_QUERY
+    else:
+        raise ValueError("entity_type must be either 'NODE' or 'RELATIONSHIP'")
+    if len(ids) != len(embeddings):
+        raise ValueError("ids and embeddings must be the same length")
+    if not all(len(embedding) == len(embeddings[0]) for embedding in embeddings):
+        raise ValueError("All embeddings must be of the same size")
+    try:
+        parameters = {
+            "rows": [
+                {"id": id, "embedding": embedding}
+                for id, embedding in zip(ids, embeddings)
+            ],
+            "embedding_property": embedding_property,
+        }
+        driver.execute_query(
+            query_=query, parameters_=parameters, database_=neo4j_database
+        )
+    except neo4j.exceptions.ClientError as e:
+        raise Neo4jInsertionError(
+            f"Upserting vectors to Neo4j failed: {e.message}"
+        ) from e
+
+
 def upsert_vector(
     driver: neo4j.Driver,
     node_id: int,
@@ -253,6 +336,9 @@ def upsert_vector(
     neo4j_database: Optional[str] = None,
 ) -> None:
     """
+    .. warning::
+        'upsert_vector' is deprecated and will be removed in a future version, please use 'upsert_vectors' instead.
+
     This method constructs a Cypher query and executes it to upsert (insert or update) a vector property on a specific node.
 
     Example:
@@ -286,6 +372,11 @@ def upsert_vector(
     Raises:
         Neo4jInsertionError: If upserting of the vector fails.
     """
+    warnings.warn(
+        "'upsert_vector' is deprecated and will be removed in a future version, please use 'upsert_vectors' instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     try:
         parameters = {
             "node_element_id": node_id,
@@ -309,6 +400,9 @@ def upsert_vector_on_relationship(
     neo4j_database: Optional[str] = None,
 ) -> None:
     """
+    .. warning::
+        'upsert_vector_on_relationship' is deprecated and will be removed in a future version, please use 'upsert_vectors' instead.
+
     This method constructs a Cypher query and executes it to upsert (insert or update) a vector property on a specific relationship.
 
     Example:
@@ -342,6 +436,11 @@ def upsert_vector_on_relationship(
     Raises:
         Neo4jInsertionError: If upserting of the vector fails.
     """
+    warnings.warn(
+        "'upsert_vector_on_relationship' is deprecated and will be removed in a future version, please use 'upsert_vectors' instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     try:
         parameters = {
             "rel_element_id": rel_id,
@@ -365,6 +464,9 @@ async def async_upsert_vector(
     neo4j_database: Optional[str] = None,
 ) -> None:
     """
+    .. warning::
+        'async_upsert_vector' is deprecated and will be removed in a future version.
+
     This method constructs a Cypher query and asynchronously executes it
     to upsert (insert or update) a vector property on a specific node.
 
@@ -399,6 +501,11 @@ async def async_upsert_vector(
     Raises:
         Neo4jInsertionError: If upserting of the vector fails.
     """
+    warnings.warn(
+        "'async_upsert_vector' is deprecated and will be removed in a future version.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     try:
         parameters = {
             "node_id": node_id,
@@ -422,6 +529,9 @@ async def async_upsert_vector_on_relationship(
     neo4j_database: Optional[str] = None,
 ) -> None:
     """
+    .. warning::
+        'async_upsert_vector_on_relationship' is deprecated and will be removed in a future version.
+
     This method constructs a Cypher query and asynchronously executes it
     to upsert (insert or update) a vector property on a specific relationship.
 
@@ -456,6 +566,11 @@ async def async_upsert_vector_on_relationship(
     Raises:
         Neo4jInsertionError: If upserting of the vector fails.
     """
+    warnings.warn(
+        "'async_upsert_vector_on_relationship' is deprecated and will be removed in a future version.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     try:
         parameters = {
             "rel_id": rel_id,

src/neo4j_graphrag/neo4j_queries.py

@@ -107,6 +107,7 @@
     "RETURN elementId(rel)"
 )
 
+# Deprecated, remove along with upsert_vector
 UPSERT_VECTOR_ON_NODE_QUERY = (
     "MATCH (n) "
     "WHERE elementId(n) = $node_element_id "
@@ -115,6 +116,16 @@
     "RETURN n"
 )
 
+UPSERT_VECTORS_ON_NODE_QUERY = (
+    "UNWIND $rows AS row "
+    "MATCH (n) "
+    "WHERE elementId(n) = row.id "
+    "WITH n, row "
+    "CALL db.create.setNodeVectorProperty(n, $embedding_property, row.embedding) "
+    "RETURN n"
+)
+
+# Deprecated, remove along with upsert_vector_on_relationship
 UPSERT_VECTOR_ON_RELATIONSHIP_QUERY = (
     "MATCH ()-[r]->() "
     "WHERE elementId(r) = $rel_element_id "
@@ -123,6 +134,15 @@
     "RETURN r"
 )
 
+UPSERT_VECTORS_ON_RELATIONSHIP_QUERY = (
+    "UNWIND $rows AS row "
+    "MATCH ()-[r]->() "
+    "WHERE elementId(r) = row.id "
+    "WITH r, row "
+    "CALL db.create.setRelationshipVectorProperty(r, $embedding_property, row.embedding) "
+    "RETURN r"
+)
+
 
 def _get_hybrid_query(neo4j_version_is_5_23_or_above: bool) -> str:
     """