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

Make the embedding field name configurable for the ElasticSearchVectorStore

Browse Source 
Signed-off-by: jonghoon park <dev@jonghoonpark.com>
This commit is contained in:
jonghoon park
2025-03-31 20:18:18 +09:00
committed by Ilayaperumal Gopinathan
parent a75b27f99b
commit 4f4da3076a
5 changed files with 77 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreAutoConfiguration.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2024 the original author or authors.
* Copyright 2023-2025 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -44,6 +44,7 @@ import org.springframework.util.StringUtils;
* @author Josh Long
* @author Christian Tzolov
* @author Soby Chacko
* @author Jonghoon Park
* @since 1.0.0
*/
@AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class)
@@ -76,6 +77,9 @@ public class ElasticsearchVectorStoreAutoConfiguration {
if (properties.getSimilarity() != null) {
elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity());
}
if (properties.getEmbeddingFieldName() != null) {
elasticsearchVectorStoreOptions.setEmbeddingFieldName(properties.getEmbeddingFieldName());
}
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
.options(elasticsearchVectorStoreOptions)

16
auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreProperties.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2024 the original author or authors.
* Copyright 2023-2025 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@ import org.springframework.boot.context.properties.ConfigurationProperties;
* @author Eddú Meléndez
* @author Wei Jiang
* @author Josh Long
* @author Jonghoon Park
* @since 1.0.0
*/
@ConfigurationProperties(prefix = "spring.ai.vectorstore.elasticsearch")
@@ -46,6 +47,11 @@ public class ElasticsearchVectorStoreProperties extends CommonVectorStorePropert
*/
private SimilarityFunction similarity;
/**
* The name of the vector field to search against
*/
private String embeddingFieldName = "embedding";
public String getIndexName() {
return this.indexName;
}
@@ -70,4 +76,12 @@ public class ElasticsearchVectorStoreProperties extends CommonVectorStorePropert
this.similarity = similarity;
}
public String getEmbeddingFieldName() {
return embeddingFieldName;
}
public void setEmbeddingFieldName(String embeddingFieldName) {
this.embeddingFieldName = embeddingFieldName;
}
}

34
vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java
View File

@@ -142,6 +142,7 @@ import org.springframework.util.Assert;
* @author Christian Tzolov
* @author Thomas Vitale
* @author Ilayaperumal Gopinathan
* @author Jonghoon Park
* @since 1.0.0
*/
public class ElasticsearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -188,11 +189,12 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
for (Document document : documents) {
ElasticSearchDocument doc = new ElasticSearchDocument(document.getId(), document.getText(),
document.getMetadata(), embeddings.get(documents.indexOf(document)));
bulkRequestBuilder.operations(
op -> op.index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(doc)));
for (int i = 0; i < embeddings.size(); i++) {
Document document = documents.get(i);
float[] embedding = embeddings.get(i);
bulkRequestBuilder.operations(op -> op.index(idx -> idx.index(this.options.getIndexName())
.id(document.getId())
.document(getDocument(document, embedding, this.options.getEmbeddingFieldName()))));
}
BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build());
if (bulkRequest.errors()) {
@@ -205,6 +207,13 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
}
}
private Object getDocument(Document document, float[] embedding, String embeddingFieldName) {
Assert.notNull(document.getText(), "document's text must not be null");
return Map.of("id", document.getId(), "content", document.getText(), "metadata", document.getMetadata(),
embeddingFieldName, embedding);
}
@Override
public void doDelete(List<String> idList) {
BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
@@ -263,7 +272,7 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
.similarity(finalThreshold)
.k(searchRequest.getTopK())
.field("embedding")
.field(this.options.getEmbeddingFieldName())
.numCandidates((int) (1.5 * searchRequest.getTopK()))
.filter(fl -> fl
.queryString(qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression())))))
@@ -321,7 +330,7 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
try {
this.elasticsearchClient.indices()
.create(cr -> cr.index(this.options.getIndexName())
.mappings(map -> map.properties("embedding",
.mappings(map -> map.properties(this.options.getEmbeddingFieldName(),
p -> p.denseVector(dv -> dv.similarity(this.options.getSimilarity().toString())
.dims(this.options.getDimensions())))));
}
@@ -370,17 +379,6 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
return new Builder(restClient, embeddingModel);
}
/**
* The representation of {@link Document} along with its embedding.
*
* @param id The id of the document
* @param content The content of the document
* @param metadata The metadata of the document
* @param embedding The vectors representing the content of the document
*/
public record ElasticSearchDocument(String id, String content, Map<String, Object> metadata, float[] embedding) {
}
public static class Builder extends AbstractVectorStoreBuilder<Builder> {
private final RestClient restClient;

16
vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreOptions.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2024 the original author or authors.
* Copyright 2023-2025 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@ package org.springframework.ai.vectorstore.elasticsearch;
* https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html
*
* @author Wei Jiang
* @author Jonghoon Park
* @since 1.0.0
*/
public class ElasticsearchVectorStoreOptions {
@@ -40,6 +41,11 @@ public class ElasticsearchVectorStoreOptions {
*/
private SimilarityFunction similarity = SimilarityFunction.cosine;
/**
* The name of the vector field to search against
*/
private String embeddingFieldName = "embedding";
public String getIndexName() {
return this.indexName;
}
@@ -64,4 +70,12 @@ public class ElasticsearchVectorStoreOptions {
this.similarity = similarity;
}
public String getEmbeddingFieldName() {
return embeddingFieldName;
}
public void setEmbeddingFieldName(String embeddingFieldName) {
this.embeddingFieldName = embeddingFieldName;
}
}

41
vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java
View File

@@ -123,10 +123,11 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests {
});
}
@Test
public void addAndDeleteDocumentsTest() {
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "cosine", "custom_embedding_field" })
public void addAndDeleteDocumentsTest(String vectorStoreBeanName) {
getContextRunner().run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
ElasticsearchVectorStore.class);
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
@@ -156,12 +157,12 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests {
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
public void addAndSearchTest(String similarityFunction) {
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
public void addAndSearchTest(String vectorStoreBeanName) {
getContextRunner().run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
ElasticsearchVectorStore.class);
vectorStore.add(this.documents);
@@ -193,11 +194,11 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests {
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
public void searchWithFilters(String similarityFunction) {
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
public void searchWithFilters(String vectorStoreBeanName) {
getContextRunner().run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
ElasticsearchVectorStore.class);
var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
@@ -307,11 +308,11 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests {
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
public void documentUpdateTest(String similarityFunction) {
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
public void documentUpdateTest(String vectorStoreBeanName) {
getContextRunner().run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
ElasticsearchVectorStore.class);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -365,10 +366,10 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests {
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
public void searchThresholdTest(String similarityFunction) {
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
public void searchThresholdTest(String vectorStoreBeanName) {
getContextRunner().run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
ElasticsearchVectorStore.class);
vectorStore.add(this.documents);
@@ -503,6 +504,16 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests {
.build();
}
@Bean("vectorStore_custom_embedding_field")
public ElasticsearchVectorStore vectorStoreCustomField(EmbeddingModel embeddingModel, RestClient restClient) {
ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
options.setEmbeddingFieldName("custom_embedding_field");
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
.initializeSchema(true)
.options(options)
.build();
}
@Bean
public EmbeddingModel embeddingModel() {
return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build());