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

Allow configuration of vector index name.

Browse Source 
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.
This commit is contained in:
Gerrit Meier
2024-02-06 22:42:35 +01:00
committed by Christian Tzolov
parent dc04327dcf
commit 254b8632cd
4 changed files with 40 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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

4
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfigurationIT.java
View File

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

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

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