From 8e04a240ba133dd0d5555135b6e8074fae61b64f Mon Sep 17 00:00:00 2001 From: Gerrit Meier Date: Wed, 15 Jan 2025 16:09:37 +0100 Subject: [PATCH] Neo4j module: Determine default embedding dimension from model. In cases where no custom size is set, derive the size by the given embedding model. Had to migrate the embeddingDimension Spring Boot property from int to Integer to introduce the null check for the fluent config. Everything else would have been just noisy. Closes #977 Signed-off-by: Gerrit Meier --- .../Neo4jVectorStoreAutoConfiguration.java | 3 +- .../neo4j/Neo4jVectorStoreProperties.java | 4 +-- .../vectorstore/neo4j/Neo4jVectorStore.java | 7 +++-- .../vectorstore/neo4j/Neo4jVectorStoreIT.java | 28 +++++++++++++++++++ 4 files changed, 36 insertions(+), 6 deletions(-) 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 5d3d114c145..47be9532ac8 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 @@ -64,7 +64,8 @@ public Neo4jVectorStore vectorStore(Driver driver, EmbeddingModel embeddingModel .customObservationConvention(customObservationConvention.getIfAvailable(() -> null)) .batchingStrategy(batchingStrategy) .databaseName(properties.getDatabaseName()) - .embeddingDimension(properties.getEmbeddingDimension()) + .embeddingDimension(properties.getEmbeddingDimension() != null ? properties.getEmbeddingDimension() + : embeddingModel.dimensions()) .distanceType(properties.getDistanceType()) .label(properties.getLabel()) .embeddingProperty(properties.getEmbeddingProperty()) 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 52606c754b1..72dcd1d176c 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 @@ -33,7 +33,7 @@ public class Neo4jVectorStoreProperties extends CommonVectorStoreProperties { private String databaseName; - private int embeddingDimension = Neo4jVectorStore.DEFAULT_EMBEDDING_DIMENSION; + private Integer embeddingDimension; private Neo4jVectorStore.Neo4jDistanceType distanceType = Neo4jVectorStore.Neo4jDistanceType.COSINE; @@ -55,7 +55,7 @@ public void setDatabaseName(String databaseName) { this.databaseName = databaseName; } - public int getEmbeddingDimension() { + public Integer getEmbeddingDimension() { return this.embeddingDimension; } diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java index c9b915f686d..82fb37363e4 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java @@ -133,6 +133,7 @@ */ public class Neo4jVectorStore extends AbstractObservationVectorStore implements InitializingBean { + @Deprecated(forRemoval = true) public static final int DEFAULT_EMBEDDING_DIMENSION = 1536; public static final int DEFAULT_TRANSACTION_SIZE = 10_000; @@ -182,7 +183,7 @@ protected Neo4jVectorStore(Builder builder) { this.driver = builder.driver; this.sessionConfig = builder.sessionConfig; - this.embeddingDimension = builder.embeddingDimension; + this.embeddingDimension = builder.embeddingDimension.orElseGet(() -> builder.getEmbeddingModel().dimensions()); this.distanceType = builder.distanceType; this.embeddingProperty = SchemaNames.sanitize(builder.embeddingProperty).orElseThrow(); this.label = SchemaNames.sanitize(builder.label).orElseThrow(); @@ -372,7 +373,7 @@ public static class Builder extends AbstractVectorStoreBuilder { private SessionConfig sessionConfig = SessionConfig.defaultConfig(); - private int embeddingDimension = DEFAULT_EMBEDDING_DIMENSION; + private Optional embeddingDimension = Optional.empty(); private Neo4jDistanceType distanceType = Neo4jDistanceType.COSINE; @@ -425,7 +426,7 @@ public Builder sessionConfig(SessionConfig sessionConfig) { */ public Builder embeddingDimension(int dimension) { Assert.isTrue(dimension >= 1, "Dimension has to be positive"); - this.embeddingDimension = dimension; + this.embeddingDimension = Optional.of(dimension); return this; } diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java index 5260b34f930..a10d0f19dee 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java @@ -28,6 +28,7 @@ import org.neo4j.driver.AuthTokens; import org.neo4j.driver.Driver; import org.neo4j.driver.GraphDatabase; +import org.springframework.context.annotation.Primary; import org.testcontainers.containers.Neo4jContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; @@ -301,16 +302,43 @@ void ensureIdIndexGetsCreated() { .isTrue()); } + @Test + void vectorIndexDimensionsDefaultAndOverwriteWorks() { + this.contextRunner.run(context -> { + var result = context.getBean(Driver.class) + .executableQuery( + "SHOW VECTOR INDEXES yield name, options return name, options['indexConfig']['vector.dimensions'] as dimensions") + .execute() + .records() + .stream() + .map(r -> r.get("name").asString() + r.get("dimensions").asInt()) + .toList(); + assertThat(result).containsExactlyInAnyOrder("secondIndex123", "spring-ai-document-index1536"); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { @Bean + @Primary public VectorStore vectorStore(Driver driver, EmbeddingModel embeddingModel) { return Neo4jVectorStore.builder(driver, embeddingModel).initializeSchema(true).build(); } + @Bean + public VectorStore vectorStoreWithCustomDimension(Driver driver, EmbeddingModel embeddingModel) { + + return Neo4jVectorStore.builder(driver, embeddingModel) + .initializeSchema(true) + .indexName("secondIndex") + .embeddingProperty("somethingElse") + .embeddingDimension(123) + .build(); + } + @Bean public Driver driver() { return GraphDatabase.driver(neo4jContainer.getBoltUrl(),