From 254b8632cd0a4235fd2d519ed2de51a70f6da841 Mon Sep 17 00:00:00 2001 From: Gerrit Meier Date: Tue, 6 Feb 2024 22:42:35 +0100 Subject: [PATCH] Allow configuration of vector index name. At the moment, it is not possible to configure SpringAI to use an existing index in the database. This commit enables the user to provide the index name for auto configuration or builder usage. --- .../Neo4jVectorStoreAutoConfiguration.java | 1 + .../neo4j/Neo4jVectorStoreProperties.java | 10 ++++++ .../Neo4jVectorStoreAutoConfigurationIT.java | 4 ++- .../ai/vectorstore/Neo4jVectorStore.java | 31 ++++++++++++++++--- 4 files changed, 40 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 60ababd29..b412530bc 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 @@ -56,6 +56,7 @@ public class Neo4jVectorStoreAutoConfiguration { .withDistanceType(properties.getDistanceType()) .withLabel(properties.getLabel()) .withEmbeddingProperty(properties.getEmbeddingProperty()) + .withIndexName(properties.getIndexName()) .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 29c11ee02..1ff83eeb8 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 @@ -37,6 +37,8 @@ public class Neo4jVectorStoreProperties { private String embeddingProperty = Neo4jVectorStore.DEFAULT_EMBEDDING_PROPERTY; + private String indexName = Neo4jVectorStore.DEFAULT_INDEX_NAME; + public String getDatabaseName() { return databaseName; } @@ -77,4 +79,12 @@ public class Neo4jVectorStoreProperties { this.embeddingProperty = embeddingProperty; } + public String getIndexName() { + return this.indexName; + } + + public void setIndexName(String indexName) { + this.indexName = indexName; + } + } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfigurationIT.java index cfe683f67..23393a859 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfigurationIT.java @@ -66,11 +66,13 @@ public class Neo4jVectorStoreAutoConfigurationIT { void addAndSearch() { contextRunner .withPropertyValues("spring.ai.vectorstore.neo4j.label=my_test_label", - "spring.ai.vectorstore.neo4j.embeddingDimension=384") + "spring.ai.vectorstore.neo4j.embeddingDimension=384", + "spring.ai.vectorstore.neo4j.indexName=customIndexName") .run(context -> { var properties = context.getBean(Neo4jVectorStoreProperties.class); assertThat(properties.getLabel()).isEqualTo("my_test_label"); assertThat(properties.getEmbeddingDimension()).isEqualTo(384); + assertThat(properties.getIndexName()).isEqualTo("customIndexName"); VectorStore vectorStore = context.getBean(VectorStore.class); vectorStore.add(documents); 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 633b2b4c5..d284fb84b 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 @@ -69,6 +69,8 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { private final String quotedLabel; + private final String indexName; + /** * Start building a new configuration. * @return The entry point for creating a new configuration. @@ -97,6 +99,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { this.label = builder.label; this.embeddingProperty = builder.embeddingProperty; this.quotedLabel = SchemaNames.sanitize(this.label).orElseThrow(); + this.indexName = builder.indexName; } public static class Builder { @@ -111,6 +114,8 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { private String embeddingProperty = DEFAULT_EMBEDDING_PROPERTY; + private String indexName = DEFAULT_INDEX_NAME; + private Builder() { } @@ -182,6 +187,21 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { return this; } + /** + * Configures the vector index to be used. Defaults to + * {@literal spring-ai-document-index}. + * @param newIndexName The name of the index to be used for storing and + * searching data. + * @return this builder + */ + public Builder withIndexName(String newIndexName) { + + Assert.hasText(newIndexName, "Index name may not be null or blank"); + + this.indexName = newIndexName; + return this; + } + /** * {@return the immutable configuration} */ @@ -198,7 +218,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { public static final String DEFAULT_LABEL = "Document"; - private static final String INDEX_NAME = "spring-ai-document-index"; + public static final String DEFAULT_INDEX_NAME = "spring-ai-document-index"; public static final String DEFAULT_EMBEDDING_PROPERTY = "embedding"; @@ -275,7 +295,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { WHERE score >= $threshold RETURN node, score """, - Map.of("indexName", INDEX_NAME, "numberOfNearestNeighbours", request.getTopK(), + Map.of("indexName", this.config.indexName, "numberOfNearestNeighbours", request.getTopK(), "embeddingValue", embedding, "threshold", request.getSimilarityThreshold())) .list(Neo4jVectorStore::recordToDocument); } @@ -292,7 +312,8 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { .consume(); var vectorIndexExists = session - .run("SHOW INDEXES YIELD name WHERE name = $name RETURN count(*) > 0", Map.of("name", INDEX_NAME)) + .run("SHOW INDEXES YIELD name WHERE name = $name RETURN count(*) > 0", + Map.of("name", this.config.indexName)) .single() .get(0) .asBoolean(); @@ -300,7 +321,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { if (!vectorIndexExists) { var statement = "CALL db.index.vector.createNodeIndex($indexName, $label, $embeddingProperty, $embeddingDimension, $distanceType)"; session.run(statement, - Map.of("indexName", INDEX_NAME, "label", this.config.label, "embeddingProperty", + Map.of("indexName", this.config.indexName, "label", this.config.label, "embeddingProperty", this.config.embeddingProperty, "embeddingDimension", this.config.embeddingDimension, "distanceType", this.config.distanceType.name)) .consume(); @@ -323,7 +344,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { document.getMetadata().forEach((k, v) -> properties.put("metadata." + k, Values.value(v))); row.put("properties", properties); - row.put(DEFAULT_EMBEDDING_PROPERTY, Values.value(toFloatArray(embedding))); + row.put(this.config.embeddingProperty, Values.value(toFloatArray(embedding))); return row; }