Skip to content

Commit bf91743

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 b525309 commit bf91743

File tree

5 files changed

+88
-31
lines changed

5 files changed

+88
-31
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.getFieldName() != null) {
77+
elasticsearchVectorStoreOptions.setFieldName(properties.getFieldName());
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 fieldName = "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 getFieldName() {
80+
return fieldName;
81+
}
82+
83+
public void setFieldName(String fieldName) {
84+
this.fieldName = fieldName;
85+
}
86+
7387
}

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

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,7 @@
145145
* @author Christian Tzolov
146146
* @author Thomas Vitale
147147
* @author Ilayaperumal Gopinathan
148+
* @author Jonghoon Park
148149
* @since 1.0.0
149150
*/
150151
public class ElasticsearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -191,11 +192,12 @@ public void doAdd(List<Document> documents) {
191192
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
192193
this.batchingStrategy);
193194

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)));
195+
for (int i = 0; i < embeddings.size(); i++) {
196+
Document document = documents.get(i);
197+
float[] embedding = embeddings.get(i);
198+
bulkRequestBuilder.operations(op -> op.index(idx -> idx.index(this.options.getIndexName())
199+
.id(document.getId())
200+
.document(getDocument(document, embedding, this.options.getFieldName()))));
199201
}
200202
BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build());
201203
if (bulkRequest.errors()) {
@@ -208,6 +210,16 @@ public void doAdd(List<Document> documents) {
208210
}
209211
}
210212

213+
private Object getDocument(Document document, float[] embedding, String fieldName) {
214+
Assert.notNull(document.getText(), "document's text must not be null");
215+
if (fieldName.equals("embedding")) {
216+
return new ElasticSearchDocument(document.getId(), document.getText(), document.getMetadata(), embedding);
217+
}
218+
219+
return Map.of("id", document.getId(), "content", document.getText(), "metadata", document.getMetadata(),
220+
fieldName, embedding);
221+
}
222+
211223
@Override
212224
public Optional<Boolean> doDelete(List<String> idList) {
213225
BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
@@ -264,7 +276,7 @@ public List<Document> doSimilaritySearch(SearchRequest searchRequest) {
264276
.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
265277
.similarity(finalThreshold)
266278
.k((long) searchRequest.getTopK())
267-
.field("embedding")
279+
.field(this.options.getFieldName())
268280
.numCandidates((long) (1.5 * searchRequest.getTopK()))
269281
.filter(fl -> fl
270282
.queryString(qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression())))))
@@ -322,7 +334,7 @@ private void createIndexMapping() {
322334
try {
323335
this.elasticsearchClient.indices()
324336
.create(cr -> cr.index(this.options.getIndexName())
325-
.mappings(map -> map.properties("embedding",
337+
.mappings(map -> map.properties(this.options.getFieldName(),
326338
p -> p.denseVector(dv -> dv.similarity(this.options.getSimilarity().toString())
327339
.dims(this.options.getDimensions())))));
328340
}

‎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 fieldName = "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 getFieldName() {
74+
return fieldName;
75+
}
76+
77+
public void setFieldName(String fieldName) {
78+
this.fieldName = fieldName;
79+
}
80+
6781
}

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

Lines changed: 34 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -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_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_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_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_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_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_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_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,6 +584,16 @@ public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingModel embeddingMo
581584
.build();
582585
}
583586

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

0 commit comments

Comments
 (0)