diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java index 689fba372..28be1a7d1 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java @@ -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); diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java index 0347e6310..5b13b56c4 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java @@ -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; + } + } 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 03a224d08..7837ecda3 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 @@ -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(); @@ -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)); } } diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java index ef0b433d5..a258d8153 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java @@ -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 documents = List.of(