Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor: Use Neo4j transactional functions.

Browse Source 
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.
This commit is contained in:
Michael J. Simons
2024-11-01 07:40:14 +01:00
committed by Christian Tzolov
parent cd15b65c7a
commit 560315ccfc
1 changed files with 22 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java
View File

@@ -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();
}
}