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

Allow more customization for Neo4j store (id and constraint).

Browse Source 
Unrelated to this change, the Neo4j test version increased to be current.
This commit is contained in:
Gerrit Meier
2024-03-14 22:23:26 +01:00
committed by Christian Tzolov
parent 152420fcc4
commit 7c6bbef3ec
4 changed files with 93 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java
View File

@@ -45,6 +45,8 @@ public class Neo4jVectorStoreAutoConfiguration {
.withLabel(properties.getLabel())
.withEmbeddingProperty(properties.getEmbeddingProperty())
.withIndexName(properties.getIndexName())
.withIdProperty(properties.getIdProperty())
.withConstraintName(properties.getConstraintName())
.build();
return new Neo4jVectorStore(driver, embeddingClient, config);

20
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java
View File

@@ -38,6 +38,10 @@ public class Neo4jVectorStoreProperties {
private String indexName = Neo4jVectorStore.DEFAULT_INDEX_NAME;
private String idProperty = Neo4jVectorStore.DEFAULT_ID_PROPERTY;
private String constraintName = Neo4jVectorStore.DEFAULT_CONSTRAINT_NAME;
public String getDatabaseName() {
return this.databaseName;
}
@@ -86,4 +90,20 @@ public class Neo4jVectorStoreProperties {
this.indexName = indexName;
}
public String getIdProperty() {
return this.idProperty;
}
public void setIdProperty(String idProperty) {
this.idProperty = idProperty;
}
public String getConstraintName() {
return this.constraintName;
}
public void setConstraintName(String constraintName) {
this.constraintName = constraintName;
}
}

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

@@ -63,14 +63,19 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
private final Neo4jDistanceType distanceType;
private final String label;
private final String embeddingProperty;
private final String quotedLabel;
private final String label;
private final String indexName;
// needed for similarity search call
private final String indexNameNotSanitized;
private final String idProperty;
private final String constraintName;
/**
* Start building a new configuration.
* @return The entry point for creating a new configuration.
@@ -96,10 +101,12 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
.orElseGet(SessionConfig::defaultConfig);
this.embeddingDimension = builder.embeddingDimension;
this.distanceType = builder.distanceType;
this.label = builder.label;
this.embeddingProperty = builder.embeddingProperty;
this.quotedLabel = SchemaNames.sanitize(this.label).orElseThrow();
this.indexName = builder.indexName;
this.embeddingProperty = SchemaNames.sanitize(builder.embeddingProperty).orElseThrow();
this.label = SchemaNames.sanitize(builder.label).orElseThrow();
this.indexNameNotSanitized = builder.indexName;
this.indexName = SchemaNames.sanitize(builder.indexName, true).orElseThrow();
this.constraintName = SchemaNames.sanitize(builder.constraintName).orElseThrow();
this.idProperty = SchemaNames.sanitize(builder.idProperty).orElseThrow();
}
public static class Builder {
@@ -116,6 +123,10 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
private String indexName = DEFAULT_INDEX_NAME;
private String idProperty = DEFAULT_ID_PROPERTY;
private String constraintName = DEFAULT_CONSTRAINT_NAME;
private Builder() {
}
@@ -202,6 +213,35 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
return this;
}
/**
* Configures the id property to be used. Defaults to {@literal id}.
* @param newIdProperty The name of the id property of the {@link Document}
* entity
* @return this builder
*/
public Builder withIdProperty(String newIdProperty) {
Assert.hasText(newIdProperty, "Id property may not be null or blank");
this.idProperty = newIdProperty;
return this;
}
/**
* Configures the constraint name to be used. Defaults to
* {@literal Document_unique_idx}.
* @param newConstraintName The name of the unique constraint for the id
* property.
* @return this builder
*/
public Builder withConstraintName(String newConstraintName) {
Assert.hasText(newConstraintName, "Constraint name may not be null or blank");
this.constraintName = newConstraintName;
return this;
}
/**
* {@return the immutable configuration}
*/
@@ -222,6 +262,10 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
public static final String DEFAULT_EMBEDDING_PROPERTY = "embedding";
public static final String DEFAULT_ID_PROPERTY = "id";
public static final String DEFAULT_CONSTRAINT_NAME = DEFAULT_LABEL + "_unique_idx";
private final Neo4jVectorFilterExpressionConverter filterExpressionConverter = new Neo4jVectorFilterExpressionConverter();
private final Driver driver;
@@ -249,16 +293,16 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
try (var session = this.driver.session()) {
var statement = """
UNWIND $rows AS row
MERGE (u:%s {id: row.id})
MERGE (u:%s {%2$s: row.id})
ON CREATE
SET u += row.properties
ON MATCH
SET u = {}
SET u.id = row.id,
SET u.%2$s = row.id,
u += row.properties
WITH row, u
CALL db.create.setNodeVectorProperty(u, $embeddingProperty, row.embedding)
""".formatted(this.config.quotedLabel);
""".formatted(this.config.label, this.config.idProperty);
session.run(statement, Map.of("rows", rows, "embeddingProperty", this.config.embeddingProperty)).consume();
}
}
@@ -268,10 +312,12 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
try (var session = this.driver.session(this.config.sessionConfig)) {
var summary = session.run("""
MATCH (n:%s) WHERE n.id IN $ids
CALL { WITH n DETACH DELETE n } IN TRANSACTIONS OF $transactionSize ROWS
""".formatted(this.config.quotedLabel), Map.of("ids", idList, "transactionSize", 10_000))
var summary = session
.run("""
MATCH (n:%s) WHERE n.%s IN $ids
CALL { WITH n DETACH DELETE n } IN TRANSACTIONS OF $transactionSize ROWS
""".formatted(this.config.label, this.config.idProperty),
Map.of("ids", idList, "transactionSize", 10_000))
.consume();
return Optional.of(idList.size() == summary.counters().nodesDeleted());
}
@@ -297,10 +343,9 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
RETURN node, score""".formatted(condition);
return session
.run(query,
Map.of("indexName", this.config.indexName, "numberOfNearestNeighbours", request.getTopK(),
"embeddingValue", embedding, "threshold", request.getSimilarityThreshold()))
.list(Neo4jVectorStore::recordToDocument);
.run(query, Map.of("indexName", this.config.indexNameNotSanitized, "numberOfNearestNeighbours",
request.getTopK(), "embeddingValue", embedding, "threshold", request.getSimilarityThreshold()))
.list(this::recordToDocument);
}
}
@@ -310,8 +355,8 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
try (var session = this.driver.session(this.config.sessionConfig)) {
session
.run("CREATE CONSTRAINT %s IF NOT EXISTS FOR (n:%s) REQUIRE n.id IS UNIQUE".formatted(
SchemaNames.sanitize(this.config.label + "_unique_idx").orElseThrow(), this.config.quotedLabel))
.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 = """
@@ -320,9 +365,8 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
`vector.dimensions`: %d,
`vector.similarity_function`: '%s'
}}
""".formatted(SchemaNames.sanitize(this.config.indexName, true).orElseThrow(),
this.config.quotedLabel, this.config.embeddingProperty, this.config.embeddingDimension,
this.config.distanceType.name);
""".formatted(this.config.indexName, this.config.label, this.config.embeddingProperty,
this.config.embeddingDimension, this.config.distanceType.name);
session.run(statement).consume();
session.run("CALL db.awaitIndexes()").consume();
}
@@ -355,7 +399,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
return embeddingFloat;
}
private static Document recordToDocument(org.neo4j.driver.Record neoRecord) {
private Document recordToDocument(org.neo4j.driver.Record neoRecord) {
var node = neoRecord.get("node").asNode();
var score = neoRecord.get("score").asFloat();
var metaData = new HashMap<String, Object>();
@@ -366,7 +410,8 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
}
});
return new Document(node.get("id").asString(), node.get("text").asString(), Map.copyOf(metaData));
return new Document(node.get(this.config.idProperty).asString(), node.get("text").asString(),
Map.copyOf(metaData));
}
}

2
vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java
View File

@@ -58,7 +58,7 @@ class Neo4jVectorStoreIT {
// creation
// function.
@Container
static Neo4jContainer<?> neo4jContainer = new Neo4jContainer<>(DockerImageName.parse("neo4j:5.15"))
static Neo4jContainer<?> neo4jContainer = new Neo4jContainer<>(DockerImageName.parse("neo4j:5.18"))
.withRandomPassword();
List<Document> documents = List.of(