From ee66ce2b0da8ed5a0bd62a8c95743885d9ab5b82 Mon Sep 17 00:00:00 2001
From: Will Tai <wtaisen@gmail.com>
Date: Thu, 21 Nov 2024 16:35:51 +0000
Subject: [PATCH] Update deprecated Cypher query string for Delete query

---
 CHANGELOG.md                                  |  4 +++
 .../vectorstores/neo4j_vector.py              | 22 +++++++++++----
 .../docker-compose/neo4j.yml                  |  2 +-
 .../unit_tests/vectorstores/test_neo4j.py     | 28 +++++++++++++++++++
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 71cbe95..bf372d2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,3 +7,7 @@
 ### Added
 
 - Migrated all Neo4j-related code, including tests and integration tests, from the LangChain Community package to this package.
+
+Updated
+Improved subquery building for import queries to support different Neo4j versions.
+Added unit tests for Neo4j version-specific query construction.
\ No newline at end of file
diff --git a/libs/neo4j/langchain_neo4j/vectorstores/neo4j_vector.py b/libs/neo4j/langchain_neo4j/vectorstores/neo4j_vector.py
index 77b30ad..0932b0d 100644
--- a/libs/neo4j/langchain_neo4j/vectorstores/neo4j_vector.py
+++ b/libs/neo4j/langchain_neo4j/vectorstores/neo4j_vector.py
@@ -594,17 +594,29 @@ def __init__(
         if pre_delete_collection:
             from neo4j.exceptions import DatabaseError
 
-            self.query(
-                f"MATCH (n:`{self.node_label}`) "
-                "CALL (n) { DETACH DELETE n } "
-                "IN TRANSACTIONS OF 10000 ROWS;"
-            )
+            delete_query = self._build_delete_query()
+            self.query(delete_query)
             # Delete index
             try:
                 self.query(f"DROP INDEX {self.index_name}")
             except DatabaseError:  # Index didn't exist yet
                 pass
 
+    def _build_delete_query(self) -> str:
+        if self.neo4j_version_is_5_23_or_above:
+            query = (
+                f"MATCH (n:`{self.node_label}`) "
+                "CALL (n) { DETACH DELETE n } "
+                "IN TRANSACTIONS OF 10000 ROWS;"
+            )
+        else:
+            query = (
+                f"MATCH (n:`{self.node_label}`) "
+                "CALL { WITH n DETACH DELETE n } "
+                "IN TRANSACTIONS OF 10000 ROWS;"
+            )
+        return query
+
     def query(
         self,
         query: str,
diff --git a/libs/neo4j/tests/integration_tests/docker-compose/neo4j.yml b/libs/neo4j/tests/integration_tests/docker-compose/neo4j.yml
index eda9f8b..d1ca1d2 100644
--- a/libs/neo4j/tests/integration_tests/docker-compose/neo4j.yml
+++ b/libs/neo4j/tests/integration_tests/docker-compose/neo4j.yml
@@ -1,7 +1,7 @@
 version: "3.8"
 services:
   neo4j:
-    image: neo4j:5.24.2
+    image: neo4j:5.22
     restart: on-failure:0
     hostname: neo4j-test
     container_name: neo4j-test
diff --git a/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py b/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py
index 0d13d48..c6d2aed 100644
--- a/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py
+++ b/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py
@@ -189,3 +189,31 @@ def test_build_import_query_version_below_5_23(mock_vector_store: Neo4jVector) -
     actual_query = mock_vector_store._build_import_query()
 
     assert actual_query == expected_query
+
+
+def test_build_delete_query_version_is_or_above_5_23(
+    mock_vector_store: Neo4jVector,
+) -> None:
+    mock_vector_store.neo4j_version_is_5_23_or_above = True
+    expected_query = (
+        f"MATCH (n:`{mock_vector_store.node_label}`) "
+        "CALL (n) { DETACH DELETE n } "
+        "IN TRANSACTIONS OF 10000 ROWS;"
+    )
+
+    actual_query = mock_vector_store._build_delete_query()
+
+    assert actual_query == expected_query
+
+
+def test_build_delete_query_version_below_5_23(mock_vector_store: Neo4jVector) -> None:
+    mock_vector_store.neo4j_version_is_5_23_or_above = False
+    expected_query = (
+        f"MATCH (n:`{mock_vector_store.node_label}`) "
+        "CALL { WITH n DETACH DELETE n } "
+        "IN TRANSACTIONS OF 10000 ROWS;"
+    )
+
+    actual_query = mock_vector_store._build_delete_query()
+
+    assert actual_query == expected_query