Adds an upsert vector function (neo4j#60)

* Added an upsert_vector function

* Added test_upsert_vector unit tests

* Added mypy type hints back in

* Removed node_label from upsert_vector

* Added Neo4jInsertionError to api.rst docs

* Updated CHANGELOG.md

Author: alexthomas93

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Added
 -   Stopped embeddings from being returned when searching with `VectorRetriever`. Added `nodeLabels` and `id` to the metadata of `VectorRetriever` results.
+-   Added `upsert_vector` utility function for attaching vectors to node properties.
+-   Introduced `Neo4jInsertionError` for handling insertion failures in Neo4j.
 
 ## 0.2.0
 

docs/source/api.rst

@@ -139,6 +139,13 @@ Neo4jIndexError
    :show-inheritance:
 
 
+Neo4jInsertionError
+===============
+
+.. autoclass:: neo4j_genai.exceptions.Neo4jInsertionError
+   :show-inheritance:
+
+
 Neo4jVersionError
 =================
 

src/neo4j_genai/exceptions.py

@@ -61,6 +61,12 @@ class Neo4jIndexError(Neo4jGenAiError):
     pass
 
 
+class Neo4jInsertionError(Neo4jGenAiError):
+    """Exception raised when inserting data into the Neo4j database fails."""
+
+    pass
+
+
 class Neo4jVersionError(Neo4jGenAiError):
     """Exception raised when Neo4j version does not meet minimum requirements."""
 

src/neo4j_genai/indexes.py

@@ -13,13 +13,14 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import logging
+from typing import Literal
+
 import neo4j
 from pydantic import ValidationError
 
-from .exceptions import Neo4jIndexError
-from .types import VectorIndexModel, FulltextIndexModel
-import logging
-from typing import Literal
+from .exceptions import Neo4jIndexError, Neo4jInsertionError
+from .types import FulltextIndexModel, VectorIndexModel
 
 logger = logging.getLogger(__name__)
 
@@ -143,3 +144,39 @@ def drop_index_if_exists(driver: neo4j.Driver, name: str) -> None:
         driver.execute_query(query, parameters)
     except neo4j.exceptions.ClientError as e:
         raise Neo4jIndexError(f"Dropping Neo4j index failed: {e}")
+
+
+def upsert_vector(
+    driver: neo4j.Driver,
+    node_id: int,
+    vector_prop: str,
+    vector: list[float],
+) -> None:
+    """
+    This method constructs a Cypher query and executes it to upsert (insert or update) a vector property on a specific node.
+
+    Args:
+        driver (neo4j.Driver): Neo4j Python driver instance.
+        node_id (int): The id of the node.
+        vector_prop (str): The name of the property to store the vector in.
+        vector (list[float]): The vector to store.
+
+    Raises:
+        Neo4jInsertError: If upserting of the vector fails.
+    """
+    try:
+        query = """
+        MATCH (n)
+        WHERE elementId(n) = $id
+        WITH n
+        CALL db.create.setNodeVectorProperty(n, $vector_prop, $vector)
+        RETURN n
+        """
+        parameters = {
+            "id": node_id,
+            "vector_prop": vector_prop,
+            "vector": vector,
+        }
+        driver.execute_query(query, parameters)
+    except neo4j.exceptions.ClientError as e:
+        raise Neo4jInsertionError(f"Upserting vector to Neo4j failed: {e}")

tests/unit/test_indexes.py

@@ -12,15 +12,16 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+from unittest.mock import MagicMock
+
 import neo4j.exceptions
 import pytest
-
-from neo4j_genai.exceptions import Neo4jIndexError
-from unittest.mock import MagicMock
+from neo4j_genai.exceptions import Neo4jIndexError, Neo4jInsertionError
 from neo4j_genai.indexes import (
+    create_fulltext_index,
     create_vector_index,
     drop_index_if_exists,
-    create_fulltext_index,
+    upsert_vector,
 )
 
 
@@ -161,3 +162,38 @@ def test_create_fulltext_index_ensure_escaping(driver: MagicMock) -> None:
         create_query,
         {"name": "my-complicated-`-index"},
     )
+
+
+def test_upsert_vector_happy_path(driver: MagicMock) -> None:
+    id = 1
+    vector_prop = "embedding"
+    vector = [1.0, 2.0, 3.0]
+
+    upsert_vector(driver, id, vector_prop, vector)
+
+    upsert_query = """
+        MATCH (n)
+        WHERE elementId(n) = $id
+        WITH n
+        CALL db.create.setNodeVectorProperty(n, $vector_prop, $vector)
+        RETURN n
+        """
+
+    driver.execute_query.assert_called_once_with(
+        upsert_query,
+        {"id": id, "vector_prop": vector_prop, "vector": vector},
+    )
+
+
+def test_upsert_vector_raises_error_with_neo4j_insertion_error(
+    driver: MagicMock,
+) -> None:
+    id = 1
+    vector_prop = "embedding"
+    vector = [1.0, 2.0, 3.0]
+    driver.execute_query.side_effect = neo4j.exceptions.ClientError
+
+    with pytest.raises(Neo4jInsertionError) as excinfo:
+        upsert_vector(driver, id, vector_prop, vector)
+
+    assert "Upserting vector to Neo4j failed" in str(excinfo)