Skip to content

Commit 4f4da30

Browse files
dev-jonghoonparkilayaperumalg
authored andcommitted
Make the embedding field name configurable for the ElasticSearchVectorStore
Signed-off-by: jonghoon park <dev@jonghoonpark.com>
1 parent a75b27f commit 4f4da30

File tree

5 files changed

+77
-36
lines changed

5 files changed

+77
-36
lines changed

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/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.
@@ -44,6 +44,7 @@
4444
* @author Josh Long
4545
* @author Christian Tzolov
4646
* @author Soby Chacko
47+
* @author Jonghoon Park
4748
* @since 1.0.0
4849
*/
4950
@AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class)
@@ -76,6 +77,9 @@ ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properti
7677
if (properties.getSimilarity() != null) {
7778
elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity());
7879
}
80+
if (properties.getEmbeddingFieldName() != null) {
81+
elasticsearchVectorStoreOptions.setEmbeddingFieldName(properties.getEmbeddingFieldName());
82+
}
7983

8084
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
8185
.options(elasticsearchVectorStoreOptions)

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/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: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -123,10 +123,11 @@ protected void executeTest(Consumer<VectorStore> testFunction) {
123123
});
124124
}
125125

126-
@Test
127-
public void addAndDeleteDocumentsTest() {
126+
@ParameterizedTest(name = "{0} : {displayName} ")
127+
@ValueSource(strings = { "cosine", "custom_embedding_field" })
128+
public void addAndDeleteDocumentsTest(String vectorStoreBeanName) {
128129
getContextRunner().run(context -> {
129-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
130+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
130131
ElasticsearchVectorStore.class);
131132
ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
132133

@@ -156,12 +157,12 @@ public void addAndDeleteDocumentsTest() {
156157
}
157158

158159
@ParameterizedTest(name = "{0} : {displayName} ")
159-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
160-
public void addAndSearchTest(String similarityFunction) {
160+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
161+
public void addAndSearchTest(String vectorStoreBeanName) {
161162

162163
getContextRunner().run(context -> {
163164

164-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
165+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
165166
ElasticsearchVectorStore.class);
166167

167168
vectorStore.add(this.documents);
@@ -193,11 +194,11 @@ public void addAndSearchTest(String similarityFunction) {
193194
}
194195

195196
@ParameterizedTest(name = "{0} : {displayName} ")
196-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
197-
public void searchWithFilters(String similarityFunction) {
197+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
198+
public void searchWithFilters(String vectorStoreBeanName) {
198199

199200
getContextRunner().run(context -> {
200-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
201+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
201202
ElasticsearchVectorStore.class);
202203

203204
var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
@@ -307,11 +308,11 @@ public void searchWithFilters(String similarityFunction) {
307308
}
308309

309310
@ParameterizedTest(name = "{0} : {displayName} ")
310-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
311-
public void documentUpdateTest(String similarityFunction) {
311+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
312+
public void documentUpdateTest(String vectorStoreBeanName) {
312313

313314
getContextRunner().run(context -> {
314-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
315+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
315316
ElasticsearchVectorStore.class);
316317

317318
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -365,10 +366,10 @@ public void documentUpdateTest(String similarityFunction) {
365366
}
366367

367368
@ParameterizedTest(name = "{0} : {displayName} ")
368-
@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
369-
public void searchThresholdTest(String similarityFunction) {
369+
@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
370+
public void searchThresholdTest(String vectorStoreBeanName) {
370371
getContextRunner().run(context -> {
371-
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
372+
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
372373
ElasticsearchVectorStore.class);
373374

374375
vectorStore.add(this.documents);
@@ -503,6 +504,16 @@ public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingModel embeddingMo
503504
.build();
504505
}
505506

507+
@Bean("vectorStore_custom_embedding_field")
508+
public ElasticsearchVectorStore vectorStoreCustomField(EmbeddingModel embeddingModel, RestClient restClient) {
509+
ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
510+
options.setEmbeddingFieldName("custom_embedding_field");
511+
return ElasticsearchVectorStore.builder(restClient, embeddingModel)
512+
.initializeSchema(true)
513+
.options(options)
514+
.build();
515+
}
516+
506517
@Bean
507518
public EmbeddingModel embeddingModel() {
508519
return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build());

0 commit comments

Comments
 (0)