Skip to content

Commit 264b132

Browse files
Make the embedding field name configurable for the ElasticSearchVectorStore
Signed-off-by: jonghoon park <dev@jonghoonpark.com>
1 parent b7dcfc1 commit 264b132

File tree

5 files changed

+80
-41
lines changed

5 files changed

+80
-41
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 & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@
142142
* @author Christian Tzolov
143143
* @author Thomas Vitale
144144
* @author Ilayaperumal Gopinathan
145+
* @author Jonghoon Park
145146
* @since 1.0.0
146147
*/
147148
public class ElasticsearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -188,11 +189,12 @@ public void doAdd(List<Document> documents) {
188189
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
189190
this.batchingStrategy);
190191

191-
for (Document document : documents) {
192-
ElasticSearchDocument doc = new ElasticSearchDocument(document.getId(), document.getText(),
193-
document.getMetadata(), embeddings.get(documents.indexOf(document)));
194-
bulkRequestBuilder.operations(
195-
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()))));
196198
}
197199
BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build());
198200
if (bulkRequest.errors()) {
@@ -205,6 +207,13 @@ public void doAdd(List<Document> documents) {
205207
}
206208
}
207209

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+
208217
@Override
209218
public void doDelete(List<String> idList) {
210219
BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
@@ -263,7 +272,7 @@ public List<Document> doSimilaritySearch(SearchRequest searchRequest) {
263272
.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
264273
.similarity(finalThreshold)
265274
.k(searchRequest.getTopK())
266-
.field("embedding")
275+
.field(this.options.getEmbeddingFieldName())
267276
.numCandidates((int) (1.5 * searchRequest.getTopK()))
268277
.filter(fl -> fl
269278
.queryString(qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression())))))
@@ -321,7 +330,7 @@ private void createIndexMapping() {
321330
try {
322331
this.elasticsearchClient.indices()
323332
.create(cr -> cr.index(this.options.getIndexName())
324-
.mappings(map -> map.properties("embedding",
333+
.mappings(map -> map.properties(this.options.getEmbeddingFieldName(),
325334
p -> p.denseVector(dv -> dv.similarity(this.options.getSimilarity().toString())
326335
.dims(this.options.getDimensions())))));
327336
}
@@ -370,17 +379,6 @@ public static Builder builder(RestClient restClient, EmbeddingModel embeddingMod
370379
return new Builder(restClient, embeddingModel);
371380
}
372381

373-
/**
374-
* The representation of {@link Document} along with its embedding.
375-
*
376-
* @param id The id of the document
377-
* @param content The content of the document
378-
* @param metadata The metadata of the document
379-
* @param embedding The vectors representing the content of the document
380-
*/
381-
public record ElasticSearchDocument(String id, String content, Map<String, Object> metadata, float[] embedding) {
382-
}
383-
384382
public static class Builder extends AbstractVectorStoreBuilder<Builder> {
385383

386384
private final RestClient restClient;

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: 29 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
4646
import org.junit.jupiter.params.ParameterizedTest;
4747
import org.junit.jupiter.params.provider.ValueSource;
48+
import org.springframework.ai.model.SimpleApiKey;
4849
import org.testcontainers.elasticsearch.ElasticsearchContainer;
4950
import org.testcontainers.junit.jupiter.Container;
5051
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -57,10 +58,6 @@
5758
import org.springframework.ai.test.vectorstore.BaseVectorStoreTests;
5859
import org.springframework.ai.vectorstore.SearchRequest;
5960
import org.springframework.ai.vectorstore.VectorStore;
60-
import org.springframework.ai.vectorstore.filter.Filter.Expression;
61-
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
62-
import org.springframework.ai.vectorstore.filter.Filter.Key;
63-
import org.springframework.ai.vectorstore.filter.Filter.Value;
6461
import org.springframework.boot.SpringBootConfiguration;
6562
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
6663
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -127,10 +124,11 @@ protected void executeTest(Consumer<VectorStore> testFunction) {
127124
});
128125
}
129126

130-
@Test
131-
public void addAndDeleteDocumentsTest() {
127+
@ParameterizedTest(name = "{0} : {displayName} ")
128+
@ValueSource(strings = { "cosine", "custom_embedding_field" })
129+
public void addAndDeleteDocumentsTest(String vectorStoreBeanName) {
132130
getContextRunner().run(context -> {
133-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
131+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
134132
ElasticsearchVectorStore.class);
135133
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
136134

@@ -160,12 +158,12 @@ public void addAndDeleteDocumentsTest() {
160158
}
161159

162160
@ParameterizedTest(name = "{0} : {displayName} ")
163-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
164-
public void addAndSearchTest(String similarityFunction) {
161+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
162+
public void addAndSearchTest(String vectorStoreBeanName) {
165163

166164
getContextRunner().run(context -> {
167165

168-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
166+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
169167
ElasticsearchVectorStore.class);
170168

171169
vectorStore.add(this.documents);
@@ -197,11 +195,11 @@ public void addAndSearchTest(String similarityFunction) {
197195
}
198196

199197
@ParameterizedTest(name = "{0} : {displayName} ")
200-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
201-
public void searchWithFilters(String similarityFunction) {
198+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
199+
public void searchWithFilters(String vectorStoreBeanName) {
202200

203201
getContextRunner().run(context -> {
204-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
202+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
205203
ElasticsearchVectorStore.class);
206204

207205
var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
@@ -311,11 +309,11 @@ public void searchWithFilters(String similarityFunction) {
311309
}
312310

313311
@ParameterizedTest(name = "{0} : {displayName} ")
314-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
315-
public void documentUpdateTest(String similarityFunction) {
312+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
313+
public void documentUpdateTest(String vectorStoreBeanName) {
316314

317315
getContextRunner().run(context -> {
318-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
316+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
319317
ElasticsearchVectorStore.class);
320318

321319
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -369,10 +367,10 @@ public void documentUpdateTest(String similarityFunction) {
369367
}
370368

371369
@ParameterizedTest(name = "{0} : {displayName} ")
372-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
373-
public void searchThresholdTest(String similarityFunction) {
370+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
371+
public void searchThresholdTest(String vectorStoreBeanName) {
374372
getContextRunner().run(context -> {
375-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
373+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
376374
ElasticsearchVectorStore.class);
377375

378376
vectorStore.add(this.documents);
@@ -507,9 +505,20 @@ public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingModel embeddingMo
507505
.build();
508506
}
509507

508+
@Bean("vectorStore_custom_embedding_field")
509+
public ElasticsearchVectorStore vectorStoreCustomField(EmbeddingModel embeddingModel, RestClient restClient) {
510+
ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
511+
options.setEmbeddingFieldName("custom_embedding_field");
512+
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
513+
.initializeSchema(true)
514+
.options(options)
515+
.build();
516+
}
517+
510518
@Bean
511519
public EmbeddingModel embeddingModel() {
512-
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
520+
return new OpenAiEmbeddingModel(
521+
OpenAiApi.builder().apiKey(new SimpleApiKey(System.getenv("OPENAI_API_KEY"))).build());
513522
}
514523

515524
@Bean

0 commit comments

Comments
 (0)