Skip to content

Commit 729ab32

Browse files
dev-jonghoonparkdev-jonghoonpark
committed
Make the embedding field name configurable for the ElasticSearchVectorStore
Signed-off-by: jonghoon park <[email protected]>
1 parent b466159 commit 729ab32

File tree

5 files changed

+88
-36
lines changed

5 files changed

+88
-36
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: 37 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
4545
import org.junit.jupiter.params.ParameterizedTest;
4646
import org.junit.jupiter.params.provider.ValueSource;
47+
import org.springframework.ai.model.SimpleApiKey;
4748
import org.testcontainers.elasticsearch.ElasticsearchContainer;
4849
import org.testcontainers.junit.jupiter.Container;
4950
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -54,7 +55,6 @@
5455
import org.springframework.ai.openai.OpenAiEmbeddingModel;
5556
import org.springframework.ai.openai.api.OpenAiApi;
5657
import org.springframework.ai.vectorstore.SearchRequest;
57-
import org.springframework.ai.vectorstore.filter.Filter;
5858
import org.springframework.ai.vectorstore.filter.Filter.Expression;
5959
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
6060
import org.springframework.ai.vectorstore.filter.Filter.Key;
@@ -117,10 +117,11 @@ void cleanDatabase() {
117117
});
118118
}
119119

120-
@Test
121-
public void addAndDeleteDocumentsTest() {
120+
@ParameterizedTest(name = "{0} : {displayName} ")
121+
@ValueSource(strings = { "cosine", "custom_embedding_field" })
122+
public void addAndDeleteDocumentsTest(String vectorStoreBeanName) {
122123
getContextRunner().run(context -> {
123-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
124+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
124125
ElasticsearchVectorStore.class);
125126
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
126127

@@ -149,10 +150,11 @@ public void addAndDeleteDocumentsTest() {
149150
});
150151
}
151152

152-
@Test
153-
public void deleteDocumentsByFilterExpressionTest() {
153+
@ParameterizedTest(name = "{0} : {displayName} ")
154+
@ValueSource(strings = { "cosine", "custom_embedding_field" })
155+
public void deleteDocumentsByFilterExpressionTest(String vectorStoreBeanName) {
154156
getContextRunner().run(context -> {
155-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
157+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
156158
ElasticsearchVectorStore.class);
157159
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
158160

@@ -202,10 +204,11 @@ public void deleteDocumentsByFilterExpressionTest() {
202204
});
203205
}
204206

205-
@Test
206-
public void deleteWithStringFilterExpressionTest() {
207+
@ParameterizedTest(name = "{0} : {displayName} ")
208+
@ValueSource(strings = { "cosine", "custom_embedding_field" })
209+
public void deleteWithStringFilterExpressionTest(String vectorStoreBeanName) {
207210
getContextRunner().run(context -> {
208-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
211+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
209212
ElasticsearchVectorStore.class);
210213
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
211214

@@ -234,12 +237,12 @@ public void deleteWithStringFilterExpressionTest() {
234237
}
235238

236239
@ParameterizedTest(name = "{0} : {displayName} ")
237-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
238-
public void addAndSearchTest(String similarityFunction) {
240+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
241+
public void addAndSearchTest(String vectorStoreBeanName) {
239242

240243
getContextRunner().run(context -> {
241244

242-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
245+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
243246
ElasticsearchVectorStore.class);
244247

245248
vectorStore.add(this.documents);
@@ -271,11 +274,11 @@ public void addAndSearchTest(String similarityFunction) {
271274
}
272275

273276
@ParameterizedTest(name = "{0} : {displayName} ")
274-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
275-
public void searchWithFilters(String similarityFunction) {
277+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
278+
public void searchWithFilters(String vectorStoreBeanName) {
276279

277280
getContextRunner().run(context -> {
278-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
281+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
279282
ElasticsearchVectorStore.class);
280283

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

387390
@ParameterizedTest(name = "{0} : {displayName} ")
388-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
389-
public void documentUpdateTest(String similarityFunction) {
391+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
392+
public void documentUpdateTest(String vectorStoreBeanName) {
390393

391394
getContextRunner().run(context -> {
392-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
395+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
393396
ElasticsearchVectorStore.class);
394397

395398
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -443,10 +446,10 @@ public void documentUpdateTest(String similarityFunction) {
443446
}
444447

445448
@ParameterizedTest(name = "{0} : {displayName} ")
446-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
447-
public void searchThresholdTest(String similarityFunction) {
449+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
450+
public void searchThresholdTest(String vectorStoreBeanName) {
448451
getContextRunner().run(context -> {
449-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
452+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
450453
ElasticsearchVectorStore.class);
451454

452455
vectorStore.add(this.documents);
@@ -581,9 +584,20 @@ public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingModel embeddingMo
581584
.build();
582585
}
583586

587+
@Bean("vectorStore_custom_embedding_field")
588+
public ElasticsearchVectorStore vectorStoreCustomField(EmbeddingModel embeddingModel, RestClient restClient) {
589+
ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
590+
options.setEmbeddingFieldName("custom_embedding_field");
591+
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
592+
.initializeSchema(true)
593+
.options(options)
594+
.build();
595+
}
596+
584597
@Bean
585598
public EmbeddingModel embeddingModel() {
586-
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
599+
return new OpenAiEmbeddingModel(
600+
OpenAiApi.builder().apiKey(new SimpleApiKey(System.getenv("OPENAI_API_KEY"))).build());
587601
}
588602

589603
@Bean

0 commit comments

Comments
 (0)