Make the embedding field name configurable for the ElasticSearchVectorStore

Signed-off-by: jonghoon park <[email protected]>

Author: dev-jonghoonpark

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 the original author or authors.
+ * Copyright 2023-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -44,6 +44,7 @@
  * @author Josh Long
  * @author Christian Tzolov
  * @author Soby Chacko
+ * @author Jonghoon Park
  * @since 1.0.0
  */
 @AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class)
@@ -76,6 +77,9 @@ ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properti
 		if (properties.getSimilarity() != null) {
 			elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity());
 		}
+		if (properties.getEmbeddingFieldName() != null) {
+			elasticsearchVectorStoreOptions.setEmbeddingFieldName(properties.getEmbeddingFieldName());
+		}
 
 		return ElasticsearchVectorStore.builder(restClient, embeddingModel)
 			.options(elasticsearchVectorStoreOptions)

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreProperties.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 the original author or authors.
+ * Copyright 2023-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@
  * @author Eddú Meléndez
  * @author Wei Jiang
  * @author Josh Long
+ * @author Jonghoon Park
  * @since 1.0.0
  */
 @ConfigurationProperties(prefix = "spring.ai.vectorstore.elasticsearch")
@@ -46,6 +47,11 @@ public class ElasticsearchVectorStoreProperties extends CommonVectorStorePropert
 	 */
 	private SimilarityFunction similarity;
 
+	/**
+	 * The name of the vector field to search against
+	 */
+	private String embeddingFieldName = "embedding";
+
 	public String getIndexName() {
 		return this.indexName;
 	}
@@ -70,4 +76,12 @@ public void setSimilarity(SimilarityFunction similarity) {
 		this.similarity = similarity;
 	}
 
+	public String getEmbeddingFieldName() {
+		return embeddingFieldName;
+	}
+
+	public void setEmbeddingFieldName(String embeddingFieldName) {
+		this.embeddingFieldName = embeddingFieldName;
+	}
+
 }

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

@@ -142,6 +142,7 @@
  * @author Christian Tzolov
  * @author Thomas Vitale
  * @author Ilayaperumal Gopinathan
+ * @author Jonghoon Park
  * @since 1.0.0
  */
 public class ElasticsearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -188,11 +189,12 @@ public void doAdd(List<Document> documents) {
 		List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
 				this.batchingStrategy);
 
-		for (Document document : documents) {
-			ElasticSearchDocument doc = new ElasticSearchDocument(document.getId(), document.getText(),
-					document.getMetadata(), embeddings.get(documents.indexOf(document)));
-			bulkRequestBuilder.operations(
-					op -> op.index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(doc)));
+		for (int i = 0; i < embeddings.size(); i++) {
+			Document document = documents.get(i);
+			float[] embedding = embeddings.get(i);
+			bulkRequestBuilder.operations(op -> op.index(idx -> idx.index(this.options.getIndexName())
+				.id(document.getId())
+				.document(getDocument(document, embedding, this.options.getEmbeddingFieldName()))));
 		}
 		BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build());
 		if (bulkRequest.errors()) {
@@ -205,6 +207,13 @@ public void doAdd(List<Document> documents) {
 		}
 	}
 
+	private Object getDocument(Document document, float[] embedding, String embeddingFieldName) {
+		Assert.notNull(document.getText(), "document's text must not be null");
+
+		return Map.of("id", document.getId(), "content", document.getText(), "metadata", document.getMetadata(),
+				embeddingFieldName, embedding);
+	}
+
 	@Override
 	public void doDelete(List<String> idList) {
 		BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
@@ -263,7 +272,7 @@ public List<Document> doSimilaritySearch(SearchRequest searchRequest) {
 				.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
 					.similarity(finalThreshold)
 					.k(searchRequest.getTopK())
-					.field("embedding")
+					.field(this.options.getEmbeddingFieldName())
 					.numCandidates((int) (1.5 * searchRequest.getTopK()))
 					.filter(fl -> fl
 						.queryString(qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression())))))
@@ -321,7 +330,7 @@ private void createIndexMapping() {
 		try {
 			this.elasticsearchClient.indices()
 				.create(cr -> cr.index(this.options.getIndexName())
-					.mappings(map -> map.properties("embedding",
+					.mappings(map -> map.properties(this.options.getEmbeddingFieldName(),
 							p -> p.denseVector(dv -> dv.similarity(this.options.getSimilarity().toString())
 								.dims(this.options.getDimensions())))));
 		}
@@ -370,17 +379,6 @@ public static Builder builder(RestClient restClient, EmbeddingModel embeddingMod
 		return new Builder(restClient, embeddingModel);
 	}
 
-	/**
-	 * The representation of {@link Document} along with its embedding.
-	 *
-	 * @param id The id of the document
-	 * @param content The content of the document
-	 * @param metadata The metadata of the document
-	 * @param embedding The vectors representing the content of the document
-	 */
-	public record ElasticSearchDocument(String id, String content, Map<String, Object> metadata, float[] embedding) {
-	}
-
 	public static class Builder extends AbstractVectorStoreBuilder<Builder> {
 
 		private final RestClient restClient;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 the original author or authors.
+ * Copyright 2023-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
  * https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html
  *
  * @author Wei Jiang
+ * @author Jonghoon Park
  * @since 1.0.0
  */
 public class ElasticsearchVectorStoreOptions {
@@ -40,6 +41,11 @@ public class ElasticsearchVectorStoreOptions {
 	 */
 	private SimilarityFunction similarity = SimilarityFunction.cosine;
 
+	/**
+	 * The name of the vector field to search against
+	 */
+	private String embeddingFieldName = "embedding";
+
 	public String getIndexName() {
 		return this.indexName;
 	}
@@ -64,4 +70,12 @@ public void setSimilarity(SimilarityFunction similarity) {
 		this.similarity = similarity;
 	}
 
+	public String getEmbeddingFieldName() {
+		return embeddingFieldName;
+	}
+
+	public void setEmbeddingFieldName(String embeddingFieldName) {
+		this.embeddingFieldName = embeddingFieldName;
+	}
+
 }

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

@@ -123,10 +123,11 @@ protected void executeTest(Consumer<VectorStore> testFunction) {
 		});
 	}
 
-	@Test
-	public void addAndDeleteDocumentsTest() {
+	@ParameterizedTest(name = "{0} : {displayName} ")
+	@ValueSource(strings = { "cosine", "custom_embedding_field" })
+	public void addAndDeleteDocumentsTest(String vectorStoreBeanName) {
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
 					ElasticsearchVectorStore.class);
 			ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
 
@@ -156,12 +157,12 @@ public void addAndDeleteDocumentsTest() {
 	}
 
 	@ParameterizedTest(name = "{0} : {displayName} ")
-	@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
-	public void addAndSearchTest(String similarityFunction) {
+	@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
+	public void addAndSearchTest(String vectorStoreBeanName) {
 
 		getContextRunner().run(context -> {
 
-			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
 					ElasticsearchVectorStore.class);
 
 			vectorStore.add(this.documents);
@@ -193,11 +194,11 @@ public void addAndSearchTest(String similarityFunction) {
 	}
 
 	@ParameterizedTest(name = "{0} : {displayName} ")
-	@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
-	public void searchWithFilters(String similarityFunction) {
+	@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
+	public void searchWithFilters(String vectorStoreBeanName) {
 
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
 					ElasticsearchVectorStore.class);
 
 			var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
@@ -307,11 +308,11 @@ public void searchWithFilters(String similarityFunction) {
 	}
 
 	@ParameterizedTest(name = "{0} : {displayName} ")
-	@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
-	public void documentUpdateTest(String similarityFunction) {
+	@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
+	public void documentUpdateTest(String vectorStoreBeanName) {
 
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
 					ElasticsearchVectorStore.class);
 
 			Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -365,10 +366,10 @@ public void documentUpdateTest(String similarityFunction) {
 	}
 
 	@ParameterizedTest(name = "{0} : {displayName} ")
-	@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
-	public void searchThresholdTest(String similarityFunction) {
+	@ValueSource(strings = { "cosine", "l2_norm", "dot_product", "custom_embedding_field" })
+	public void searchThresholdTest(String vectorStoreBeanName) {
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + vectorStoreBeanName,
 					ElasticsearchVectorStore.class);
 
 			vectorStore.add(this.documents);
@@ -503,6 +504,16 @@ public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingModel embeddingMo
 				.build();
 		}
 
+		@Bean("vectorStore_custom_embedding_field")
+		public ElasticsearchVectorStore vectorStoreCustomField(EmbeddingModel embeddingModel, RestClient restClient) {
+			ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
+			options.setEmbeddingFieldName("custom_embedding_field");
+			return ElasticsearchVectorStore.builder(restClient, embeddingModel)
+				.initializeSchema(true)
+				.options(options)
+				.build();
+		}
+
 		@Bean
 		public EmbeddingModel embeddingModel() {
 			return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build());