Skip to content

Commit a567120

Browse files
dev-jonghoonparkdev-jonghoonpark
committed
Make the embedding field name configurable for the ElasticSearchVectorStore
Signed-off-by: jonghoon park <dev@jonghoonpark.com>
1 parent b466159 commit a567120

File tree

5 files changed

+85
-35
lines changed

5 files changed

+85
-35
lines changed

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -42,6 +42,7 @@
4242
* @author Josh Long
4343
* @author Christian Tzolov
4444
* @author Soby Chacko
45+
* @author Jonghoon Park
4546
* @since 1.0.0
4647
*/
4748
@AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class)
@@ -72,6 +73,9 @@ ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properti
7273
if (properties.getSimilarity() != null) {
7374
elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity());
7475
}
76+
if (properties.getEmbeddingFieldName() != null) {
77+
elasticsearchVectorStoreOptions.setEmbeddingFieldName(properties.getEmbeddingFieldName());
78+
}
7579

7680
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
7781
.options(elasticsearchVectorStoreOptions)

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@
2626
* @author Eddú Meléndez
2727
* @author Wei Jiang
2828
* @author Josh Long
29+
* @author Jonghoon Park
2930
* @since 1.0.0
3031
*/
3132
@ConfigurationProperties(prefix = "spring.ai.vectorstore.elasticsearch")
@@ -46,6 +47,11 @@ public class ElasticsearchVectorStoreProperties extends CommonVectorStorePropert
4647
*/
4748
private SimilarityFunction similarity;
4849

50+
/**
51+
* The name of the vector field to search against
52+
*/
53+
private String embeddingFieldName = "embedding";
54+
4955
public String getIndexName() {
5056
return this.indexName;
5157
}
@@ -70,4 +76,12 @@ public void setSimilarity(SimilarityFunction similarity) {
7076
this.similarity = similarity;
7177
}
7278

79+
public String getEmbeddingFieldName() {
80+
return embeddingFieldName;
81+
}
82+
83+
public void setEmbeddingFieldName(String embeddingFieldName) {
84+
this.embeddingFieldName = embeddingFieldName;
85+
}
86+
7387
}

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

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import co.elastic.clients.elasticsearch.ElasticsearchClient;
2727
import co.elastic.clients.elasticsearch.core.BulkRequest;
2828
import co.elastic.clients.elasticsearch.core.BulkResponse;
29-
import co.elastic.clients.elasticsearch.core.DeleteByQueryResponse;
3029
import co.elastic.clients.elasticsearch.core.SearchResponse;
3130
import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
3231
import co.elastic.clients.elasticsearch.core.search.Hit;
@@ -41,10 +40,8 @@
4140

4241
import org.springframework.ai.document.Document;
4342
import org.springframework.ai.document.DocumentMetadata;
44-
import org.springframework.ai.embedding.BatchingStrategy;
4543
import org.springframework.ai.embedding.EmbeddingModel;
4644
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
47-
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
4845
import org.springframework.ai.model.EmbeddingUtils;
4946
import org.springframework.ai.observation.conventions.VectorStoreProvider;
5047
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
@@ -145,6 +142,7 @@
145142
* @author Christian Tzolov
146143
* @author Thomas Vitale
147144
* @author Ilayaperumal Gopinathan
145+
* @author Jonghoon Park
148146
* @since 1.0.0
149147
*/
150148
public class ElasticsearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -191,11 +189,12 @@ public void doAdd(List<Document> documents) {
191189
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
192190
this.batchingStrategy);
193191

194-
for (Document document : documents) {
195-
ElasticSearchDocument doc = new ElasticSearchDocument(document.getId(), document.getText(),
196-
document.getMetadata(), embeddings.get(documents.indexOf(document)));
197-
bulkRequestBuilder.operations(
198-
op -> op.index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(doc)));
192+
for (int i = 0; i < embeddings.size(); i++) {
193+
Document document = documents.get(i);
194+
float[] embedding = embeddings.get(i);
195+
bulkRequestBuilder.operations(op -> op.index(idx -> idx.index(this.options.getIndexName())
196+
.id(document.getId())
197+
.document(getDocument(document, embedding, this.options.getEmbeddingFieldName()))));
199198
}
200199
BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build());
201200
if (bulkRequest.errors()) {
@@ -208,6 +207,13 @@ public void doAdd(List<Document> documents) {
208207
}
209208
}
210209

210+
private Object getDocument(Document document, float[] embedding, String embeddingFieldName) {
211+
Assert.notNull(document.getText(), "document's text must not be null");
212+
213+
return Map.of("id", document.getId(), "content", document.getText(), "metadata", document.getMetadata(),
214+
embeddingFieldName, embedding);
215+
}
216+
211217
@Override
212218
public Optional<Boolean> doDelete(List<String> idList) {
213219
BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
@@ -264,7 +270,7 @@ public List<Document> doSimilaritySearch(SearchRequest searchRequest) {
264270
.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
265271
.similarity(finalThreshold)
266272
.k(searchRequest.getTopK())
267-
.field("embedding")
273+
.field(this.options.getEmbeddingFieldName())
268274
.numCandidates((int) (1.5 * searchRequest.getTopK()))
269275
.filter(fl -> fl
270276
.queryString(qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression())))))
@@ -322,7 +328,7 @@ private void createIndexMapping() {
322328
try {
323329
this.elasticsearchClient.indices()
324330
.create(cr -> cr.index(this.options.getIndexName())
325-
.mappings(map -> map.properties("embedding",
331+
.mappings(map -> map.properties(this.options.getEmbeddingFieldName(),
326332
p -> p.denseVector(dv -> dv.similarity(this.options.getSimilarity().toString())
327333
.dims(this.options.getDimensions())))));
328334
}

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
2121
* https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html
2222
*
2323
* @author Wei Jiang
24+
* @author Jonghoon Park
2425
* @since 1.0.0
2526
*/
2627
public class ElasticsearchVectorStoreOptions {
@@ -40,6 +41,11 @@ public class ElasticsearchVectorStoreOptions {
4041
*/
4142
private SimilarityFunction similarity = SimilarityFunction.cosine;
4243

44+
/**
45+
* The name of the vector field to search against
46+
*/
47+
private String embeddingFieldName = "embedding";
48+
4349
public String getIndexName() {
4450
return this.indexName;
4551
}
@@ -64,4 +70,12 @@ public void setSimilarity(SimilarityFunction similarity) {
6470
this.similarity = similarity;
6571
}
6672

73+
public String getEmbeddingFieldName() {
74+
return embeddingFieldName;
75+
}
76+
77+
public void setEmbeddingFieldName(String embeddingFieldName) {
78+
this.embeddingFieldName = embeddingFieldName;
79+
}
80+
6781
}

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

Lines changed: 34 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@
5454
import org.springframework.ai.openai.OpenAiEmbeddingModel;
5555
import org.springframework.ai.openai.api.OpenAiApi;
5656
import org.springframework.ai.vectorstore.SearchRequest;
57-
import org.springframework.ai.vectorstore.filter.Filter;
5857
import org.springframework.ai.vectorstore.filter.Filter.Expression;
5958
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
6059
import org.springframework.ai.vectorstore.filter.Filter.Key;
@@ -117,10 +116,11 @@ void cleanDatabase() {
117116
});
118117
}
119118

120-
@Test
121-
public void addAndDeleteDocumentsTest() {
119+
@ParameterizedTest(name = "{0} : {displayName} ")
120+
@ValueSource(strings = { "cosine", "custom_field" })
121+
public void addAndDeleteDocumentsTest(String vectorStoreBeanName) {
122122
getContextRunner().run(context -> {
123-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
123+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
124124
ElasticsearchVectorStore.class);
125125
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
126126

@@ -149,10 +149,11 @@ public void addAndDeleteDocumentsTest() {
149149
});
150150
}
151151

152-
@Test
153-
public void deleteDocumentsByFilterExpressionTest() {
152+
@ParameterizedTest(name = "{0} : {displayName} ")
153+
@ValueSource(strings = { "cosine", "custom_field" })
154+
public void deleteDocumentsByFilterExpressionTest(String vectorStoreBeanName) {
154155
getContextRunner().run(context -> {
155-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
156+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
156157
ElasticsearchVectorStore.class);
157158
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
158159

@@ -202,10 +203,11 @@ public void deleteDocumentsByFilterExpressionTest() {
202203
});
203204
}
204205

205-
@Test
206-
public void deleteWithStringFilterExpressionTest() {
206+
@ParameterizedTest(name = "{0} : {displayName} ")
207+
@ValueSource(strings = { "cosine", "custom_field" })
208+
public void deleteWithStringFilterExpressionTest(String vectorStoreBeanName) {
207209
getContextRunner().run(context -> {
208-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
210+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
209211
ElasticsearchVectorStore.class);
210212
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
211213

@@ -234,12 +236,12 @@ public void deleteWithStringFilterExpressionTest() {
234236
}
235237

236238
@ParameterizedTest(name = "{0} : {displayName} ")
237-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
238-
public void addAndSearchTest(String similarityFunction) {
239+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_field" })
240+
public void addAndSearchTest(String vectorStoreBeanName) {
239241

240242
getContextRunner().run(context -> {
241243

242-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
244+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
243245
ElasticsearchVectorStore.class);
244246

245247
vectorStore.add(this.documents);
@@ -271,11 +273,11 @@ public void addAndSearchTest(String similarityFunction) {
271273
}
272274

273275
@ParameterizedTest(name = "{0} : {displayName} ")
274-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
275-
public void searchWithFilters(String similarityFunction) {
276+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_field" })
277+
public void searchWithFilters(String vectorStoreBeanName) {
276278

277279
getContextRunner().run(context -> {
278-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
280+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
279281
ElasticsearchVectorStore.class);
280282

281283
var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
@@ -385,11 +387,11 @@ public void searchWithFilters(String similarityFunction) {
385387
}
386388

387389
@ParameterizedTest(name = "{0} : {displayName} ")
388-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
389-
public void documentUpdateTest(String similarityFunction) {
390+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_field" })
391+
public void documentUpdateTest(String vectorStoreBeanName) {
390392

391393
getContextRunner().run(context -> {
392-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
394+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
393395
ElasticsearchVectorStore.class);
394396

395397
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -443,10 +445,10 @@ public void documentUpdateTest(String similarityFunction) {
443445
}
444446

445447
@ParameterizedTest(name = "{0} : {displayName} ")
446-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
447-
public void searchThresholdTest(String similarityFunction) {
448+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_field" })
449+
public void searchThresholdTest(String vectorStoreBeanName) {
448450
getContextRunner().run(context -> {
449-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
451+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
450452
ElasticsearchVectorStore.class);
451453

452454
vectorStore.add(this.documents);
@@ -581,6 +583,16 @@ public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingModel embeddingMo
581583
.build();
582584
}
583585

586+
@Bean("vectorStore_custom_field")
587+
public ElasticsearchVectorStore vectorStoreCustomField(EmbeddingModel embeddingModel, RestClient restClient) {
588+
ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
589+
options.setEmbeddingFieldName("custom_field");
590+
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
591+
.initializeSchema(true)
592+
.options(options)
593+
.build();
594+
}
595+
584596
@Bean
585597
public EmbeddingModel embeddingModel() {
586598
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));

0 commit comments

Comments
 (0)