diff --git a/pom.xml b/pom.xml index 8fc852901..a16eaee71 100644 --- a/pom.xml +++ b/pom.xml @@ -27,7 +27,6 @@ models/spring-ai-vertex-ai-palm2 models/spring-ai-vertex-ai-gemini models/spring-ai-anthropic - spring-ai-test spring-ai-spring-boot-autoconfigure spring-ai-spring-boot-starters/spring-ai-starter-openai @@ -67,6 +66,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-mongodb-atlas-store spring-ai-spring-boot-testcontainers spring-ai-spring-boot-starters/spring-ai-starter-anthropic + vector-stores/spring-ai-elasticsearch-store diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index 7457c2798..6d72deb07 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -182,7 +182,13 @@ org.springframework.ai - spring-ai-mongodb-atlas-store + spring-ai-mongodb-atlas-store + ${project.version} + + + + org.springframework.ai + spring-ai-elasticsearch-store ${project.version} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/FilterExpressionConverter.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/FilterExpressionConverter.java index 31a9f6232..127f2dc92 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/FilterExpressionConverter.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/FilterExpressionConverter.java @@ -15,8 +15,6 @@ */ package org.springframework.ai.vectorstore.filter; -import org.springframework.ai.vectorstore.filter.Filter; - /** * Converters a generic, portable {@link Filter.Expression} into a * {@link org.springframework.ai.vectorstore.VectorStore} specific expression language diff --git a/vector-stores/spring-ai-elasticsearch-store/pom.xml b/vector-stores/spring-ai-elasticsearch-store/pom.xml new file mode 100644 index 000000000..11aac09f9 --- /dev/null +++ b/vector-stores/spring-ai-elasticsearch-store/pom.xml @@ -0,0 +1,78 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-elasticsearch-store + jar + Spring AI Vector Store - Elasticsearch + Spring AI Elasticsearch Vector Store + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + 4.0.3 + + + + + org.springframework.ai + spring-ai-core + ${parent.version} + + + + co.elastic.clients + elasticsearch-java + 8.12.2 + + + + + org.springframework.ai + spring-ai-openai + ${parent.version} + test + + + + + org.springframework.ai + spring-ai-test + ${parent.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.testcontainers + elasticsearch + test + + + + org.testcontainers + junit-jupiter + ${testcontainers.version} + test + + + + + diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverter.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverter.java new file mode 100644 index 000000000..e0ce7d620 --- /dev/null +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverter.java @@ -0,0 +1,150 @@ +/* + * Copyright 2023 - 2024 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.vectorstore; + +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.Filter.Expression; +import org.springframework.ai.vectorstore.filter.Filter.Key; +import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter; + +import java.text.ParseException; +import java.text.SimpleDateFormat; +import java.util.Date; +import java.util.List; +import java.util.TimeZone; +import java.util.regex.Pattern; + +/** + * @author Jemin Huh + * @since 1.0.0 + */ +public class ElasticsearchAiSearchFilterExpressionConverter extends AbstractFilterExpressionConverter { + + private static final Pattern DATE_FORMAT_PATTERN = Pattern.compile("\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z"); + + private final SimpleDateFormat dateFormat; + + public ElasticsearchAiSearchFilterExpressionConverter() { + this.dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'"); + this.dateFormat.setTimeZone(TimeZone.getTimeZone("UTC")); + } + + @Override + protected void doExpression(Expression expression, StringBuilder context) { + if (expression.type() == Filter.ExpressionType.IN || expression.type() == Filter.ExpressionType.NIN) { + context.append(getOperationSymbol(expression)); + context.append("("); + this.convertOperand(expression.left(), context); + this.convertOperand(expression.right(), context); + context.append(")"); + } + else { + this.convertOperand(expression.left(), context); + context.append(getOperationSymbol(expression)); + this.convertOperand(expression.right(), context); + } + } + + @Override + protected void doStartValueRange(Filter.Value listValue, StringBuilder context) { + } + + @Override + protected void doEndValueRange(Filter.Value listValue, StringBuilder context) { + } + + @Override + protected void doAddValueRangeSpitter(Filter.Value listValue, StringBuilder context) { + context.append(" OR "); + } + + private String getOperationSymbol(Expression exp) { + return switch (exp.type()) { + case AND -> " AND "; + case OR -> " OR "; + case EQ, IN -> ""; + case NE -> " NOT "; + case LT -> "<"; + case LTE -> "<="; + case GT -> ">"; + case GTE -> ">="; + case NIN -> "NOT "; + default -> throw new RuntimeException("Not supported expression type: " + exp.type()); + }; + } + + @Override + public void doKey(Key key, StringBuilder context) { + var identifier = hasOuterQuotes(key.key()) ? removeOuterQuotes(key.key()) : key.key(); + var prefixedIdentifier = withMetaPrefix(identifier); + context.append(prefixedIdentifier.trim()).append(":"); + } + + public String withMetaPrefix(String identifier) { + return "metadata." + identifier; + } + + @Override + protected void doValue(Filter.Value filterValue, StringBuilder context) { + if (filterValue.value() instanceof List list) { + int c = 0; + for (Object v : list) { + context.append(v); + if (c++ < list.size() - 1) { + this.doAddValueRangeSpitter(filterValue, context); + } + } + } + else { + this.doSingleValue(filterValue.value(), context); + } + } + + @Override + protected void doSingleValue(Object value, StringBuilder context) { + if (value instanceof Date date) { + context.append(this.dateFormat.format(date)); + } + else if (value instanceof String text) { + if (DATE_FORMAT_PATTERN.matcher(text).matches()) { + try { + Date date = this.dateFormat.parse(text); + context.append(this.dateFormat.format(date)); + } + catch (ParseException e) { + throw new IllegalArgumentException("Invalid date type:" + text, e); + } + } + else { + context.append(text); + } + } + else { + context.append(value); + } + } + + @Override + public void doStartGroup(Filter.Group group, StringBuilder context) { + context.append("("); + } + + @Override + public void doEndGroup(Filter.Group group, StringBuilder context) { + context.append(")"); + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java new file mode 100644 index 000000000..75a41a580 --- /dev/null +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java @@ -0,0 +1,215 @@ +/* + * Copyright 2023 - 2024 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.vectorstore; + +import co.elastic.clients.elasticsearch.ElasticsearchClient; +import co.elastic.clients.elasticsearch._types.query_dsl.Query; +import co.elastic.clients.elasticsearch.core.BulkRequest; +import co.elastic.clients.elasticsearch.core.BulkResponse; +import co.elastic.clients.elasticsearch.core.search.Hit; +import co.elastic.clients.elasticsearch.indices.CreateIndexResponse; +import co.elastic.clients.json.JsonData; +import co.elastic.clients.json.jackson.JacksonJsonpMapper; +import co.elastic.clients.transport.endpoints.BooleanResponse; +import co.elastic.clients.transport.rest_client.RestClientTransport; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.elasticsearch.client.RestClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.util.Assert; + +import java.io.IOException; +import java.io.StringReader; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.stream.Collectors; + +/** + * @author Jemin Huh + * @since 1.0.0 + */ +public class ElasticsearchVectorStore implements VectorStore, InitializingBean { + + // divided by 2 to get score in the range [0, 1] + public static final String COSINE_SIMILARITY_FUNCTION = "(cosineSimilarity(params.query_vector, 'embedding') + 1.0) / 2"; + + private static final Logger logger = LoggerFactory.getLogger(ElasticsearchVectorStore.class); + + private static final String INDEX_NAME = "spring-ai-document-index"; + + private final EmbeddingClient embeddingClient; + + private final ElasticsearchClient elasticsearchClient; + + private final String index; + + private final FilterExpressionConverter filterExpressionConverter; + + private String similarityFunction; + + public ElasticsearchVectorStore(RestClient restClient, EmbeddingClient embeddingClient) { + this(INDEX_NAME, restClient, embeddingClient); + } + + public ElasticsearchVectorStore(String index, RestClient restClient, EmbeddingClient embeddingClient) { + Objects.requireNonNull(embeddingClient, "RestClient must not be null"); + Objects.requireNonNull(embeddingClient, "EmbeddingClient must not be null"); + this.elasticsearchClient = new ElasticsearchClient(new RestClientTransport(restClient, new JacksonJsonpMapper( + new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)))); + this.embeddingClient = embeddingClient; + this.index = index; + this.filterExpressionConverter = new ElasticsearchAiSearchFilterExpressionConverter(); + // the potential functions for vector fields at + // https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-script-score-query.html#vector-functions + this.similarityFunction = COSINE_SIMILARITY_FUNCTION; + } + + public ElasticsearchVectorStore withSimilarityFunction(String similarityFunction) { + this.similarityFunction = similarityFunction; + return this; + } + + @Override + public void add(List documents) { + BulkRequest.Builder builkRequestBuilder = new BulkRequest.Builder(); + for (Document document : documents) { + if (Objects.isNull(document.getEmbedding()) || document.getEmbedding().isEmpty()) { + logger.debug("Calling EmbeddingClient for document id = " + document.getId()); + document.setEmbedding(this.embeddingClient.embed(document)); + } + builkRequestBuilder + .operations(op -> op.index(idx -> idx.index(this.index).id(document.getId()).document(document))); + } + bulkRequest(builkRequestBuilder.build()); + } + + @Override + public Optional delete(List idList) { + BulkRequest.Builder builkRequestBuilder = new BulkRequest.Builder(); + for (String id : idList) + builkRequestBuilder.operations(op -> op.delete(idx -> idx.index(this.index).id(id))); + return Optional.of(bulkRequest(builkRequestBuilder.build()).errors()); + } + + private BulkResponse bulkRequest(BulkRequest bulkRequest) { + try { + return this.elasticsearchClient.bulk(bulkRequest); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Override + public List similaritySearch(SearchRequest searchRequest) { + Assert.notNull(searchRequest, "The search request must not be null."); + return similaritySearch(this.embeddingClient.embed(searchRequest.getQuery()), searchRequest.getTopK(), + Double.valueOf(searchRequest.getSimilarityThreshold()).floatValue(), + searchRequest.getFilterExpression()); + } + + public List similaritySearch(List embedding, int topK, double similarityThreshold, + Filter.Expression filterExpression) { + return similaritySearch(new co.elastic.clients.elasticsearch.core.SearchRequest.Builder() + .query(getElasticsearchSimilarityQuery(embedding, filterExpression)) + .size(topK) + .minScore(similarityThreshold) + .build()); + } + + private Query getElasticsearchSimilarityQuery(List embedding, Filter.Expression filterExpression) { + return Query.of(queryBuilder -> queryBuilder.scriptScore(scriptScoreQueryBuilder -> scriptScoreQueryBuilder + .query(queryBuilder2 -> queryBuilder2.queryString(queryStringQuerybuilder -> queryStringQuerybuilder + .query(getElasticsearchQueryString(filterExpression)))) + .script(scriptBuilder -> scriptBuilder + .inline(inlineScriptBuilder -> inlineScriptBuilder.source(this.similarityFunction) + .params("query_vector", JsonData.of(embedding)))))); + } + + private String getElasticsearchQueryString(Filter.Expression filterExpression) { + return Objects.isNull(filterExpression) ? "*" + : this.filterExpressionConverter.convertExpression(filterExpression); + + } + + private List similaritySearch(co.elastic.clients.elasticsearch.core.SearchRequest searchRequest) { + try { + return this.elasticsearchClient.search(searchRequest, Document.class) + .hits() + .hits() + .stream() + .map(this::toDocument) + .collect(Collectors.toList()); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + private Document toDocument(Hit hit) { + Document document = hit.source(); + document.getMetadata().put("distance", 1 - hit.score().floatValue()); + return document; + } + + public boolean exists(String targetIndex) { + try { + BooleanResponse response = this.elasticsearchClient.indices() + .exists(existRequestBuilder -> existRequestBuilder.index(targetIndex)); + return response.value(); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + public CreateIndexResponse createIndexMapping(String index, String mappingJson) { + try { + return this.elasticsearchClient.indices() + .create(createIndexBuilder -> createIndexBuilder.index(index) + .mappings(typeMappingBuilder -> typeMappingBuilder.withJson(new StringReader(mappingJson)))); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Override + public void afterPropertiesSet() { + if (!exists(this.index)) { + createIndexMapping(this.index, """ + { + "properties": { + "embedding": { + "type": "dense_vector", + "dims": 1536, + "index": true, + "similarity": "cosine" + } + } + } + """); + } + } + +} diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverterTest.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverterTest.java new file mode 100644 index 000000000..636761cf6 --- /dev/null +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverterTest.java @@ -0,0 +1,117 @@ +/* + * Copyright 2023 - 2024 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.vectorstore; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; + +import java.util.Date; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.*; + +class ElasticsearchAiSearchFilterExpressionConverterTest { + + final FilterExpressionConverter converter = new ElasticsearchAiSearchFilterExpressionConverter(); + + @Test + public void testDate() { + String vectorExpr = converter.convertExpression(new Filter.Expression(EQ, new Filter.Key("activationDate"), + new Filter.Value(new Date(1704637752148L)))); + assertThat(vectorExpr).isEqualTo("metadata.activationDate:2024-01-07T14:29:12Z"); + + vectorExpr = converter.convertExpression( + new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value("1970-01-01T00:00:02Z"))); + assertThat(vectorExpr).isEqualTo("metadata.activationDate:1970-01-01T00:00:02Z"); + } + + @Test + public void testEQ() { + String vectorExpr = converter + .convertExpression(new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG"))); + assertThat(vectorExpr).isEqualTo("metadata.country:BG"); + } + + @Test + public void tesEqAndGte() { + String vectorExpr = converter.convertExpression(new Filter.Expression(AND, + new Filter.Expression(EQ, new Filter.Key("genre"), new Filter.Value("drama")), + new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020)))); + assertThat(vectorExpr).isEqualTo("metadata.genre:drama AND metadata.year:>=2020"); + } + + @Test + public void tesIn() { + String vectorExpr = converter.convertExpression(new Filter.Expression(IN, new Filter.Key("genre"), + new Filter.Value(List.of("comedy", "documentary", "drama")))); + assertThat(vectorExpr).isEqualTo("(metadata.genre:comedy OR documentary OR drama)"); + } + + @Test + public void testNe() { + String vectorExpr = converter.convertExpression( + new Filter.Expression(OR, new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020)), + new Filter.Expression(AND, + new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")), + new Filter.Expression(NE, new Filter.Key("city"), new Filter.Value("Sofia"))))); + assertThat(vectorExpr).isEqualTo("metadata.year:>=2020 OR metadata.country:BG AND metadata.city: NOT Sofia"); + } + + @Test + public void testGroup() { + String vectorExpr = converter.convertExpression(new Filter.Expression(AND, + new Filter.Group(new Filter.Expression(OR, + new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020)), + new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")))), + new Filter.Expression(NIN, new Filter.Key("city"), new Filter.Value(List.of("Sofia", "Plovdiv"))))); + assertThat(vectorExpr) + .isEqualTo("(metadata.year:>=2020 OR metadata.country:BG) AND NOT (metadata.city:Sofia OR Plovdiv)"); + } + + @Test + public void tesBoolean() { + String vectorExpr = converter.convertExpression(new Filter.Expression(AND, + new Filter.Expression(AND, new Filter.Expression(EQ, new Filter.Key("isOpen"), new Filter.Value(true)), + new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020))), + new Filter.Expression(IN, new Filter.Key("country"), new Filter.Value(List.of("BG", "NL", "US"))))); + + assertThat(vectorExpr) + .isEqualTo("metadata.isOpen:true AND metadata.year:>=2020 AND (metadata.country:BG OR NL OR US)"); + } + + @Test + public void testDecimal() { + String vectorExpr = converter.convertExpression(new Filter.Expression(AND, + new Filter.Expression(GTE, new Filter.Key("temperature"), new Filter.Value(-15.6)), + new Filter.Expression(LTE, new Filter.Key("temperature"), new Filter.Value(20.13)))); + + assertThat(vectorExpr).isEqualTo("metadata.temperature:>=-15.6 AND metadata.temperature:<=20.13"); + } + + @Test + public void testComplexIdentifiers() { + String vectorExpr = converter + .convertExpression(new Filter.Expression(EQ, new Filter.Key("\"country 1 2 3\""), new Filter.Value("BG"))); + assertThat(vectorExpr).isEqualTo("metadata.country 1 2 3:BG"); + + vectorExpr = converter + .convertExpression(new Filter.Expression(EQ, new Filter.Key("'country 1 2 3'"), new Filter.Value("BG"))); + assertThat(vectorExpr).isEqualTo("metadata.country 1 2 3:BG"); + } + +} diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java new file mode 100644 index 000000000..c277ff1cb --- /dev/null +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java @@ -0,0 +1,376 @@ +/* + * Copyright 2023 - 2024 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.vectorstore; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.time.ZonedDateTime; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.TimeUnit; + +import org.apache.http.HttpHost; +import org.awaitility.Awaitility; +import org.elasticsearch.client.RestClient; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.testcontainers.elasticsearch.ElasticsearchContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.shaded.com.fasterxml.jackson.databind.ObjectMapper; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.openai.OpenAiEmbeddingClient; +import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.core.io.DefaultResourceLoader; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.hasSize; + +@Testcontainers +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") +class ElasticsearchVectorStoreIT { + + @Container + private static final ElasticsearchContainer elasticsearchContainer = new ElasticsearchContainer( + "docker.elastic.co/elasticsearch/elasticsearch:8.12.2") + .withEnv("xpack.security.enabled", "false"); + + private static final String DEFAULT = "default cosine similarity"; + + protected final ObjectMapper objectMapper = new ObjectMapper(); + + private List documents = List.of( + new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")), + new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()), + new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2"))); + + @BeforeAll + public static void beforeAll() { + Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS); + Awaitility.setDefaultPollDelay(Duration.ZERO); + Awaitility.setDefaultTimeout(Duration.ofMinutes(1)); + } + + private String getText(String uri) { + var resource = new DefaultResourceLoader().getResource(uri); + try { + return resource.getContentAsString(StandardCharsets.UTF_8); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + private ApplicationContextRunner getContextRunner() { + return new ApplicationContextRunner().withUserConfiguration(TestApplication.class); + } + + @BeforeEach + void cleanDatabase() { + getContextRunner().run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + vectorStore.delete(List.of("_all")); + }); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { DEFAULT, """ + double value = dotProduct(params.query_vector, 'embedding'); + return sigmoid(1, Math.E, -value); + """, "1 / (1 + l1norm(params.query_vector, 'embedding'))", + "1 / (1 + l2norm(params.query_vector, 'embedding'))" }) + public void addAndSearchTest(String similarityFunction) { + + getContextRunner().run(context -> { + ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class); + + if (!DEFAULT.equals(similarityFunction)) { + vectorStore.withSimilarityFunction(similarityFunction); + } + + vectorStore.add(documents); + + Awaitility.await() + .until(() -> vectorStore + .similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThreshold(0)), + hasSize(1)); + + List results = vectorStore + .similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThreshold(0)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); + assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).hasSize(2); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(Document::getId).toList()); + + Awaitility.await() + .until(() -> vectorStore + .similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThreshold(0)), + hasSize(0)); + }); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { DEFAULT, """ + double value = dotProduct(params.query_vector, 'embedding'); + return sigmoid(1, Math.E, -value); + """, "1 / (1 + l1norm(params.query_vector, 'embedding'))", + "1 / (1 + l2norm(params.query_vector, 'embedding'))" }) + public void searchWithFilters(String similarityFunction) { + + getContextRunner().run(context -> { + ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class); + + if (!DEFAULT.equals(similarityFunction)) { + vectorStore.withSimilarityFunction(similarityFunction); + } + + 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))); + var nlDocument = new Document("2", "The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "NL", "activationDate", new Date(2000))); + var bgDocument2 = new Document("3", "The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2023, "activationDate", new Date(3000))); + + vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(5)), hasSize(3)); + + List results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'NL'")); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG'")); + + assertThat(results).hasSize(2); + assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG' && year == 2020")); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country in ['BG']")); + + assertThat(results).hasSize(2); + assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country in ['BG','NL']")); + + assertThat(results).hasSize(3); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country not in ['BG']")); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("NOT(country not in ['BG'])")); + + assertThat(results).hasSize(2); + assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression( + "activationDate > " + ZonedDateTime.parse("1970-01-01T00:00:02Z").toInstant().toEpochMilli())); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument2.getId()); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(Document::getId).toList()); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(1)), hasSize(0)); + }); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { DEFAULT, """ + double value = dotProduct(params.query_vector, 'embedding'); + return sigmoid(1, Math.E, -value); + """, "1 / (1 + l1norm(params.query_vector, 'embedding'))", + "1 / (1 + l2norm(params.query_vector, 'embedding'))" }) + public void documentUpdateTest(String similarityFunction) { + + getContextRunner().run(context -> { + ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class); + if (!DEFAULT.equals(similarityFunction)) { + vectorStore.withSimilarityFunction(similarityFunction); + } + + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Map.of("meta1", "meta1")); + vectorStore.add(List.of(document)); + + Awaitility.await() + .until(() -> vectorStore + .similaritySearch(SearchRequest.query("Spring").withSimilarityThreshold(0).withTopK(5)), + hasSize(1)); + + List results = vectorStore + .similaritySearch(SearchRequest.query("Spring").withSimilarityThreshold(0).withTopK(5)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).containsKey("meta1"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", Map.of("meta2", "meta2")); + + vectorStore.add(List.of(sameIdDocument)); + SearchRequest fooBarSearchRequest = SearchRequest.query("FooBar").withTopK(5); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch(fooBarSearchRequest).get(0).getContent(), + equalTo("The World is Big and Salvation Lurks Around the Corner")); + + results = vectorStore.similaritySearch(fooBarSearchRequest); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + // Remove all documents from the store + vectorStore.delete(List.of(document.getId())); + + Awaitility.await().until(() -> vectorStore.similaritySearch(fooBarSearchRequest), hasSize(0)); + + }); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { DEFAULT, """ + double value = dotProduct(params.query_vector, 'embedding'); + return sigmoid(1, Math.E, -value); + """, "1 / (1 + l1norm(params.query_vector, 'embedding'))", + "1 / (1 + l2norm(params.query_vector, 'embedding'))" }) + public void searchThresholdTest(String similarityFunction) { + + getContextRunner().run(context -> { + ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class); + if (!DEFAULT.equals(similarityFunction)) { + vectorStore.withSimilarityFunction(similarityFunction); + } + + vectorStore.add(documents); + + SearchRequest query = SearchRequest.query("Great Depression") + .withTopK(50) + .withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL); + + Awaitility.await().until(() -> vectorStore.similaritySearch(query), hasSize(3)); + + List fullResult = vectorStore.similaritySearch(query); + + List distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList(); + + assertThat(distances).hasSize(3); + + float threshold = (distances.get(0) + distances.get(1)) / 2; + + List results = vectorStore.similaritySearch( + SearchRequest.query("Great Depression").withTopK(50).withSimilarityThreshold(1 - threshold)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); + assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(Document::getId).toList()); + + Awaitility.await() + .until(() -> vectorStore + .similaritySearch(SearchRequest.query("Great Depression").withTopK(50).withSimilarityThreshold(0)), + hasSize(0)); + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Bean + public ElasticsearchVectorStore vectorStore(EmbeddingClient embeddingClient) { + return new ElasticsearchVectorStore( + RestClient.builder(HttpHost.create(elasticsearchContainer.getHttpHostAddress())).build(), + embeddingClient); + } + + @Bean + public EmbeddingClient embeddingClient() { + return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); + } + + } + +}