diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreAutoConfiguration.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreAutoConfiguration.java index 5327f46ff..56fc50c54 100644 --- a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreAutoConfiguration.java +++ b/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 @@ import org.springframework.util.StringUtils; * @author Josh Long * @author Christian Tzolov * @author Soby Chacko + * @author Jonghoon Park * @since 1.0.0 */ @AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class) @@ -76,6 +77,9 @@ public class ElasticsearchVectorStoreAutoConfiguration { if (properties.getSimilarity() != null) { elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity()); } + if (properties.getEmbeddingFieldName() != null) { + elasticsearchVectorStoreOptions.setEmbeddingFieldName(properties.getEmbeddingFieldName()); + } return ElasticsearchVectorStore.builder(restClient, embeddingModel) .options(elasticsearchVectorStoreOptions) diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreProperties.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreProperties.java index 6b4454834..944250e11 100644 --- a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch/src/main/java/org/springframework/ai/vectorstore/elasticsearch/autoconfigure/ElasticsearchVectorStoreProperties.java +++ b/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 @@ import org.springframework.boot.context.properties.ConfigurationProperties; * @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 class ElasticsearchVectorStoreProperties extends CommonVectorStorePropert this.similarity = similarity; } + public String getEmbeddingFieldName() { + return embeddingFieldName; + } + + public void setEmbeddingFieldName(String embeddingFieldName) { + this.embeddingFieldName = embeddingFieldName; + } + } diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java index 98fa84d68..f09194632 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java @@ -142,6 +142,7 @@ import org.springframework.util.Assert; * @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 class ElasticsearchVectorStore extends AbstractObservationVectorStore imp List 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 class ElasticsearchVectorStore extends AbstractObservationVectorStore imp } } + 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 idList) { BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder(); @@ -263,7 +272,7 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp .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 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp 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 class ElasticsearchVectorStore extends AbstractObservationVectorStore imp 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 metadata, float[] embedding) { - } - public static class Builder extends AbstractVectorStoreBuilder { private final RestClient restClient; diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreOptions.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreOptions.java index 2958b6e66..8b558bc35 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreOptions.java +++ b/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 @@ package org.springframework.ai.vectorstore.elasticsearch; * 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 class ElasticsearchVectorStoreOptions { this.similarity = similarity; } + public String getEmbeddingFieldName() { + return embeddingFieldName; + } + + public void setEmbeddingFieldName(String embeddingFieldName) { + this.embeddingFieldName = embeddingFieldName; + } + } diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java index 14d037f6f..f1f706539 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java @@ -123,10 +123,11 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests { }); } - @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 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests { } @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 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests { } @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 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests { } @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 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests { } @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 @@ class ElasticsearchVectorStoreIT extends BaseVectorStoreTests { .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());