Refactor id variables in neo4j queries

neo4j · Oct 25, 2024 · 846a6af · 846a6af
1 parent bc8540e
commit 846a6af
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/src/neo4j_graphrag/indexes.py b/src/neo4j_graphrag/indexes.py
@@ -288,7 +288,7 @@ def upsert_vector(
     """
     try:
         parameters = {
-            "id": node_id,
+            "node_id": node_id,
             "embedding_property": embedding_property,
             "vector": vector,
         }
@@ -344,7 +344,7 @@ def upsert_vector_on_relationship(
     """
     try:
         parameters = {
-            "id": rel_id,
+            "rel_id": rel_id,
             "embedding_property": embedding_property,
             "vector": vector,
         }
@@ -401,7 +401,7 @@ async def async_upsert_vector(
     """
     try:
         parameters = {
-            "id": node_id,
+            "node_id": node_id,
             "embedding_property": embedding_property,
             "vector": vector,
         }
@@ -458,7 +458,7 @@ async def async_upsert_vector_on_relationship(
     """
     try:
         parameters = {
-            "id": rel_id,
+            "rel_id": rel_id,
             "embedding_property": embedding_property,
             "vector": vector,
         }

diff --git a/src/neo4j_graphrag/neo4j_queries.py b/src/neo4j_graphrag/neo4j_queries.py
@@ -100,15 +100,15 @@
 
 UPSERT_VECTOR_ON_NODE_QUERY = (
     "MATCH (n) "
-    "WHERE elementId(n) = $id "
+    "WHERE elementId(n) = $node_id "
     "WITH n "
     "CALL db.create.setNodeVectorProperty(n, $embedding_property, $vector) "
     "RETURN n"
 )
 
 UPSERT_VECTOR_ON_RELATIONSHIP_QUERY = (
     "MATCH ()-[r]->() "
-    "WHERE elementId(r) = $id "
+    "WHERE elementId(r) = $rel_id "
     "WITH r "
     "CALL db.create.setRelationshipVectorProperty(r, $embedding_property, $vector) "
     "RETURN r"

diff --git a/tests/unit/test_indexes.py b/tests/unit/test_indexes.py
@@ -219,15 +219,15 @@ def test_upsert_vector_happy_path(driver: MagicMock) -> None:
 
     upsert_query = (
         "MATCH (n) "
-        "WHERE elementId(n) = $id "
+        "WHERE elementId(n) = $node_id "
         "WITH n "
         "CALL db.create.setNodeVectorProperty(n, $embedding_property, $vector) "
         "RETURN n"
     )
 
     driver.execute_query.assert_called_once_with(
         upsert_query,
-        {"id": id, "embedding_property": embedding_property, "vector": vector},
+        {"node_id": id, "embedding_property": embedding_property, "vector": vector},
         database_=None,
     )
 
@@ -241,15 +241,15 @@ def test_upsert_vector_on_relationship_happy_path(driver: MagicMock) -> None:
 
     upsert_query = (
         "MATCH ()-[r]->() "
-        "WHERE elementId(r) = $id "
+        "WHERE elementId(r) = $rel_id "
         "WITH r "
         "CALL db.create.setRelationshipVectorProperty(r, $embedding_property, $vector) "
         "RETURN r"
     )
 
     driver.execute_query.assert_called_once_with(
         upsert_query,
-        {"id": id, "embedding_property": embedding_property, "vector": vector},
+        {"rel_id": id, "embedding_property": embedding_property, "vector": vector},
         database_=None,
     )