From 560315ccfc4ac823d2d7c3cc2a1983e22e58992e Mon Sep 17 00:00:00 2001 From: "Michael J. Simons" Date: Fri, 1 Nov 2024 07:40:14 +0100 Subject: [PATCH] refactor: Use Neo4j transactional functions. This change uses `session.executeWrite` and `session.executeRead` to interact with the database. Those methods use build-in retries for several Neo4j specific error states. Doing will improve the overall user experience in this case as the no other external retry mechanism should be used. Doing so is fine, as the Neo4j vector store is not subject to Springs general transaction management. --- .../ai/vectorstore/Neo4jVectorStore.java | 38 +++++++++++-------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java index 21f2e8c8a..702dc9f56 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java @@ -119,7 +119,9 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements WITH row, u CALL db.create.setNodeVectorProperty(u, $embeddingProperty, row.embedding) """.formatted(this.config.label, this.config.idProperty); - session.run(statement, Map.of("rows", rows, "embeddingProperty", this.config.embeddingProperty)).consume(); + session.executeWrite( + tx -> tx.run(statement, Map.of("rows", rows, "embeddingProperty", this.config.embeddingProperty)) + .consume()); } } @@ -128,6 +130,8 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements try (var session = this.driver.session(this.config.sessionConfig)) { + // Those queries with internal, cypher based transaction management cannot be + // run with executeWrite var summary = session .run(""" MATCH (n:%s) WHERE n.%s IN $ids @@ -158,10 +162,10 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements WHERE %s RETURN node, score""".formatted(condition); - return session + return session.executeRead(tx -> tx .run(query, Map.of("indexName", this.config.indexNameNotSanitized, "numberOfNearestNeighbours", request.getTopK(), "embeddingValue", embedding, "threshold", request.getSimilarityThreshold())) - .list(this::recordToDocument); + .list(this::recordToDocument)); } } @@ -174,20 +178,22 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements try (var session = this.driver.session(this.config.sessionConfig)) { - session - .run("CREATE CONSTRAINT %s IF NOT EXISTS FOR (n:%s) REQUIRE n.%s IS UNIQUE" - .formatted(this.config.constraintName, this.config.label, this.config.idProperty)) - .consume(); + session.executeWriteWithoutResult(tx -> { + tx.run("CREATE CONSTRAINT %s IF NOT EXISTS FOR (n:%s) REQUIRE n.%s IS UNIQUE" + .formatted(this.config.constraintName, this.config.label, this.config.idProperty)).consume(); - var statement = """ - CREATE VECTOR INDEX %s IF NOT EXISTS FOR (n:%s) ON (n.%s) - OPTIONS {indexConfig: { - `vector.dimensions`: %d, - `vector.similarity_function`: '%s' - }} - """.formatted(this.config.indexName, this.config.label, this.config.embeddingProperty, - this.config.embeddingDimension, this.config.distanceType.name); - session.run(statement).consume(); + var statement = """ + CREATE VECTOR INDEX %s IF NOT EXISTS FOR (n:%s) ON (n.%s) + OPTIONS {indexConfig: { + `vector.dimensions`: %d, + `vector.similarity_function`: '%s' + }} + """.formatted(this.config.indexName, this.config.label, this.config.embeddingProperty, + this.config.embeddingDimension, this.config.distanceType.name); + tx.run(statement).consume(); + }); + + // Bad idea to retry this... session.run("CALL db.awaitIndexes()").consume(); } }