rebase

Author: l-trotta

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java

@@ -28,6 +28,8 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.util.StringUtils;
 
+import java.util.Objects;
+
 /**
  * @author Eddú Meléndez
  * @author Wei Jiang
@@ -52,10 +54,7 @@ ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properti
 		if (properties.getDimensions() != null) {
 			elasticsearchVectorStoreOptions.setDimensions(properties.getDimensions());
 		}
-		if (properties.isDenseVectorIndexing() != null) {
-			elasticsearchVectorStoreOptions.setDenseVectorIndexing(properties.isDenseVectorIndexing());
-		}
-		if (StringUtils.hasText(properties.getSimilarity())) {
+		if (properties.getSimilarity() != null) {
 			elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity());
 		}
 

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java

@@ -16,6 +16,7 @@
 package org.springframework.ai.autoconfigure.vectorstore.elasticsearch;
 
 import org.springframework.ai.autoconfigure.CommonVectorStoreProperties;
+import org.springframework.ai.vectorstore.SimilarityFunction;
 import org.springframework.boot.context.properties.ConfigurationProperties;
 
 /**
@@ -37,15 +38,10 @@ public class ElasticsearchVectorStoreProperties extends CommonVectorStorePropert
 	 */
 	private Integer dimensions;
 
-	/**
-	 * Whether to use dense vector indexing.
-	 */
-	private Boolean denseVectorIndexing;
-
 	/**
 	 * The similarity function to use.
 	 */
-	private String similarity;
+	private SimilarityFunction similarity;
 
 	public String getIndexName() {
 		return this.indexName;
@@ -63,19 +59,11 @@ public void setDimensions(Integer dimensions) {
 		this.dimensions = dimensions;
 	}
 
-	public Boolean isDenseVectorIndexing() {
-		return denseVectorIndexing;
-	}
-
-	public void setDenseVectorIndexing(Boolean denseVectorIndexing) {
-		this.denseVectorIndexing = denseVectorIndexing;
-	}
-
-	public String getSimilarity() {
+	public SimilarityFunction getSimilarity() {
 		return similarity;
 	}
 
-	public void setSimilarity(String similarity) {
+	public void setSimilarity(SimilarityFunction similarity) {
 		this.similarity = similarity;
 	}
 

vector-stores/spring-ai-elasticsearch-store/pom.xml

@@ -35,6 +35,7 @@
 		<dependency>
 			<groupId>co.elastic.clients</groupId>
 			<artifactId>elasticsearch-java</artifactId>
+			<version>8.13.2</version>
 		</dependency>
 
 		<!-- TESTING -->

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

@@ -16,10 +16,12 @@
 package org.springframework.ai.vectorstore;
 
 import co.elastic.clients.elasticsearch.ElasticsearchClient;
-import co.elastic.clients.elasticsearch._types.mapping.TypeMapping;
+import co.elastic.clients.elasticsearch._types.mapping.DenseVectorProperty;
+import co.elastic.clients.elasticsearch._types.mapping.Property;
 import co.elastic.clients.elasticsearch.core.BulkRequest;
 import co.elastic.clients.elasticsearch.core.BulkResponse;
 import co.elastic.clients.elasticsearch.core.SearchResponse;
+import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
 import co.elastic.clients.elasticsearch.core.search.Hit;
 import co.elastic.clients.elasticsearch.indices.CreateIndexResponse;
 import co.elastic.clients.json.jackson.JacksonJsonpMapper;
@@ -43,6 +45,7 @@
 import java.util.stream.Collectors;
 
 import static java.lang.Math.sqrt;
+import static org.springframework.ai.vectorstore.SimilarityFunction.l2_norm;
 
 /**
  * @author Jemin Huh
@@ -59,8 +62,6 @@ public class ElasticsearchVectorStore implements VectorStore, InitializingBean {
 
 	private final ElasticsearchVectorStoreOptions options;
 
-	private String similarityFunction = SIMILARITY_DEFAULT;
-
 	private final FilterExpressionConverter filterExpressionConverter;
 
 	private String similarityFunction;
@@ -85,18 +86,18 @@ public ElasticsearchVectorStore(ElasticsearchVectorStoreOptions options, RestCli
 
 	@Override
 	public void add(List<Document> documents) {
-		BulkRequest.Builder builkRequestBuilder = new BulkRequest.Builder();
+		BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
 
 		for (Document document : documents) {
 			if (Objects.isNull(document.getEmbedding()) || document.getEmbedding().isEmpty()) {
 				logger.debug("Calling EmbeddingModel for document id = " + document.getId());
 				document.setEmbedding(this.embeddingModel.embed(document));
 			}
-			builkRequestBuilder.operations(op -> op
+			bulkRequestBuilder.operations(op -> op
 				.index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(document)));
 		}
 
-		BulkResponse bulkRequest = bulkRequest(builkRequestBuilder.build());
+		BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build());
 
 		if (bulkRequest.errors()) {
 			List<BulkResponseItem> bulkResponseItems = bulkRequest.items();
@@ -110,10 +111,10 @@ public void add(List<Document> documents) {
 
 	@Override
 	public Optional<Boolean> delete(List<String> idList) {
-		BulkRequest.Builder builkRequestBuilder = new BulkRequest.Builder();
+		BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
 		for (String id : idList)
-			builkRequestBuilder.operations(op -> op.delete(idx -> idx.index(this.options.getIndexName()).id(id)));
-		return Optional.of(bulkRequest(builkRequestBuilder.build()).errors());
+			bulkRequestBuilder.operations(op -> op.delete(idx -> idx.index(this.options.getIndexName()).id(id)));
+		return Optional.of(bulkRequest(bulkRequestBuilder.build()).errors());
 	}
 
 	private BulkResponse bulkRequest(BulkRequest bulkRequest) {
@@ -131,7 +132,7 @@ public List<Document> similaritySearch(SearchRequest searchRequest) {
 		try {
 			float threshold = (float) searchRequest.getSimilarityThreshold();
 			// reverting l2_norm distance to its original value
-			if (similarityFunction.equals("l2_norm")) {
+			if (options.getSimilarity().equals(l2_norm)) {
 				threshold = 1 - threshold;
 			}
 			final float finalThreshold = threshold;
@@ -141,11 +142,11 @@ public List<Document> similaritySearch(SearchRequest searchRequest) {
 				.toList();
 
 			SearchResponse<Document> res = elasticsearchClient.search(
-					sr -> sr.index(this.index)
+					sr -> sr.index(options.getIndexName())
 						.knn(knn -> knn.queryVector(vectors)
 							.similarity(finalThreshold)
 							.k(searchRequest.getTopK())
-							.field(EMBEDDING)
+							.field("embedding")
 							.numCandidates((long) (1.5 * searchRequest.getTopK()))
 							.filter(fl -> fl.queryString(
 									qs -> qs.query(getElasticsearchQueryString(searchRequest.getFilterExpression()))))),
@@ -174,8 +175,8 @@ private Document toDocument(Hit<Document> hit) {
 	// more info on score/distance calculation
 	// https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html#knn-similarity-search
 	private float calculateDistance(Float score) {
-		switch (similarityFunction) {
-			case "l2_norm":
+		switch (options.getSimilarity()) {
+			case l2_norm:
 				// the returned value of l2_norm is the opposite of the other functions
 				// (closest to zero means more accurate), so to make it consistent
 				// with the other functions the reverse is returned applying a "1-"
@@ -187,9 +188,9 @@ private float calculateDistance(Float score) {
 		}
 	}
 
-	public boolean existsIndex() {
+	public boolean indexExists() {
 		try {
-			return this.elasticsearchClient.indices().exists(ex -> ex.index(this.index)).value();
+			return this.elasticsearchClient.indices().exists(ex -> ex.index(options.getIndexName())).value();
 		}
 		catch (IOException e) {
 			throw new RuntimeException(e);
@@ -199,18 +200,9 @@ public boolean existsIndex() {
 	private CreateIndexResponse createIndexMapping() {
 		try {
 			return this.elasticsearchClient.indices()
-				.create(createIndexBuilder -> createIndexBuilder.index(options.getIndexName())
-					.mappings(typeMappingBuilder -> {
-						typeMappingBuilder.properties("embedding",
-								new Property.Builder()
-									.denseVector(new DenseVectorProperty.Builder().dims(options.getDimensions())
-										.similarity(options.getSimilarity())
-										.index(options.isDenseVectorIndexing())
-										.build())
-									.build());
-
-						return typeMappingBuilder;
-					}));
+				.create(cr -> cr.index(options.getIndexName())
+					.mappings(map -> map.properties("embedding", p -> p.denseVector(
+							dv -> dv.similarity(options.getSimilarity().toString()).dims(options.getDimensions())))));
 		}
 		catch (IOException e) {
 			throw new RuntimeException(e);

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

@@ -34,15 +34,10 @@ public class ElasticsearchVectorStoreOptions {
 	 */
 	private int dimensions = 1536;
 
-	/**
-	 * Whether to use dense vector indexing.
-	 */
-	private boolean denseVectorIndexing = true;
-
 	/**
 	 * The similarity function to use.
 	 */
-	private String similarity = "cosine";
+	private SimilarityFunction similarity = SimilarityFunction.cosine;
 
 	public String getIndexName() {
 		return indexName;
@@ -60,19 +55,11 @@ public void setDimensions(int dims) {
 		this.dimensions = dims;
 	}
 
-	public boolean isDenseVectorIndexing() {
-		return denseVectorIndexing;
-	}
-
-	public void setDenseVectorIndexing(boolean denseVectorIndexing) {
-		this.denseVectorIndexing = denseVectorIndexing;
-	}
-
-	public String getSimilarity() {
+	public SimilarityFunction getSimilarity() {
 		return similarity;
 	}
 
-	public void setSimilarity(String similarity) {
+	public void setSimilarity(SimilarityFunction similarity) {
 		this.similarity = similarity;
 	}
 

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

@@ -0,0 +1,12 @@
+package org.springframework.ai.vectorstore;
+
+/*
+https://www.elastic.co/guide/en/elasticsearch/reference/master/dense-vector.html
+max_inner_product is currently not supported because the distance value is not
+normalized and would not comply with the requirement of being between 0 and 1
+*/
+public enum SimilarityFunction {
+
+	l2_norm, dot_product, cosine
+
+}

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

@@ -26,7 +26,7 @@
 import java.util.concurrent.TimeUnit;
 
 import co.elastic.clients.elasticsearch.ElasticsearchClient;
-import co.elastic.clients.elasticsearch._types.mapping.TypeMapping;
+import co.elastic.clients.elasticsearch.cat.indices.IndicesRecord;
 import co.elastic.clients.json.jackson.JacksonJsonpMapper;
 import co.elastic.clients.transport.rest_client.RestClientTransport;
 import com.fasterxml.jackson.databind.DeserializationFeature;
@@ -93,29 +93,15 @@ private ApplicationContextRunner getContextRunner() {
 		return new ApplicationContextRunner().withUserConfiguration(TestApplication.class);
 	}
 
-	private void prepareMapping(String similarityFunction, ElasticsearchVectorStore vectorStore) {
-		if (!similarityFunction.equals("cosine")) { // cosine is the default similarity
-													// function, no need for custom
-													// mapping
-
-			// vector dimension 1536 is openAI specific
-			TypeMapping mapping = TypeMapping.of(tm -> tm.properties("embedding",
-					p -> p.denseVector(dv -> dv.dims(1536).index(true).similarity(similarityFunction))));
-
-			vectorStore.createIndexMapping(mapping);
-		}
-	}
 
 	@BeforeEach
 	void cleanDatabase() {
 		getContextRunner().run(context -> {
-			VectorStore vectorStore = context.getBean(VectorStore.class);
-			vectorStore.delete(List.of("_all"));
-			// deleting index so that it can be recreated with new mapping
-			// containing a different similarity function
+			// deleting indices and data before following tests
 			ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
-			if (elasticsearchClient.indices().exists(ex -> ex.index("spring-ai-document-index")).value()) {
-				elasticsearchClient.indices().delete(del -> del.index("spring-ai-document-index"));
+			List indices = elasticsearchClient.cat().indices().valueBody().stream().map(IndicesRecord::index).toList();
+			if(!indices.isEmpty()) {
+				elasticsearchClient.indices().delete(del -> del.index(indices));
 			}
 		});
 	}
@@ -125,9 +111,8 @@ void cleanDatabase() {
 	public void addAndSearchTest(String similarityFunction) {
 
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class);
 
-			prepareMapping(similarityFunction, vectorStore);
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_"+similarityFunction, ElasticsearchVectorStore.class);
 
 			vectorStore.add(documents);
 
@@ -162,9 +147,7 @@ public void addAndSearchTest(String similarityFunction) {
 	public void searchWithFilters(String similarityFunction) {
 
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class);
-
-			prepareMapping(similarityFunction, vectorStore);
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_"+similarityFunction, ElasticsearchVectorStore.class);
 
 			var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
 					Map.of("country", "BG", "year", 2020, "activationDate", new Date(1000)));
@@ -260,9 +243,7 @@ public void searchWithFilters(String similarityFunction) {
 	public void documentUpdateTest(String similarityFunction) {
 
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class);
-
-			prepareMapping(similarityFunction, vectorStore);
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_"+similarityFunction, ElasticsearchVectorStore.class);
 
 			Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
 					Map.of("meta1", "meta1"));
@@ -314,9 +295,7 @@ public void documentUpdateTest(String similarityFunction) {
 	@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
 	public void searchThresholdTest(String similarityFunction) {
 		getContextRunner().run(context -> {
-			ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class);
-
-			prepareMapping(similarityFunction, vectorStore);
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_"+similarityFunction, ElasticsearchVectorStore.class);
 
 			vectorStore.add(documents);
 
@@ -355,11 +334,27 @@ public void searchThresholdTest(String similarityFunction) {
 	@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
 	public static class TestApplication {
 
-		@Bean
-		public ElasticsearchVectorStore vectorStore(EmbeddingClient embeddingClient, RestClient restClient) {
+		@Bean("vectorStore_cosine")
+		public ElasticsearchVectorStore vectorStoreDefault(EmbeddingClient embeddingClient, RestClient restClient) {
 			return new ElasticsearchVectorStore(restClient, embeddingClient);
 		}
 
+		@Bean("vectorStore_l2_norm")
+		public ElasticsearchVectorStore vectorStoreL2(EmbeddingClient embeddingClient, RestClient restClient) {
+			ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
+			options.setIndexName("index_l2");
+			options.setSimilarity(SimilarityFunction.l2_norm);
+			return new ElasticsearchVectorStore(options,restClient, embeddingClient);
+		}
+
+		@Bean("vectorStore_dot_product")
+		public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingClient embeddingClient, RestClient restClient) {
+			ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
+			options.setIndexName("index_dot_product");
+			options.setSimilarity(SimilarityFunction.dot_product);
+			return new ElasticsearchVectorStore(options,restClient, embeddingClient);
+		}
+
 		@Bean
 		public EmbeddingModel embeddingModel() {
 			return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
@@ -377,5 +372,4 @@ ElasticsearchClient elasticsearchClient(RestClient restClient) {
 		}
 
 	}
-
 }