diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/pom.xml b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/pom.xml new file mode 100644 index 00000000000..3ee702e55b5 --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/pom.xml @@ -0,0 +1,130 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai-parent + 2.0.0-SNAPSHOT + ../../../pom.xml + + spring-ai-autoconfigure-vector-store-infinispan + jar + Spring AI Auto Configuration for Infinispan vector store + Spring AI Auto Configuration for Infinispan 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 + + + + + + org.infinispan + infinispan-bom + ${infinispan.version} + pom + import + + + + + + + + org.springframework.ai + spring-ai-infinispan-store + ${project.parent.version} + true + + + org.springframework.boot + spring-boot-starter + + + org.infinispan + infinispan-spring-boot4-starter-remote + true + + + org.springframework.boot + spring-boot-configuration-processor + true + + + org.springframework.boot + spring-boot-autoconfigure-processor + true + + + + org.springframework.ai + spring-ai-test + ${project.parent.version} + test + + + org.springframework.boot + spring-boot-starter-test + test + + + org.springframework.boot + spring-boot-testcontainers + test + + + org.infinispan + testcontainers-infinispan + test + + + org.testcontainers + testcontainers-junit-jupiter + test + + + org.awaitility + awaitility + test + + + org.springframework.ai + spring-ai-transformers + ${project.parent.version} + test + + + + + + + + + + + + + + + diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreAutoConfiguration.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreAutoConfiguration.java new file mode 100644 index 00000000000..381bae1ed08 --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreAutoConfiguration.java @@ -0,0 +1,69 @@ +/* + * 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. + * 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.infinispan.autoconfigure; + +import io.micrometer.observation.ObservationRegistry; +import org.infinispan.client.hotrod.RemoteCacheManager; + +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.SpringAIVectorStoreTypes; +import org.springframework.ai.vectorstore.infinispan.InfinispanVectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +/** + * {@link AutoConfiguration Auto-configuration} for Infinispan Vector Store. + * + * @author Katia Aresti + */ +@AutoConfiguration +@ConditionalOnClass({ InfinispanVectorStore.class, EmbeddingModel.class }) +@EnableConfigurationProperties(InfinispanVectorStoreProperties.class) +@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.GEMFIRE, + matchIfMissing = true) +public class InfinispanVectorStoreAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public InfinispanVectorStore infinispanVectorStore(EmbeddingModel embeddingModel, + InfinispanVectorStoreProperties properties, RemoteCacheManager infinispanClient, + ObjectProvider observationRegistry, + ObjectProvider customObservationConvention) { + + return InfinispanVectorStore.builder(infinispanClient, embeddingModel) + .createStore(properties.isCreateStore()) + .registerSchema(properties.isRegisterSchema()) + .schemaFileName(properties.getSchemaFileName()) + .storeName(properties.getStoreName()) + .storeConfig(properties.getStoreConfig()) + .observationRegistry(observationRegistry.getIfAvailable()) + .customObservationConvention(customObservationConvention.getIfAvailable()) + .distance(properties.getDistance()) + .similarity(properties.getSimilarity()) + .packageName(properties.getPackageName()) + .springAiItemName(properties.getItemName()) + .metadataItemName(properties.getMetadataItemName()) + .build(); + } + +} diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreProperties.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreProperties.java new file mode 100644 index 00000000000..6103856f043 --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreProperties.java @@ -0,0 +1,135 @@ +/* + * 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. + * 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.infinispan.autoconfigure; + +import org.springframework.ai.vectorstore.properties.CommonVectorStoreProperties; +import org.springframework.boot.context.properties.ConfigurationProperties; + +import static org.springframework.ai.vectorstore.infinispan.autoconfigure.InfinispanVectorStoreProperties.CONFIG_PREFIX; + +/** + * Configuration properties for Infinispan Vector Store. + */ +@ConfigurationProperties(prefix = CONFIG_PREFIX) +public class InfinispanVectorStoreProperties extends CommonVectorStoreProperties { + + /** + * Configuration prefix for Spring AI VectorStore Infinispan. + */ + public static final String CONFIG_PREFIX = "spring.ai.vectorstore.infinispan"; + + private Boolean registerSchema; + + private Boolean createStore; + + private String storeName; + + private String storeConfig; + + private Integer distance; + + private String similarity; + + private String schemaFileName; + + private String packageName; + + private String itemName; + + private String metadataItemName; + + public String getStoreName() { + return storeName; + } + + public void setStoreName(String storeName) { + this.storeName = storeName; + } + + public String getStoreConfig() { + return storeConfig; + } + + public void setStoreConfig(String storeConfig) { + this.storeConfig = storeConfig; + } + + public Integer getDistance() { + return distance; + } + + public void setDistance(Integer distance) { + this.distance = distance; + } + + public String getSimilarity() { + return similarity; + } + + public void setSimilarity(String similarity) { + this.similarity = similarity; + } + + public String getSchemaFileName() { + return schemaFileName; + } + + public void setSchemaFileName(String schemaFileName) { + this.schemaFileName = schemaFileName; + } + + public String getPackageName() { + return packageName; + } + + public void setPackageName(String packageName) { + this.packageName = packageName; + } + + public String getItemName() { + return itemName; + } + + public void setItemName(String itemName) { + this.itemName = itemName; + } + + public String getMetadataItemName() { + return metadataItemName; + } + + public void setMetadataItemName(String metadataItemName) { + this.metadataItemName = metadataItemName; + } + + public Boolean isRegisterSchema() { + return registerSchema; + } + + public void setRegisterSchema(Boolean registerSchema) { + this.registerSchema = registerSchema; + } + + public Boolean isCreateStore() { + return createStore; + } + + public void setCreateStore(boolean createStore) { + this.createStore = createStore; + } + +} diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports new file mode 100644 index 00000000000..0d9ec659b4d --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -0,0 +1,16 @@ +# +# Copyright 2025-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. +# 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. +# +org.springframework.ai.vectorstore.infinispan.autoconfigure.InfinispanVectorStoreAutoConfiguration diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/test/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreAutoConfigurationIT.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/test/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreAutoConfigurationIT.java new file mode 100644 index 00000000000..b76d1acf0ab --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan/src/test/java/org/springframework/ai/vectorstore/infinispan/autoconfigure/InfinispanVectorStoreAutoConfigurationIT.java @@ -0,0 +1,227 @@ +/* + * 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.infinispan.autoconfigure; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import io.micrometer.observation.tck.TestObservationRegistry; +import org.awaitility.Awaitility; +import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.client.hotrod.RemoteCacheManager; +import org.infinispan.commons.marshall.ProtoStreamMarshaller; +import org.infinispan.commons.util.Version; +import org.infinispan.protostream.schema.Schema; +import org.infinispan.spring.starter.remote.InfinispanRemoteAutoConfiguration; +import org.infinispan.testcontainers.InfinispanContainer; +import org.jetbrains.annotations.NotNull; +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.observation.conventions.VectorStoreProvider; +import org.springframework.ai.test.vectorstore.ObservationTestUtil; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.util.ResourceUtils; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.ai.vectorstore.infinispan.InfinispanVectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.core.io.DefaultResourceLoader; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.hamcrest.Matchers.hasSize; + +/** + * @author Katia Aresti + */ +@Testcontainers +class InfinispanVectorStoreAutoConfigurationIT { + + @Container + private static final InfinispanContainer infinispanContainer = new InfinispanContainer( + InfinispanContainer.IMAGE_BASENAME + ":" + Version.getVersion()); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(InfinispanRemoteAutoConfiguration.class, + InfinispanVectorStoreAutoConfiguration.class)) + .withUserConfiguration(Config.class) + .withPropertyValues("spring.ai.vectorstore.infinispan.distance=" + 10, + "infinispan.remote.server-list=" + serverList(), + "infinispan.remote.auth-username=" + InfinispanContainer.DEFAULT_USERNAME, + // Needs the marshalling property until fix + // https://github.com/infinispan/infinispan/issues/16440 + "infinispan.remote.marshaller=" + ProtoStreamMarshaller.class.getName(), + "infinispan.remote.auth-password=" + InfinispanContainer.DEFAULT_PASSWORD); + + List documents = List.of( + new Document(ResourceUtils.getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")), + new Document(ResourceUtils.getText("classpath:/test/data/time.shelter.txt")), new Document( + ResourceUtils.getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad"))); + + @Test + public void addAndSearchTest() { + this.contextRunner.run(context -> { + InfinispanVectorStore vectorStore = context.getBean(InfinispanVectorStore.class); + TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class); + + vectorStore.add(this.documents); + + ObservationTestUtil.assertObservationRegistry(observationRegistry, VectorStoreProvider.INFINISPAN, + VectorStoreObservationContext.Operation.ADD); + observationRegistry.clear(); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThreshold(0).build()), + hasSize(1)); + + observationRegistry.clear(); + + List results = vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThreshold(0).build()); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).hasSize(1); + + ObservationTestUtil.assertObservationRegistry(observationRegistry, VectorStoreProvider.INFINISPAN, + VectorStoreObservationContext.Operation.QUERY); + observationRegistry.clear(); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(Document::getId).toList()); + + ObservationTestUtil.assertObservationRegistry(observationRegistry, VectorStoreProvider.INFINISPAN, + VectorStoreObservationContext.Operation.DELETE); + observationRegistry.clear(); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThreshold(0).build()), + hasSize(0)); + }); + } + + @Test + public void propertiesTest() { + new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(InfinispanVectorStoreAutoConfiguration.class, + InfinispanRemoteAutoConfiguration.class)) + .withUserConfiguration(Config.class) + .withPropertyValues("infinispan.remote.server-list=" + serverList(), + "infinispan.remote.auth-username=" + InfinispanContainer.DEFAULT_USERNAME, + "infinispan.remote.auth-password=" + InfinispanContainer.DEFAULT_PASSWORD, + "spring.ai.vectorstore.infinispan.distance=20", + "spring.ai.vectorstore.infinispan.item-name=ItemExample", + "spring.ai.vectorstore.infinispan.metadata-item-name=MetadataExample", + "spring.ai.vectorstore.infinispan.package-name=exam.pac", + "spring.ai.vectorstore.infinispan.store-name=mycoolstore", + "spring.ai.vectorstore.infinispan.schema-file-name=schemaName.proto", + "spring.ai.vectorstore.infinispan.similarity=cosine") + .run(context -> { + var properties = context.getBean(InfinispanVectorStoreProperties.class); + assertThat(properties).isNotNull(); + assertThat(properties.getDistance()).isEqualTo(20); + assertThat(properties.getItemName()).isEqualTo("ItemExample"); + assertThat(properties.getMetadataItemName()).isEqualTo("MetadataExample"); + assertThat(properties.getSimilarity()).isEqualTo("cosine"); + + InfinispanVectorStore infinispanVectorStore = context.getBean(InfinispanVectorStore.class); + assertThat(infinispanVectorStore).isNotNull(); + + RemoteCacheManager cacheManager = context.getBean(RemoteCacheManager.class); + RemoteCache cache = cacheManager.getCache("mycoolstore"); + assertThat(cache).isNotNull(); + Optional schema = cacheManager.administration().schemas().get("schemaName.proto"); + assertThat(schema).isNotEmpty(); + String schemaContent = schema.get().getContent(); + assertThat(schemaContent).contains("ItemExample"); + assertThat(schemaContent).contains("MetadataExample"); + assertThat(schemaContent).contains("package exam.pac"); + }); + } + + @Test + public void autoConfigurationDisabledWhenTypeIsNone() { + this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=none").run(context -> { + assertThat(context.getBeansOfType(InfinispanVectorStoreProperties.class)).isEmpty(); + assertThat(context.getBeansOfType(InfinispanVectorStore.class)).isEmpty(); + assertThat(context.getBeansOfType(VectorStore.class)).isEmpty(); + }); + } + + @Test + public void autoConfigurationEnabledByDefault() { + this.contextRunner.run(context -> { + assertThat(context.getBeansOfType(InfinispanVectorStoreProperties.class)).isNotEmpty(); + assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty(); + assertThat(context.getBean(VectorStore.class)).isInstanceOf(InfinispanVectorStore.class); + }); + } + + @Test + public void autoConfigurationEnabledWhenTypeIsInfinispan() { + this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=infinispan").run(context -> { + assertThat(context.getBeansOfType(InfinispanVectorStoreProperties.class)).isNotEmpty(); + assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty(); + assertThat(context.getBean(VectorStore.class)).isInstanceOf(InfinispanVectorStore.class); + }); + } + + private static @NotNull String serverList() { + return infinispanContainer.getHost() + ":" + + infinispanContainer.getMappedPort(InfinispanContainer.DEFAULT_HOTROD_PORT); + } + + 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); + } + } + + @Configuration(proxyBeanMethods = false) + static class Config { + + @Bean + public TestObservationRegistry observationRegistry() { + return TestObservationRegistry.create(); + } + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } + + } + +} diff --git a/pom.xml b/pom.xml index 190f086a801..2d1dc175faa 100644 --- a/pom.xml +++ b/pom.xml @@ -69,6 +69,7 @@ vector-stores/spring-ai-elasticsearch-store vector-stores/spring-ai-gemfire-store vector-stores/spring-ai-hanadb-store + vector-stores/spring-ai-infinispan-store vector-stores/spring-ai-mariadb-store vector-stores/spring-ai-milvus-store vector-stores/spring-ai-mongodb-atlas-store @@ -137,6 +138,7 @@ auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-couchbase auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-elasticsearch auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-gemfire + auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-infinispan auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-mariadb auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-milvus auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-mongodb-atlas @@ -333,6 +335,7 @@ 3.5.3 9.2.0 0.22.0 + 16.0.5 3.9.1 2024.5.1 diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index dde9366f78e..9cb308e1bb6 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -479,6 +479,12 @@ ${project.version} + + org.springframework.ai + spring-ai-infinispan-store + ${project.version} + + @@ -764,6 +770,12 @@ ${project.version} + + org.springframework.ai + spring-ai-autoconfigure-vector-store-infinispan + ${project.version} + + org.springframework.ai spring-ai-autoconfigure-vector-store-mariadb diff --git a/spring-ai-commons/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java b/spring-ai-commons/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java index a9c3e3f2b12..3666c4ba5c8 100644 --- a/spring-ai-commons/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java +++ b/spring-ai-commons/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java @@ -70,6 +70,11 @@ public enum VectorStoreProvider { */ HANA("hana"), + /** + * Vector store provided by Infinispan. + */ + INFINISPAN("infinispan"), + /** * Vector store provided by MariaDB. */ diff --git a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SpringAIVectorStoreTypes.java b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SpringAIVectorStoreTypes.java index 4e8f8913b7a..1403f8e6854 100644 --- a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SpringAIVectorStoreTypes.java +++ b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SpringAIVectorStoreTypes.java @@ -40,6 +40,8 @@ private SpringAIVectorStoreTypes() { public static final String HANADB = "hanadb"; + public static final String INFINISPAN = "infinispan"; + public static final String MARIADB = "mariadb"; public static final String MILVUS = "milvus"; diff --git a/vector-stores/spring-ai-infinispan-store/pom.xml b/vector-stores/spring-ai-infinispan-store/pom.xml new file mode 100644 index 00000000000..4df72a86281 --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/pom.xml @@ -0,0 +1,127 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai-parent + 2.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-infinispan-store + jar + Spring AI Vector Store - Infinispan + Spring AI Infinispan 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 + + + + 17 + 17 + + + + + + org.infinispan + infinispan-bom + ${infinispan.version} + pom + import + + + + + + + org.springframework.ai + spring-ai-vector-store + ${project.parent.version} + + + + org.infinispan + infinispan-client-hotrod + ${infinispan.version} + + + + + org.springframework.ai + spring-ai-transformers + ${project.parent.version} + test + + + + org.springframework.ai + spring-ai-test + ${project.parent.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.infinispan + testcontainers-infinispan + ${infinispan.version} + test + + + + org.testcontainers + testcontainers-junit-jupiter + test + + + + io.micrometer + micrometer-observation-test + test + + + + org.apache.commons + commons-lang3 + test + + + + + + + + + + + + + + + diff --git a/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/InfinispanFilterExpressionConverter.java b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/InfinispanFilterExpressionConverter.java new file mode 100644 index 00000000000..7af2df8488c --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/InfinispanFilterExpressionConverter.java @@ -0,0 +1,273 @@ +/* + * 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. + * 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.infinispan; + +import java.time.Instant; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.time.format.DateTimeParseException; +import java.util.Collection; +import java.util.Date; +import java.util.List; +import java.util.regex.Pattern; +import java.util.stream.Collectors; + +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter; + +class InfinispanFilterExpressionConverter 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 static final DateTimeFormatter DATE_TIME_FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss'Z'") + .withZone(ZoneOffset.UTC); + + private int i = -1; + + public String doJoin() { + StringBuilder sb = new StringBuilder(); + for (int j = 0; j <= this.i; j++) { + sb.append(" join i.metadata m").append(j); + } + return sb.toString(); + } + + @Override + protected void doExpression(Filter.Expression expression, StringBuilder context) { + switch (expression.type()) { + case AND: + context.append("(("); + doExpression((Filter.Expression) expression.left(), context); + context.append(") AND ("); + doExpression((Filter.Expression) expression.right(), context); + context.append("))"); + break; + case OR: + context.append("(("); + doExpression((Filter.Expression) expression.left(), context); + context.append(") OR ("); + doExpression((Filter.Expression) expression.right(), context); + context.append("))"); + break; + default: + doField(expression, context); + break; + } + } + + @Override + protected void doKey(Filter.Key filterKey, StringBuilder context) { + + } + + @Override + public void doStartGroup(Filter.Group group, StringBuilder context) { + context.append("("); + } + + @Override + public void doEndGroup(Filter.Group group, StringBuilder context) { + context.append(")"); + } + + private void doField(Filter.Expression expression, StringBuilder context) { + Filter.Key key = (Filter.Key) expression.left(); + Filter.Value value = (Filter.Value) expression.right(); + String result = switch (expression.type()) { + case EQ -> mapEqual(key, value, true); + case NE -> mapEqual(key, value, false); + case GT -> mapGreaterThan(key, value); + case GTE -> mapGreaterThanEqual(key, value); + case LT -> mapLessThan(key, value); + case LTE -> mapLessThanEqual(key, value); + case IN -> mapIn(key, value, true); + case NIN -> mapIn(key, value, false); + case ISNULL -> mapIsNull(key); + case ISNOTNULL -> mapIsNotNull(key); + default -> throw new UnsupportedOperationException("Unsupported value: " + expression.type()); + }; + context.append(result); + } + + private String mapIsNull(Filter.Key key) { + incrementJoin(); + String m = "m" + this.i + "."; + return "(" + metadataKey(key) + String.format( + "%svalue IS NULL and %svalue_int IS NULL and %svalue_date IS NULL and %svalue_float IS NULL and %svalue_bool IS NULL)", + m, m, m, m, m) + " OR (" + m + "name" + " NOT IN('" + key.key() + "'))"; + } + + private String mapIsNotNull(Filter.Key key) { + incrementJoin(); + String m = "m" + this.i + "."; + return metadataKey(key) + String.format( + "(%svalue IS NOT NULL or %svalue_int IS NOT NULL or %svalue_date IS NOT NULL or %svalue_float IS NOT NULL or %svalue_bool IS NOT NULL)", + m, m, m, m, m); + } + + private String mapEqual(Filter.Key key, Filter.Value value, boolean equals) { + incrementJoin(); + String filter = metadataKey(key) + computeValue(equals ? "=" : "!=", value.value()); + if (equals) { + return filter; + } + return filter + " " + addMetadataNullCheck(); + } + + private String mapGreaterThan(Filter.Key key, Filter.Value value) { + incrementJoin(); + return metadataKey(key) + computeValue(">", value.value()); + } + + private String mapGreaterThanEqual(Filter.Key key, Filter.Value value) { + incrementJoin(); + return metadataKey(key) + computeValue(">=", value.value()); + } + + private String mapLessThan(Filter.Key key, Filter.Value value) { + incrementJoin(); + return metadataKey(key) + computeValue("<", value.value()); + } + + private String mapLessThanEqual(Filter.Key key, Filter.Value value) { + incrementJoin(); + return metadataKey(key) + computeValue("<=", value.value()); + } + + private String mapIn(Filter.Key key, Filter.Value value, boolean in) { + incrementJoin(); + String inStatement; + Object first; + if (value.value() instanceof List values) { + if (values.isEmpty()) { + throw new UnsupportedOperationException("Infinispan metadata filter IN must contain values"); + } + first = values.get(0); + inStatement = formattedComparisonValues(values); + } + else { + // single value + first = value.value(); + inStatement = first instanceof String ? "'" + value.value() + "'" : value.value().toString(); + } + + String m = "m" + this.i + "."; + String inFilter = m + "value IN (" + inStatement + ")"; + if (first instanceof Integer || first instanceof Long) { + inFilter = m + "value_int IN (" + inStatement + ")"; + } + else if (first instanceof Float || first instanceof Double) { + inFilter = m + "value_float IN (" + inStatement + ")"; + } + else if (first instanceof Boolean) { + inFilter = m + "value_bool IN (" + inStatement + ")"; + } + else if (first instanceof Date) { + inFilter = m + "value_date IN (" + inStatement + ")"; + } + + if (in) { + return metadataKey(key) + inFilter; + } + + String notInFilter = m + "value NOT IN (" + inStatement + ")"; + if (first instanceof Integer || first instanceof Long) { + notInFilter = m + "value_int NOT IN (" + inStatement + ")"; + } + else if (first instanceof Float || first instanceof Double) { + notInFilter = m + "value_float NOT IN (" + inStatement + ")"; + } + + return "(" + notInFilter + metadataKeyLast(key) + ") " + "OR (" + inFilter + " and " + m + "name!='" + key.key() + + "')" + " " + addMetadataNullCheck(); + } + + private String metadataKey(Filter.Key key) { + return "m" + this.i + ".name='" + key.key() + "' and "; + } + + private String metadataKeyLast(Filter.Key key) { + return " and m" + this.i + ".name='" + key.key() + "' "; + } + + private String computeValue(String operator, Object value) { + String m = "m" + this.i + "."; + String filterQuery = ""; + if (value instanceof String text && DATE_FORMAT_PATTERN.matcher(text).matches()) { + try { + filterQuery = m + "value_date" + operator + + Instant.from(DATE_TIME_FORMATTER.parse(text)).toEpochMilli(); + } + catch (DateTimeParseException e) { + throw new IllegalArgumentException("Invalid date type:" + text, e); + } + } + else if (value instanceof Integer || value instanceof Long) { + Long longValue = getLongValue(value); + filterQuery = m + "value_int" + operator + longValue; + } + else if (value instanceof Float || value instanceof Double) { + Double doubleValue = getDoubleValue(value); + filterQuery = m + "value_float" + operator + doubleValue; + } + else if (value instanceof Date || value instanceof Instant) { + filterQuery = m + "value_date" + operator + getDate(value); + } + else if (value instanceof Boolean bool) { + filterQuery = m + "value_bool" + operator + bool.booleanValue(); + } + else { + // Any other case + filterQuery = m + "value" + operator + "'" + value + "'"; + } + return filterQuery; + } + + private long getDate(Object value) { + if (value instanceof Date date) { + return date.toInstant().toEpochMilli(); + } + if (value instanceof Instant instant) { + return instant.toEpochMilli(); + } + return 0L; + } + + private Long getLongValue(Object value) { + return value instanceof Integer ? ((Integer) value).longValue() : (Long) value; + } + + private Double getDoubleValue(Object value) { + return value instanceof Float ? ((Float) value).doubleValue() : (Double) value; + } + + private String formattedComparisonValues(Collection comparisonValues) { + String inStatement = comparisonValues.stream() + .map(s -> s instanceof String || s instanceof Date ? "'" + s + "'" : s.toString()) + .collect(Collectors.joining(", ")); + return inStatement; + } + + private String addMetadataNullCheck() { + return "OR (i.metadata is null)"; + } + + private void incrementJoin() { + this.i++; + } + +} diff --git a/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStore.java b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStore.java new file mode 100644 index 00000000000..38e4e198c11 --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStore.java @@ -0,0 +1,472 @@ +/* + * 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. + * 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.infinispan; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.client.hotrod.RemoteCacheManager; +import org.infinispan.commons.api.query.Query; +import org.infinispan.commons.configuration.StringConfiguration; +import org.infinispan.commons.marshall.ProtoStreamMarshaller; +import org.infinispan.protostream.schema.Field; +import org.infinispan.protostream.schema.Schema; +import org.infinispan.protostream.schema.Type; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.embedding.EmbeddingOptions; +import org.springframework.ai.observation.conventions.VectorStoreProvider; +import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; +import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.util.Assert; + +public class InfinispanVectorStore extends AbstractObservationVectorStore implements InitializingBean { + + /** + * Default Store Name + */ + public static final String DEFAULT_STORE_NAME = "defaultStore"; + + /** + * Default Cache Config + */ + public static final String DEFAULT_CACHE_CONFIG = "\n" + + "\n" + "\n" + + "SPRING_AI_ITEM\n" + "\n" + "\n" + + ""; + + /** + * Default package of the schema + */ + public static final String DEFAULT_PACKAGE = "dev.spring_ai"; + + /** + * Default name of the protobuf springAi item. Size will be added + */ + public static final String DEFAULT_ITEM_NAME = "SpringAiItem"; + + /** + * Default name of the protobuf metadata item. Size will be added + */ + public static final String DEFAULT_METADATA_ITEM = "SpringAiMetadata"; + + /** + * The default distance to for the search + */ + public static final int DEFAULT_DISTANCE = 3; + + /** + * Default vector similarity + */ + public static final String DEFAULT_SIMILARITY = VectorStoreSimilarityMetric.COSINE.value(); + + private final RemoteCacheManager infinispanClient; + + private final String storeName; + + private final String storeConfig; + + private final int distance; + + private final int dimension; + + private final String similarity; + + private final String schemaFileName; + + private final String packageName; + + private final String itemName; + + private final String metadataItemName; + + private final boolean registerSchema; + + private final boolean createStore; + + private final String itemFullName; + + private final String metadataFullName; + + private RemoteCache remoteCache; + + protected InfinispanVectorStore(Builder builder) { + super(builder); + Assert.notNull(builder.infinispanClient, "infinispanClientBuilder must not be null"); + Assert.notNull(builder.dimension, "dimension must not be null"); + Assert.isTrue(builder.distance == null || (builder.distance != null && builder.distance > 0), + "provided distance must be greater than 0"); + this.infinispanClient = builder.infinispanClient; + this.dimension = builder.dimension; + this.storeConfig = builder.storeConfig; + this.createStore = builder.createStore == null ? true : builder.createStore; + this.storeName = builder.storeName == null ? DEFAULT_STORE_NAME : builder.storeName; + this.distance = builder.distance == null ? DEFAULT_DISTANCE : builder.distance; + this.similarity = builder.similarity == null ? DEFAULT_SIMILARITY : builder.similarity; + this.packageName = builder.packageName == null ? DEFAULT_PACKAGE : builder.packageName; + this.itemName = builder.itemName == null ? DEFAULT_ITEM_NAME : builder.itemName; + this.metadataItemName = builder.metadataItemName == null ? DEFAULT_METADATA_ITEM : builder.metadataItemName; + this.registerSchema = builder.registerSchema == null ? true : builder.registerSchema; + this.schemaFileName = getSchemaFileName(builder); + this.itemFullName = computeProtoFullName(this.itemName); + this.metadataFullName = computeProtoFullName(this.metadataItemName); + } + + private String getSchemaFileName(Builder builder) { + if (builder.schemaFileName != null) { + return builder.schemaFileName; + } + return builder.packageName + "." + "dimension." + builder.dimension + ".proto"; + } + + private String computeProtoFullName(String name) { + return this.packageName + "." + name; + } + + @Override + public void doAdd(List documents) { + Map elements = new HashMap<>(documents.size()); + List embeddings = this.embeddingModel.embed(documents, EmbeddingOptions.builder().build(), + this.batchingStrategy); + + for (int i = 0; i < embeddings.size(); i++) { + Document document = documents.get(i); + float[] vector = embeddings.get(i); + Set metadataSet = document.getMetadata() + .entrySet() + .stream() + .map(e -> new SpringAiMetadata(e.getKey(), e.getValue())) + .collect(Collectors.toSet()); + elements.put(document.getId(), new SpringAiInfinispanItem(document.getId(), document.getText(), metadataSet, + vector, document.getMetadata())); + } + + this.remoteCache.putAll(elements); + } + + @Override + public void doDelete(List idList) { + if (idList == null || idList.isEmpty()) { + throw new IllegalArgumentException("ids cannot be null or empty"); + } + + for (String id : idList) { + this.remoteCache.remove(id); + } + } + + @Override + public void doDelete(Filter.Expression filterExpression) { + InfinispanFilterExpressionConverter filterExpressionConverter = new InfinispanFilterExpressionConverter(); + String filteringPart = filterExpressionConverter.convertExpression(filterExpression); + String joinPart = filterExpressionConverter.doJoin(); + String deleteQuery = "DELETE FROM " + this.itemFullName + " i " + joinPart + " where " + filteringPart; + Query query = this.remoteCache.query(deleteQuery); + query.execute(); + } + + @Override + public List doSimilaritySearch(SearchRequest searchRequest) { + String joinPart = ""; + String filteringPart = ""; + + if (searchRequest.hasFilterExpression()) { + InfinispanFilterExpressionConverter filterExpressionConverter = new InfinispanFilterExpressionConverter(); + filteringPart = "filtering(" + + filterExpressionConverter.convertExpression(searchRequest.getFilterExpression()) + ")"; + joinPart = filterExpressionConverter.doJoin(); + } + + var embedding = this.embeddingModel.embed(searchRequest.getQuery()); + String vectorQuery = "select i, score(i) from " + this.itemFullName + " i " + joinPart + + " where i.embedding <-> " + Arrays.toString(embedding) + "~" + this.distance + " " + filteringPart; + + Query query = this.remoteCache.query(vectorQuery); + List hits = query.maxResults(searchRequest.getTopK()).list(); + + return hits.stream().map(obj -> { + SpringAiInfinispanItem item = (SpringAiInfinispanItem) obj[0]; + Float score = (Float) obj[1]; + if (score.doubleValue() < searchRequest.getSimilarityThreshold()) { + return null; + } + + return Document.builder() + .id(item.id()) + .text(item.text()) + .metadata(item.metadataMap()) + .score(score.doubleValue()) + .build(); + }).filter(Objects::nonNull).collect(Collectors.toList()); + } + + @Override + public void afterPropertiesSet() { + Schema schema = buildSchema(); + // Register the schema and marshaller on client side + ProtoStreamMarshaller marshaller = (ProtoStreamMarshaller) this.infinispanClient.getMarshallerRegistry() + .getMarshaller(ProtoStreamMarshaller.class); + if (marshaller == null) { + throw new IllegalStateException("ProtoStreamMarshaller not found"); + } + marshaller.register(schema, new SpringAiMetadataMarshaller(this.metadataFullName), + new SpringAiItemMarshaller(this.itemFullName)); + + // Uploads the schema to the server, if necessary + if (this.registerSchema) { + this.infinispanClient.administration().schemas().createOrUpdate(schema); + } + + // Check if the schema is present + if (this.infinispanClient.administration().schemas().get(this.schemaFileName).isEmpty()) { + throw new IllegalStateException("SpringAI Schema '" + this.schemaFileName + "' not found"); + } + ProtoStreamMarshaller finalMarshaller = marshaller; + // Make sure the marshaller is Protostream on the client side + infinispanClient.getConfiguration().addRemoteCache(this.storeName, c -> { + c.marshaller(finalMarshaller); + }); + + // Get the underlying infinispan remote cache where the embeddings are stored + this.remoteCache = this.infinispanClient.getCache(this.storeName); + if (this.remoteCache == null && this.createStore) { + String infinispanCacheConfig = this.storeConfig; + if (infinispanCacheConfig == null) { + infinispanCacheConfig = DEFAULT_CACHE_CONFIG.replace("CACHE_NAME", this.storeName) + .replace("SPRING_AI_ITEM", this.itemFullName); + } + this.remoteCache = this.infinispanClient.administration() + .getOrCreateCache(this.storeName, new StringConfiguration(infinispanCacheConfig)); + } + + if (this.remoteCache == null) { + throw new IllegalStateException("Infinispan Cache '" + this.storeName + "' not found"); + } + } + + private Schema buildSchema() { + Field.Builder schemaBuilder = new Schema.Builder(this.schemaFileName).packageName(this.packageName) + // Medata Item + .addMessage(this.metadataItemName) + .addComment("@Indexed") + .addField(Type.Scalar.STRING, "name", 1) + .addComment("@Basic(projectable=true)") + .addField(Type.Scalar.STRING, "value", 2) + .addComment("@Basic(projectable=true)") + .addField(Type.Scalar.INT64, "value_int", 3) + .addComment("@Basic(projectable=true)") + .addField(Type.Scalar.DOUBLE, "value_float", 4) + .addComment("@Basic(projectable=true)") + .addField(Type.Scalar.BOOL, "value_bool", 5) + .addComment("@Basic(projectable=true)") + .addField(Type.Scalar.FIXED64, "value_date", 6) + .addComment("@Basic(projectable=true)") + // SpringAi item + .addMessage(this.itemName) + .addComment("@Indexed") + .addField(Type.Scalar.STRING, "id", 1) + .addComment("@Basic(projectable=true)") + .addField(Type.Scalar.STRING, "text", 2) + .addComment("@Basic(projectable=true)"); + + // Add metadata field + schemaBuilder.addRepeatedField(Type.create(this.metadataItemName), "metadata", 3).addComment("@Embedded"); + + // Add embedding + schemaBuilder.addRepeatedField(Type.Scalar.FLOAT, "embedding", 4) + .addComment(String.format("@Vector(dimension=%d, similarity=%s)", this.dimension, + this.similarity.toUpperCase())); + return schemaBuilder.build(); + } + + @Override + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { + return VectorStoreObservationContext.builder(VectorStoreProvider.INFINISPAN.value(), operationName) + .collectionName(this.storeName) + .dimensions(this.embeddingModel.dimensions()) + .similarityMetric(this.similarity); + } + + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.remoteCache; + return Optional.of(client); + } + + /** + * Creates a new builder instance for InfinispanVectorStore. + * @return a new InfinispanBuilder instance + */ + public static Builder builder(RemoteCacheManager infinispanClient, EmbeddingModel embeddingModel) { + return new Builder(infinispanClient, embeddingModel); + } + + public void clear() { + this.remoteCache.clear(); + } + + public static class Builder extends AbstractVectorStoreBuilder { + + /** + * Model dimension + */ + private final Integer dimension; + + private RemoteCacheManager infinispanClient; + + // Store name and configuration + private Boolean createStore; + + private String storeName; + + private String storeConfig; + + // Vector properties + private Integer distance; + + private String similarity; + + // Schema properties + private String schemaFileName; + + private String packageName; + + private String itemName; + + private String metadataItemName; + + private Boolean registerSchema; + + /** + * Infinispan store name to be used, will be created on first access + */ + public Builder storeName(String name) { + this.storeName = name; + return this; + } + + /** + * Infinispan cache config to be used, will be created on first access + */ + public Builder storeConfig(String storeConfig) { + this.storeConfig = storeConfig; + return this; + } + + /** + * Infinispan distance for knn query + */ + public Builder distance(Integer distance) { + this.distance = distance; + return this; + } + + /** + * Infinispan similarity for the embedding definition + */ + public Builder similarity(String similarity) { + this.similarity = similarity; + return this; + } + + /** + * Infinispan schema package name + */ + public Builder packageName(String packageName) { + this.packageName = packageName; + return this; + } + + /** + * Infinispan schema itemName + */ + public Builder springAiItemName(String itemName) { + this.itemName = itemName; + return this; + } + + /** + * Infinispan schema metadataItemName + */ + public Builder metadataItemName(String metadataItemName) { + this.metadataItemName = metadataItemName; + return this; + } + + /** + * Register Langchain schema in the server + */ + public Builder registerSchema(Boolean registerSchema) { + this.registerSchema = registerSchema; + return this; + } + + /** + * Create store in the server + */ + public Builder createStore(Boolean create) { + this.createStore = create; + return this; + } + + /** + * Schema file name in the server + */ + public Builder schemaFileName(String schemaFileName) { + this.schemaFileName = schemaFileName; + return this; + } + + /** + * Sets the Infinispan Hot Rod client. + * @param infinispanClient infinispan client + * @param embeddingModel the Embedding Model to be used + */ + public Builder(RemoteCacheManager infinispanClient, EmbeddingModel embeddingModel) { + super(embeddingModel); + Assert.notNull(infinispanClient, "infinispanClient must not be null"); + this.infinispanClient = infinispanClient; + this.dimension = embeddingModel.dimensions(); + } + + /** + * Builds the InfinispanVectorStore instance. + * @return a new InfinispanVectorStore instance + * @throws IllegalStateException if the builder is in an invalid state + */ + @Override + public InfinispanVectorStore build() { + return new InfinispanVectorStore(this); + } + + } + +} diff --git a/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiInfinispanItem.java b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiInfinispanItem.java new file mode 100644 index 00000000000..80bcd608166 --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiInfinispanItem.java @@ -0,0 +1,34 @@ +/* + * 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. + * 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.infinispan; + +import java.util.Map; +import java.util.Set; + +/** + * SpringAi item that is serialized for the Spring AI integration use case + * + * @param id, the id of the item + * @param text, associated text + * @param metadata, additional set of metadata + * @param embedding, the vector + * @param metadataMap, metadata as map + */ +public record SpringAiInfinispanItem(String id, String text, Set metadata, float[] embedding, + Map metadataMap) { + +} diff --git a/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiItemMarshaller.java b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiItemMarshaller.java new file mode 100644 index 00000000000..81dce9c235c --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiItemMarshaller.java @@ -0,0 +1,76 @@ +/* + * 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.infinispan; + +import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import org.infinispan.protostream.MessageMarshaller; + +/** + * Marshaller to read and write embeddings to Infinispan + */ +public class SpringAiItemMarshaller implements MessageMarshaller { + + private final String typeName; + + /** + * Constructor for the SpringAiItemMarshaller Marshaller + * @param typeName, the full type of the protobuf entity + */ + public SpringAiItemMarshaller(String typeName) { + this.typeName = typeName; + } + + @Override + public SpringAiInfinispanItem readFrom(ProtoStreamReader reader) throws IOException { + String id = reader.readString("id"); + String text = reader.readString("text"); + Set metadata = reader.readCollection("metadata", new HashSet<>(), SpringAiMetadata.class); + float[] embedding = reader.readFloats("embedding"); + + Map metadataMap = new HashMap<>(); + if (metadata != null) { + for (SpringAiMetadata meta : metadata) { + metadataMap.put(meta.name(), meta.value()); + } + } + return new SpringAiInfinispanItem(id, text, metadata, embedding, metadataMap); + } + + @Override + public void writeTo(ProtoStreamWriter writer, SpringAiInfinispanItem item) throws IOException { + writer.writeString("id", item.id()); + writer.writeString("text", item.text()); + writer.writeCollection("metadata", item.metadata(), SpringAiMetadata.class); + writer.writeFloats("embedding", item.embedding()); + } + + @Override + public Class getJavaClass() { + return SpringAiInfinispanItem.class; + } + + @Override + public String getTypeName() { + return this.typeName; + } + +} diff --git a/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiMetadata.java b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiMetadata.java new file mode 100644 index 00000000000..925626a111f --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiMetadata.java @@ -0,0 +1,26 @@ +/* + * 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.infinispan; + +/** + * Spring AI Metadata item that is serialized for the Spring AI integration use case + * + * @param name, the name of the metadata + * @param value, the value of the metadata + */ +public record SpringAiMetadata(String name, Object value) { +} diff --git a/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiMetadataMarshaller.java b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiMetadataMarshaller.java new file mode 100644 index 00000000000..040ebaf6ed9 --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/SpringAiMetadataMarshaller.java @@ -0,0 +1,115 @@ +/* + * 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.infinispan; + +import java.io.IOException; +import java.util.Date; + +import org.infinispan.protostream.MessageMarshaller; + +/** + * Marshaller to read and write metadata to Infinispan + */ +public class SpringAiMetadataMarshaller implements MessageMarshaller { + + private final String typeName; + + /** + * Constructor for the LangChainMetadata Marshaller + * @param typeName, the full type of the protobuf entity + */ + public SpringAiMetadataMarshaller(String typeName) { + this.typeName = typeName; + } + + @Override + public SpringAiMetadata readFrom(ProtoStreamReader reader) throws IOException { + String name = reader.readString("name"); + String valueStr = reader.readString("value"); + Long valueInt = reader.readLong("value_int"); + Double valueFloat = reader.readDouble("value_float"); + Boolean valueBoolean = reader.readBoolean("value_bool"); + Date valueDate = reader.readDate("value_date"); + + Object value = valueStr; + if (value == null) { + value = valueInt; + } + if (value == null) { + value = valueFloat; + } + if (value == null) { + value = valueBoolean; + } + if (value == null) { + value = valueDate; + } + + return new SpringAiMetadata(name, value); + } + + @Override + public void writeTo(ProtoStreamWriter writer, SpringAiMetadata item) throws IOException { + writer.writeString("name", item.name()); + String value = null; + Long value_int = null; + Double value_float = null; + Boolean value_boolean = null; + Date value_date = null; + if (item.value() instanceof String) { + value = (String) item.value(); + } + else if (item.value() instanceof Integer) { + value_int = ((Integer) item.value()).longValue(); + } + else if (item.value() instanceof Long) { + value_int = (Long) item.value(); + } + else if (item.value() instanceof Float) { + value_float = ((Float) item.value()).doubleValue(); + } + else if (item.value() instanceof Double) { + value_float = (Double) item.value(); + } + else if (item.value() instanceof Boolean) { + value_boolean = ((Boolean) item.value()); + } + else if (item.value() instanceof Date) { + value_date = ((Date) item.value()); + } + else { + value = item.value().toString(); + } + + writer.writeString("value", value); + writer.writeLong("value_int", value_int); + writer.writeDouble("value_float", value_float); + writer.writeBoolean("value_bool", value_boolean); + writer.writeDate("value_date", value_date); + } + + @Override + public Class getJavaClass() { + return SpringAiMetadata.class; + } + + @Override + public String getTypeName() { + return this.typeName; + } + +} diff --git a/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/package-info.java b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/package-info.java new file mode 100644 index 00000000000..ac21de12f8b --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/main/java/org/springframework/ai/vectorstore/infinispan/package-info.java @@ -0,0 +1,25 @@ +/* + * 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. + * 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. + */ + +/** + * Provides the API for embedding observations. + */ +@NonNullApi +@NonNullFields +package org.springframework.ai.vectorstore.infinispan; + +import org.springframework.lang.NonNullApi; +import org.springframework.lang.NonNullFields; diff --git a/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanFilterExpressionConverterTest.java b/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanFilterExpressionConverterTest.java new file mode 100644 index 00000000000..28a9d7c11ed --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanFilterExpressionConverterTest.java @@ -0,0 +1,210 @@ +/* + * 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. + * 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.infinispan; + +import java.util.Arrays; +import java.util.Date; +import java.util.List; + +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import org.springframework.ai.vectorstore.filter.Filter; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.GT; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.GTE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.IN; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.ISNOTNULL; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.ISNULL; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.LT; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.LTE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NIN; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NOT; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR; + +public class InfinispanFilterExpressionConverterTest { + + private final InfinispanFilterExpressionConverter converter = new InfinispanFilterExpressionConverter(); + + @Test + void shouldMapNull() { + assertThat(this.converter.convertExpression(null)).isEmpty(); + } + + @ParameterizedTest + @MethodSource("stringComparisonFilterExpression") + void shouldMapStringComparisonFilterExpression(Filter.Expression expression, String expectedQuery, + String expectedJoin) { + assertQueryAndJoin(expression, expectedQuery, expectedJoin); + } + + static List stringComparisonFilterExpression() { + return Arrays.asList( + Arguments.of(new Filter.Expression(EQ, new Filter.Key("name"), new Filter.Value("John")), + "m0.name='name' and m0.value='John'", " join i.metadata m0"), + Arguments.of(new Filter.Expression(NE, new Filter.Key("status"), new Filter.Value("active")), + "m0.name='status' and m0.value!='active' OR (i.metadata is null)", " join i.metadata m0"), + Arguments.of(new Filter.Expression(GT, new Filter.Key("name"), new Filter.Value("A")), + "m0.name='name' and m0.value>'A'", " join i.metadata m0"), + Arguments.of(new Filter.Expression(GTE, new Filter.Key("name"), new Filter.Value("A")), + "m0.name='name' and m0.value>='A'", " join i.metadata m0"), + Arguments.of(new Filter.Expression(LT, new Filter.Key("name"), new Filter.Value("Z")), + "m0.name='name' and m0.value<'Z'", " join i.metadata m0"), + Arguments.of(new Filter.Expression(LTE, new Filter.Key("name"), new Filter.Value("Z")), + "m0.name='name' and m0.value<='Z'", " join i.metadata m0")); + } + + @ParameterizedTest + @MethodSource("numericComparisonFilterExpression") + void shouldMapNumericComparisonFilterExpression(Filter.Expression expression, String expectedQuery, + String expectedJoin) { + assertQueryAndJoin(expression, expectedQuery, expectedJoin); + } + + static List numericComparisonFilterExpression() { + return Arrays.asList( + Arguments.of(new Filter.Expression(EQ, new Filter.Key("age"), new Filter.Value(25)), + "m0.name='age' and m0.value_int=25", " join i.metadata m0"), + Arguments.of(new Filter.Expression(EQ, new Filter.Key("age"), new Filter.Value(123L)), + "m0.name='age' and m0.value_int=123", " join i.metadata m0"), + Arguments.of(new Filter.Expression(EQ, new Filter.Key("score"), new Filter.Value(3.14f)), + "m0.name='score' and m0.value_float=3.140000104904175", " join i.metadata m0"), + Arguments.of(new Filter.Expression(EQ, new Filter.Key("price"), new Filter.Value(99.99d)), + "m0.name='price' and m0.value_float=99.99", " join i.metadata m0"), + Arguments.of(new Filter.Expression(GT, new Filter.Key("age"), new Filter.Value(18)), + "m0.name='age' and m0.value_int>18", " join i.metadata m0"), + Arguments.of(new Filter.Expression(LT, new Filter.Key("score"), new Filter.Value(4.5f)), + "m0.name='score' and m0.value_float<4.5", " join i.metadata m0")); + } + + @ParameterizedTest + @MethodSource("inFilters") + void shouldMapInFilterExpression(Filter.Expression expression, String expectedQuery, String expectedJoin) { + assertQueryAndJoin(expression, expectedQuery, expectedJoin); + } + + static List inFilters() { + return Arrays.asList( + Arguments.of( + new Filter.Expression(IN, new Filter.Key("category"), new Filter.Value(List.of("A", "B", "C"))), + "m0.name='category' and m0.value IN ('A', 'B', 'C')", " join i.metadata m0"), + Arguments.of(new Filter.Expression(IN, new Filter.Key("status"), new Filter.Value(List.of(1, 2, 3))), + "m0.name='status' and m0.value_int IN (1, 2, 3)", " join i.metadata m0"), + Arguments.of( + new Filter.Expression(IN, new Filter.Key("score"), new Filter.Value(List.of(1.1f, 2.2f, 3.3f))), + "m0.name='score' and m0.value_float IN (1.1, 2.2, 3.3)", " join i.metadata m0"), + Arguments.of( + new Filter.Expression(IN, new Filter.Key("score"), new Filter.Value(List.of(5.1d, 6.2d, 7.3d))), + "m0.name='score' and m0.value_float IN (5.1, 6.2, 7.3)", " join i.metadata m0")); + } + + @ParameterizedTest + @MethodSource("notInFilters") + void shouldMapNotInFilter(Filter.Expression expression, String expectedQuery, String expectedJoin) { + assertQueryAndJoin(expression, expectedQuery, expectedJoin); + } + + static List notInFilters() { + return Arrays.asList(Arguments.of( + new Filter.Expression(NIN, new Filter.Key("category"), new Filter.Value(List.of("X", "Y", "Z"))), + "(m0.value NOT IN ('X', 'Y', 'Z') and m0.name='category' ) OR (m0.value IN ('X', 'Y', 'Z') and m0.name!='category') OR (i.metadata is null)", + " join i.metadata m0"), + Arguments.of(new Filter.Expression(NIN, new Filter.Key("age"), new Filter.Value(List.of(2, 5, 6))), + "(m0.value_int NOT IN (2, 5, 6) and m0.name='age' ) OR (m0.value_int IN (2, 5, 6) and m0.name!='age') OR (i.metadata is null)", + " join i.metadata m0"), + Arguments.of( + new Filter.Expression(NIN, new Filter.Key("score"), new Filter.Value(List.of(1d, 3d, 4.4d))), + "(m0.value_float NOT IN (1.0, 3.0, 4.4) and m0.name='score' ) OR (m0.value_float IN (1.0, 3.0, 4.4) and m0.name!='score') OR (i.metadata is null)", + " join i.metadata m0")); + } + + @Test + void mapAndExpressions() { + Filter.Expression ageEq = new Filter.Expression(EQ, new Filter.Key("age"), new Filter.Value(25)); + Filter.Expression sizeNotEq = new Filter.Expression(GT, new Filter.Key("size"), new Filter.Value(170)); + Filter.Expression andExpression = new Filter.Expression(AND, ageEq, sizeNotEq); + + String filter = this.converter.convertExpression(andExpression); + String join = this.converter.doJoin(); + + assertThat(filter).isEqualTo("((m0.name='age' and m0.value_int=25) AND (m1.name='size' and m1.value_int>170))"); + assertThat(join).isEqualTo(" join i.metadata m0 join i.metadata m1"); + } + + @Test + void mapOrExpressions() { + Filter.Expression ageEq = new Filter.Expression(EQ, new Filter.Key("age"), new Filter.Value(25)); + Filter.Expression sizeNotEq = new Filter.Expression(GT, new Filter.Key("size"), new Filter.Value(170)); + Filter.Expression andExpression = new Filter.Expression(OR, ageEq, sizeNotEq); + + String filter = this.converter.convertExpression(andExpression); + String join = this.converter.doJoin(); + + assertThat(filter).isEqualTo("((m0.name='age' and m0.value_int=25) OR (m1.name='size' and m1.value_int>170))"); + assertThat(join).isEqualTo(" join i.metadata m0 join i.metadata m1"); + } + + @Test + public void shouldTransformNotExpression() { + String filter = this.converter.convertExpression( + new Filter.Expression(NOT, new Filter.Expression(EQ, new Filter.Key("age"), new Filter.Value(25)))); + assertThat(filter).isEqualTo("m0.name='age' and m0.value_int!=25 OR (i.metadata is null)"); + } + + @Test + public void testDate() { + Date date = new Date(2000); + String vectorExpr = converter + .convertExpression(new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value(date))); + Assertions.assertThat(vectorExpr) + .isEqualTo("m0.name='activationDate' and m0.value_date=" + date.toInstant().toEpochMilli()); + + vectorExpr = converter.convertExpression( + new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value("1970-01-01T00:00:02Z"))); + Assertions.assertThat(vectorExpr) + .isEqualTo("m1.name='activationDate' and m1.value_date=" + date.toInstant().toEpochMilli()); + } + + @Test + public void testNullNotNull() { + String vectorExpr = converter + .convertExpression(new Filter.Expression(ISNULL, new Filter.Key("activationDate"))); + Assertions.assertThat(vectorExpr) + .isEqualTo( + "(m0.name='activationDate' and m0.value IS NULL and m0.value_int IS NULL and m0.value_date IS NULL and m0.value_float IS NULL and m0.value_bool IS NULL) OR (m0.name NOT IN('activationDate'))"); + + vectorExpr = converter.convertExpression(new Filter.Expression(ISNOTNULL, new Filter.Key("activationDate"))); + Assertions.assertThat(vectorExpr) + .isEqualTo( + "m1.name='activationDate' and (m1.value IS NOT NULL or m1.value_int IS NOT NULL or m1.value_date IS NOT NULL or m1.value_float IS NOT NULL or m1.value_bool IS NOT NULL)"); + } + + private void assertQueryAndJoin(Filter.Expression expression, String expectedQuery, String expectedJoin) { + String filter = this.converter.convertExpression(expression); + String join = this.converter.doJoin(); + assertThat(filter).isEqualTo(expectedQuery); + assertThat(join).isEqualTo(expectedJoin); + } + +} diff --git a/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStoreIT.java b/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStoreIT.java new file mode 100644 index 00000000000..13f781adaca --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStoreIT.java @@ -0,0 +1,485 @@ +/* + * 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. + * 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.infinispan; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.function.Consumer; + +import org.awaitility.Awaitility; +import org.infinispan.client.hotrod.RemoteCache; +import org.infinispan.client.hotrod.RemoteCacheManager; +import org.infinispan.commons.util.Version; +import org.infinispan.testcontainers.InfinispanContainer; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder; +import org.springframework.boot.SpringBootConfiguration; +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 +public class InfinispanVectorStoreIT extends BaseVectorStoreTests { + + @Container + static InfinispanContainer infinispanContainer = new InfinispanContainer( + InfinispanContainer.IMAGE_BASENAME + ":" + Version.getVersion()); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class); + + 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"))); + + public static String getText(String uri) { + var resource = new DefaultResourceLoader().getResource(uri); + try { + return resource.getContentAsString(StandardCharsets.UTF_8); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @BeforeEach + void cleanDatabase() { + this.contextRunner.run(context -> context.getBean(InfinispanVectorStore.class).clear()); + } + + @Override + protected void executeTest(Consumer testFunction) { + this.contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + + @Test + public void addAndDeleteDocumentsTest() { + executeTest(vectorStore -> { + RemoteCache nativeClient = vectorStore.getNativeClient().get(); + assertThat(nativeClient.size()).isZero(); + vectorStore.add(this.documents); + assertThat(nativeClient.size()).isEqualTo(3); + vectorStore.delete(List.of("1", "2", "3")); + assertThat(nativeClient.size()).isZero(); + + }); + } + + @Test + public void addAndSearchTest() { + executeTest(vectorStore -> { + vectorStore.add(this.documents); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThresholdAll().build()), + hasSize(1)); + + List results = vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThresholdAll().build()); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(Document::getId).toList()); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThresholdAll().build()), + hasSize(0)); + + }); + + } + + @Test + public void searchWithFilersTest() { + executeTest(vectorStore -> { + 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.builder().query("The World").topK(5).similarityThresholdAll().build()), + hasSize(3)); + + List results = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("country == 'NL'") + .build()); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("country == 'BG'") + .build()); + + 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.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("country == 'BG' && year == 2020") + .build()); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("country in ['BG']") + .build()); + + 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.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("country in ['BG','NL']") + .build()); + + assertThat(results).hasSize(3); + + results = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("country not in ['BG']") + .build()); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("NOT(country not in ['BG'])") + .build()); + + 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.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("activationDate > '1970-01-01T00:00:02Z'") + .build()); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument2.getId()); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(Document::getId).toList()); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch(SearchRequest.builder().query("The World").topK(1).build()), + hasSize(0)); + + }); + + } + + @Test + public void documentUpdateTest() { + + executeTest(vectorStore -> { + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")); + + vectorStore.add(List.of(document)); + + SearchRequest springSearchRequest = SearchRequest.builder().query("Spring").topK(5).build(); + + Awaitility.await().until(() -> vectorStore.similaritySearch(springSearchRequest), hasSize(1)); + + List results = vectorStore.similaritySearch(springSearchRequest); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).containsKey("meta1"); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", + Collections.singletonMap("meta2", "meta2")); + + vectorStore.add(List.of(sameIdDocument)); + + SearchRequest fooBarSearchRequest = SearchRequest.builder().query("FooBar").topK(5).build(); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch(fooBarSearchRequest).get(0).getText(), + 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.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + + // Remove all documents from the store + vectorStore.delete(List.of(document.getId())); + Awaitility.await().until(() -> vectorStore.similaritySearch(fooBarSearchRequest), hasSize(0)); + }); + } + + @Test + public void searchThresholdTest() { + executeTest(vectorStore -> { + vectorStore.add(this.documents); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch( + SearchRequest.builder().query("Depression").topK(50).similarityThresholdAll().build()), + hasSize(3)); + + List fullResult = vectorStore + .similaritySearch(SearchRequest.builder().query("Depression").topK(5).similarityThresholdAll().build()); + + List scores = fullResult.stream().map(Document::getScore).toList(); + + assertThat(scores).hasSize(3); + + double similarityThreshold = (scores.get(0) + scores.get(1)) / 2; + + List results = vectorStore.similaritySearch(SearchRequest.builder() + .query("Depression") + .topK(5) + .similarityThreshold(similarityThreshold) + .build()); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getScore()).isGreaterThanOrEqualTo(similarityThreshold); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList()); + Awaitility.await() + .until(() -> vectorStore.similaritySearch(SearchRequest.builder().query("Hello").topK(1).build()), + hasSize(0)); + }); + } + + @Test + public void searchWithIsNullFilter() { + executeTest(vectorStore -> { + 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")); + 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.builder().query("The World").topK(5).similarityThresholdAll().build()), + hasSize(3)); + + // with text filter expression + List resultWithText = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("year IS NULL") + .build()); + + assertThat(resultWithText).hasSize(1); + assertThat(resultWithText.get(0).getId()).isEqualTo(nlDocument.getId()); + + // with filter expression builder + List resultsWithBuilder = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression(new FilterExpressionBuilder().isNull("year").build()) + .build()); + + assertThat(resultsWithBuilder).hasSize(1); + assertThat(resultsWithBuilder.get(0).getId()).isEqualTo(nlDocument.getId()); + }); + } + + @Test + public void searchWithIsNotNullFilter() { + executeTest(vectorStore -> { + 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")); + 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.builder().query("The World").topK(5).similarityThresholdAll().build()), + hasSize(3)); + + Set expectedResultSet = Set.of(bgDocument.getId(), bgDocument2.getId()); + + // with text filter expression + List resultWithText = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression("year IS NOT NULL") + .build()); + + assertThat(resultWithText).hasSize(2); + assertThat(resultWithText.get(0).getId()).isIn(expectedResultSet); + assertThat(resultWithText.get(1).getId()).isIn(expectedResultSet); + + // with filter expression builder + List resultsWithBuilder = vectorStore.similaritySearch(SearchRequest.builder() + .query("The World") + .topK(5) + .similarityThresholdAll() + .filterExpression(new FilterExpressionBuilder().isNotNull("year").build()) + .build()); + + assertThat(resultsWithBuilder).hasSize(2); + assertThat(resultsWithBuilder.get(0).getId()).isIn(expectedResultSet); + assertThat(resultsWithBuilder.get(1).getId()).isIn(expectedResultSet); + }); + } + + @Test + public void overDefaultSizeTest() { + var overDefaultSize = 12; + executeTest(vectorStore -> { + var testDocs = new ArrayList(); + for (int i = 0; i < overDefaultSize; i++) { + testDocs.add(new Document(String.valueOf(i), "Great Depression " + i, Map.of())); + } + vectorStore.add(testDocs); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThresholdAll().build()), + hasSize(1)); + + List results = vectorStore.similaritySearch(SearchRequest.builder() + .query("Great Depression") + .topK(overDefaultSize) + .similarityThresholdAll() + .build()); + + assertThat(results).hasSize(overDefaultSize); + + // Remove all documents from the store + vectorStore.delete(testDocs.stream().map(Document::getId).toList()); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch( + SearchRequest.builder().query("Great Depression").topK(1).similarityThresholdAll().build()), + hasSize(0)); + }); + } + + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + InfinispanVectorStore vectorStore = context.getBean(InfinispanVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + + @SpringBootConfiguration + public static class TestApplication { + + @Bean("vectorStore_cosine") + public InfinispanVectorStore vectorStoreDefault(EmbeddingModel embeddingModel, + RemoteCacheManager infinispanClient) { + return InfinispanVectorStore.builder(infinispanClient, embeddingModel).distance(100).build(); + } + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } + + @Bean + public RemoteCacheManager infinispanClient() { + return new RemoteCacheManager(infinispanContainer.getConnectionURI()); + } + + } + +} diff --git a/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStoreObservationIT.java b/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStoreObservationIT.java new file mode 100644 index 00000000000..3cfaf498538 --- /dev/null +++ b/vector-stores/spring-ai-infinispan-store/src/test/java/org/springframework/ai/vectorstore/infinispan/InfinispanVectorStoreObservationIT.java @@ -0,0 +1,206 @@ +/* + * 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.infinispan; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +import io.micrometer.observation.tck.TestObservationRegistry; +import io.micrometer.observation.tck.TestObservationRegistryAssert; +import org.awaitility.Awaitility; +import org.infinispan.client.hotrod.RemoteCacheManager; +import org.infinispan.commons.util.Version; +import org.infinispan.testcontainers.InfinispanContainer; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.observation.conventions.SpringAiKind; +import org.springframework.ai.observation.conventions.VectorStoreProvider; +import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; +import org.springframework.boot.SpringBootConfiguration; +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.greaterThan; + +@Testcontainers +public class InfinispanVectorStoreObservationIT { + + @Container + static InfinispanContainer infinispanContainer = new InfinispanContainer( + InfinispanContainer.IMAGE_BASENAME + ":" + Version.getVersion()); + + List documents = List.of( + new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")), + new Document(getText("classpath:/test/data/time.shelter.txt")), + new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2"))); + + public static String getText(String uri) { + var resource = new DefaultResourceLoader().getResource(uri); + try { + return resource.getContentAsString(StandardCharsets.UTF_8); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @BeforeAll + public static void beforeAll() { + Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS); + Awaitility.setDefaultPollDelay(Duration.ZERO); + Awaitility.setDefaultTimeout(Duration.ofMinutes(1)); + } + + private ApplicationContextRunner getContextRunner() { + return new ApplicationContextRunner().withUserConfiguration(Config.class); + } + + @BeforeEach + void cleanDatabase() { + getContextRunner().run(context -> { + InfinispanVectorStore store = context.getBean(InfinispanVectorStore.class); + store.clear(); + }); + } + + @Test + void observationVectorStoreAddAndQueryOperations() { + + getContextRunner().run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class); + + vectorStore.add(this.documents); + + TestObservationRegistryAssert.assertThat(observationRegistry) + .doesNotHaveAnyRemainingCurrentObservation() + .hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME) + .that() + .hasContextualNameEqualTo("%s add".formatted(VectorStoreProvider.INFINISPAN.value())) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), "add") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_SYSTEM.asString(), + VectorStoreProvider.INFINISPAN.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.SPRING_AI_KIND.asString(), + SpringAiKind.VECTOR_STORE.value()) + .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_QUERY_CONTENT.asString()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "384") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(), "observationStore") + .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_NAMESPACE.asString()) + .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_FIELD_NAME.asString()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_SEARCH_SIMILARITY_METRIC.asString(), + VectorStoreSimilarityMetric.COSINE.value()) + .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_QUERY_TOP_K.asString()) + .doesNotHaveHighCardinalityKeyValueWithKey( + HighCardinalityKeyNames.DB_VECTOR_QUERY_SIMILARITY_THRESHOLD.asString()) + + .hasBeenStarted() + .hasBeenStopped(); + + Awaitility.await() + .until(() -> vectorStore + .similaritySearch( + SearchRequest.builder().query("What is Great Depression").similarityThresholdAll().build()) + .size(), greaterThan(1)); + + observationRegistry.clear(); + + List results = vectorStore + .similaritySearch(SearchRequest.builder().query("What is Great Depression").topK(1).build()); + + assertThat(results).isNotEmpty(); + + TestObservationRegistryAssert.assertThat(observationRegistry) + .doesNotHaveAnyRemainingCurrentObservation() + .hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME) + .that() + .hasContextualNameEqualTo("%s query".formatted(VectorStoreProvider.INFINISPAN.value())) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), "query") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_SYSTEM.asString(), + VectorStoreProvider.INFINISPAN.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.SPRING_AI_KIND.asString(), + SpringAiKind.VECTOR_STORE.value()) + + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_QUERY_CONTENT.asString(), + "What is Great Depression") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "384") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(), "observationStore") + .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_NAMESPACE.asString()) + .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_FIELD_NAME.asString()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_SEARCH_SIMILARITY_METRIC.asString(), + VectorStoreSimilarityMetric.COSINE.value()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_QUERY_TOP_K.asString(), "1") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_QUERY_SIMILARITY_THRESHOLD.asString(), + "0.0") + + .hasBeenStarted() + .hasBeenStopped(); + + }); + } + + @SpringBootConfiguration + public static class Config { + + @Bean + public TestObservationRegistry observationRegistry() { + return TestObservationRegistry.create(); + } + + @Bean("vectorStore_cosine") + public InfinispanVectorStore vectorStoreDefault(EmbeddingModel embeddingModel, + RemoteCacheManager infinispanClient, TestObservationRegistry observationRegistry) { + return InfinispanVectorStore.builder(infinispanClient, embeddingModel) + .distance(100) + .observationRegistry(observationRegistry) + .customObservationConvention(null) + .storeName("observationStore") + .build(); + } + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } + + @Bean + public RemoteCacheManager infinispanClient() { + return new RemoteCacheManager(infinispanContainer.getConnectionURI()); + } + + } + +}