diff --git a/pom.xml b/pom.xml index abe272dc4..84902529c 100644 --- a/pom.xml +++ b/pom.xml @@ -30,6 +30,7 @@ vector-stores/spring-ai-cassandra-store vector-stores/spring-ai-chroma-store vector-stores/spring-ai-elasticsearch-store + vector-stores/spring-ai-gemfire-store vector-stores/spring-ai-hanadb-store vector-stores/spring-ai-milvus-store @@ -40,8 +41,9 @@ vector-stores/spring-ai-pinecone-store vector-stores/spring-ai-qdrant-store vector-stores/spring-ai-redis-store - vector-stores/spring-ai-weaviate-store + vector-stores/spring-ai-typesense-store + vector-stores/spring-ai-weaviate-store spring-ai-spring-boot-starters/spring-ai-starter-azure-store spring-ai-spring-boot-starters/spring-ai-starter-cassandra-store spring-ai-spring-boot-starters/spring-ai-starter-chroma-store @@ -55,8 +57,8 @@ spring-ai-spring-boot-starters/spring-ai-starter-pinecone-store spring-ai-spring-boot-starters/spring-ai-starter-qdrant-store spring-ai-spring-boot-starters/spring-ai-starter-redis-store + spring-ai-spring-boot-starters/spring-ai-starter-typesense-store spring-ai-spring-boot-starters/spring-ai-starter-weaviate-store - models/spring-ai-anthropic models/spring-ai-azure-openai models/spring-ai-bedrock @@ -72,7 +74,6 @@ models/spring-ai-vertex-ai-palm2 models/spring-ai-watsonx-ai models/spring-ai-zhipuai - spring-ai-spring-boot-starters/spring-ai-starter-anthropic spring-ai-spring-boot-starters/spring-ai-starter-azure-openai spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai @@ -88,7 +89,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-vertex-ai-palm2 spring-ai-spring-boot-starters/spring-ai-starter-watsonx-ai spring-ai-spring-boot-starters/spring-ai-starter-zhipuai - + VMware Inc. @@ -167,6 +168,7 @@ 11.6.1 4.5.1 1.7.1 + 0.5.0 0.0.4 diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index 19b569bc7..07249e06d 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -332,6 +332,12 @@ ${project.version} + + org.springframework.ai + spring-ai-typesense-store + ${project.version} + + org.springframework.ai spring-ai-pinecone-store-spring-boot-starter @@ -397,12 +403,19 @@ spring-ai-mongodb-atlas-store-spring-boot-starter ${project.version} + org.springframework.ai spring-ai-anthropic-spring-boot-starter ${project.version} + + org.springframework.ai + spring-ai-typesense-store-spring-boot-starter + ${project.version} + + org.springframework.ai spring-ai-spring-boot-testcontainers diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc index 6e625eaa3..52ea50314 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc @@ -68,6 +68,7 @@ *** xref:api/vectordbs/qdrant.adoc[] *** xref:api/vectordbs/redis.adoc[] *** xref:api/vectordbs/hana.adoc[SAP Hana] +*** xref:api/vectordbs/typesense.adoc[] *** xref:api/vectordbs/weaviate.adoc[] diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc index 9c9bfb8f3..f98d70b05 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc @@ -109,6 +109,7 @@ These are the available implementations of the `VectorStore` interface: * xref:api/vectordbs/qdrant.adoc[Qdrant Vector Store] - https://www.qdrant.tech/[Qdrant] vector store. * xref:api/vectordbs/redis.adoc[Redis Vector Store] - The https://redis.io/[Redis] vector store. * xref:api/vectordbs/hana.adoc[SAP Hana Vector Store] - The https://news.sap.com/2024/04/sap-hana-cloud-vector-engine-ai-with-business-context/[SAP HANA] vector store. +* xref:api/vectordbs/typesense.adoc[Typesense Vector Store] - The https://typesense.org/docs/0.24.0/api/vector-search.html[Typesense] vector store. * xref:api/vectordbs/weaviate.adoc[Weaviate Vector Store] - The https://weaviate.io/[Weaviate] vector store. * link:https://github.com/spring-projects/spring-ai/blob/main/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java[SimpleVectorStore] - A simple implementation of persistent vector storage, good for educational purposes. diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/typesense.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/typesense.adoc new file mode 100644 index 000000000..e48aaa975 --- /dev/null +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/typesense.adoc @@ -0,0 +1,242 @@ += Typesense + +This section walks you through setting up `TypesenseVectorStore` to store document embeddings and perform similarity searches. + +link:https://typesense.org[Typesense] Typesense is an open source, typo tolerant search engine that is optimized for instant sub-50ms searches, while providing an intuitive developer experience. + +== Prerequisites + +1. A Typesense instance +- link:https://typesense.org/docs/guide/install-typesense.html[Typesense Cloud] (recommended) +- link:https://hub.docker.com/r/typesense/typesense/[Docker] image _typesense/typesense:latest_ + +2. `EmbeddingClient` instance to compute the document embeddings. Several options are available: +- If required, an API key for the xref:api/embeddings.adoc#available-implementations[EmbeddingClient] to generate the embeddings stored by the `TypesenseVectorStore`. + +== Auto-configuration + +Spring AI provides Spring Boot auto-configuration for the Typesense Vector Sore. +To enable it, add the following dependency to your project's Maven `pom.xml` file: + +[source, xml] +---- + + org.springframework.ai + spring-ai-typesense-spring-boot-starter + +---- + +or to your Gradle `build.gradle` build file. + +[source,groovy] +---- +dependencies { + implementation 'org.springframework.ai:spring-ai-typesense-spring-boot-starter' +} +---- + +TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file. + +TIP: Refer to the xref:getting-started.adoc#repositories[Repositories] section to add Milestone and/or Snapshot Repositories to your build file. + +Additionally, you will need a configured `EmbeddingClient` bean. Refer to the xref:api/embeddings.adoc#available-implementations[EmbeddingClient] section for more information. + +Here is an example of the needed bean: + +[source,java] +---- +@Bean +public EmbeddingClient embeddingClient() { + // Can be any other EmbeddingClient implementation. + return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("SPRING_AI_OPENAI_API_KEY"))); +} +---- + +To connect to Typesense you need to provide access details for your instance. +A simple configuration can either be provided via Spring Boot's _application.yml_, + +[source,yaml] +---- +spring: + ai: + vectorstore: + typesense: + collectionName: "vector_store" + embeddingDimension: 1536 + client: + protocl: http + host: localhost + port: 8108 + apiKey: xyz +---- + +Please have a look at the list of xref:#_configuration_properties[configuration parameters] for the vector store to learn about the default values and configuration options. + +Now you can Auto-wire the Typesense Vector Store in your application and use it + +[source,java] +---- +@Autowired VectorStore vectorStore; + +// ... + +List documents = List.of( + new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")), + new Document("The World is Big and Salvation Lurks Around the Corner"), + new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2"))); + +// Add the documents to Typesense +vectorStore.add(documents); + +// Retrieve documents similar to a query +List results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5)); +---- + +=== Configuration properties + +You can use the following properties in your Spring Boot configuration to customize the Typesense vector store. + +|=== +|Property| Description | Default value + +|`spring.ai.vectorstore.typesense.client.protocol`| HTTP Protocol | `http` +|`spring.ai.vectorstore.typesense.client.host`| Hostname | `localhost` +|`spring.ai.vectorstore.typesense.client.port`| Port | `8108` +|`spring.ai.vectorstore.typesense.client.apiKey`| ApiKey | `xyz` +|`spring.ai.vectorstore.typesense.collectionName`| Collection Name | `vector_store` +|`spring.ai.vectorstore.typesense.embeddingDimension`| Embedding Dimension | `1536` + +|=== + +== Metadata filtering + +You can leverage the generic, portable link:https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_metadata_filters[metadata filters] with `TypesenseVectorStore` as well. + +For example, you can use either the text expression language: + +[source,java] +---- +vectorStore.similaritySearch( + SearchRequest + .query("The World") + .withTopK(TOP_K) + .withSimilarityThreshold(SIMILARITY_THRESHOLD) + .withFilterExpression("country in ['UK', 'NL'] && year >= 2020")); +---- + +or programmatically using the expression DSL: + +[source,java] +---- +FilterExpressionBuilder b = new FilterExpressionBuilder(); + +vectorStore.similaritySearch( + SearchRequest + .query("The World") + .withTopK(TOP_K) + .withSimilarityThreshold(SIMILARITY_THRESHOLD) + .withFilterExpression(b.and( + b.in("country", "UK", "NL"), + b.gte("year", 2020)).build())); +---- + +The portable filter expressions get automatically converted into link:https://typesense.org/docs/0.24.0/api/search.html#filter-parameters[Typesense Search Filters]. +For example, the following portable filter expression: + +[source,sql] +---- +country in ['UK', 'NL'] && year >= 2020 +---- + +is converted into Typesense filter: + +[source] +---- +country: ['UK', 'NL'] && year: >=2020 +---- + +== Manual configuration + +If you prefer not to use the auto-configuration, you can manually configure the Typesense Vector Store. +Add the Typesense Vector Store and Jedis dependencies + +[source,xml] +---- + + org.springframework.ai + spring-ai-typesense + +---- + +TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file. + +Then, create a `TypesenseVectorStore` bean in your Spring configuration: + +[source,java] +---- +@Bean +public VectorStore vectorStore(Client client, EmbeddingClient embeddingClient) { + + TypesenseVectorStoreConfig config = TypesenseVectorStoreConfig.builder() + .withCollectionName("test_vector_store") + .withEmbeddingDimension(embeddingClient.dimensions()) + .build(); + + return new TypesenseVectorStore(client, embeddingClient, config); +} + +@Bean +public Client typesenseClient() { + List nodes = new ArrayList<>(); + nodes + .add(new Node("http", typesenseContainer.getHost(), typesenseContainer.getMappedPort(8108).toString())); + + Configuration configuration = new Configuration(nodes, Duration.ofSeconds(5), "xyz"); + return new Client(configuration); +} +---- + +[NOTE] +==== +It is more convenient and preferred to create the `TypesenseVectorStore` as a Bean. +But if you decide to create it manually, then you must call the `TypesenseVectorStore#afterPropertiesSet()` after setting the properties and before using the client. +==== + + +Then in your main code, create some documents: + +[source,java] +---- +List documents = List.of( + new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("country", "UK", "year", 2020)), + new Document("The World is Big and Salvation Lurks Around the Corner", Map.of()), + new Document("You walk forward facing the past and you turn back toward the future.", Map.of("country", "NL", "year", 2023))); +---- + +Now add the documents to your vector store: + + +[source,java] +---- +vectorStore.add(documents); +---- + +And finally, retrieve documents similar to a query: + +[source,java] +---- +List results = vectorStore.similaritySearch( + SearchRequest + .query("Spring") + .withTopK(5)); +---- + +If all goes well, you should retrieve the document containing the text "Spring AI rocks!!". + +[NOTE] +==== +If you are not retrieveing the documents in the expected order or the search results are not as expected, check the embedding model you are using. + +Embedding models can have a significant impact on the search results (i.e. make sure if your data is in Spanish to use a Spanish or multilingual embedding model). +==== + diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index 9c8f275ae..aaba27c8d 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -296,6 +296,7 @@ true + org.springframework.ai spring-ai-minimax @@ -310,6 +311,14 @@ true + + + org.springframework.ai + spring-ai-typesense-store + ${project.parent.version} + true + + diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseConnectionDetails.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseConnectionDetails.java new file mode 100644 index 000000000..7c342996a --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseConnectionDetails.java @@ -0,0 +1,16 @@ +package org.springframework.ai.autoconfigure.vectorstore.typesense; + +import org.springframework.boot.autoconfigure.service.connection.ConnectionDetails; + +/** + * @author Pablo Sanchidrian Herrera + */ +public interface TypesenseConnectionDetails extends ConnectionDetails { + + String getHost(); + + String getProtocol(); + + String getPort(); + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseServiceClientProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseServiceClientProperties.java new file mode 100644 index 000000000..2ede36c66 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseServiceClientProperties.java @@ -0,0 +1,57 @@ +package org.springframework.ai.autoconfigure.vectorstore.typesense; + +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * @author Pablo Sanchidrian Herrera + */ +@ConfigurationProperties(TypesenseServiceClientProperties.CONFIG_PREFIX) +public class TypesenseServiceClientProperties { + + public static final String CONFIG_PREFIX = "spring.ai.vectorstore.typesense.client"; + + private String protocol = "http"; + + private String host = "localhost"; + + private String port = "8108"; + + /** + * Typesense API key. This is the default api key when the user follows the Typesense + * quick start guide. + */ + private String apiKey = "xyz"; + + public String getProtocol() { + return protocol; + } + + public void setProtocol(String protocol) { + this.protocol = protocol; + } + + public String getHost() { + return host; + } + + public void setHost(String host) { + this.host = host; + } + + public String getPort() { + return port; + } + + public void setPort(String port) { + this.port = port; + } + + public String getApiKey() { + return apiKey; + } + + public void setApiKey(String apiKey) { + this.apiKey = apiKey; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreAutoConfiguration.java new file mode 100644 index 000000000..f16a04df8 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreAutoConfiguration.java @@ -0,0 +1,84 @@ +package org.springframework.ai.autoconfigure.vectorstore.typesense; + +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.TypesenseVectorStore; +import org.springframework.ai.vectorstore.TypesenseVectorStore.TypesenseVectorStoreConfig; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; +import org.typesense.api.Client; +import org.typesense.api.Configuration; +import org.typesense.resources.Node; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; + +/** + * @author Pablo Sanchidrian Herrera + */ +@AutoConfiguration +@ConditionalOnClass({ TypesenseVectorStore.class, EmbeddingModel.class }) +@EnableConfigurationProperties({ TypesenseServiceClientProperties.class, TypesenseVectorStoreProperties.class }) +public class TypesenseVectorStoreAutoConfiguration { + + @Bean + @ConditionalOnMissingBean(TypesenseConnectionDetails.class) + TypesenseVectorStoreAutoConfiguration.PropertiesTypesenseConnectionDetails typesenseServiceClientConnectionDetails( + TypesenseServiceClientProperties properties) { + return new TypesenseVectorStoreAutoConfiguration.PropertiesTypesenseConnectionDetails(properties); + } + + @Bean + @ConditionalOnMissingBean + public VectorStore vectorStore(Client typesenseClient, EmbeddingModel embeddingClient, + TypesenseVectorStoreProperties properties) { + + TypesenseVectorStoreConfig config = TypesenseVectorStoreConfig.builder() + .withCollectionName(properties.getCollectionName()) + .withEmbeddingDimension(properties.getEmbeddingDimension()) + .build(); + + return new TypesenseVectorStore(typesenseClient, embeddingClient, config); + } + + @Bean + @ConditionalOnMissingBean + public Client typesenseClient(TypesenseServiceClientProperties clientProperties, + TypesenseConnectionDetails connectionDetails) { + List nodes = new ArrayList<>(); + nodes.add(new Node(clientProperties.getProtocol(), clientProperties.getHost(), clientProperties.getPort())); + + Configuration configuration = new Configuration(nodes, Duration.ofSeconds(5), clientProperties.getApiKey()); + return new Client(configuration); + } + + private static class PropertiesTypesenseConnectionDetails implements TypesenseConnectionDetails { + + private final TypesenseServiceClientProperties properties; + + PropertiesTypesenseConnectionDetails(TypesenseServiceClientProperties properties) { + this.properties = properties; + } + + @Override + public String getProtocol() { + return this.properties.getProtocol(); + } + + @Override + public String getHost() { + return this.properties.getHost(); + } + + @Override + public String getPort() { + return this.properties.getPort(); + } + + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreProperties.java new file mode 100644 index 000000000..006c071c4 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreProperties.java @@ -0,0 +1,40 @@ +package org.springframework.ai.autoconfigure.vectorstore.typesense; + +import org.springframework.ai.vectorstore.TypesenseVectorStore; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * @author Pablo Sanchidrian Herrera + */ +@ConfigurationProperties(TypesenseVectorStoreProperties.CONFIG_PREFIX) +public class TypesenseVectorStoreProperties { + + public static final String CONFIG_PREFIX = "spring.ai.vectorstore.typesense"; + + /** + * Typesense collection name to store the vectors. + */ + private String collectionName = TypesenseVectorStore.DEFAULT_COLLECTION_NAME; + + /** + * The dimension of the vectors to be stored in the Typesense collection. + */ + private int embeddingDimension = TypesenseVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE; + + public String getCollectionName() { + return collectionName; + } + + public void setCollectionName(String collectionName) { + this.collectionName = collectionName; + } + + public int getEmbeddingDimension() { + return embeddingDimension; + } + + public void setEmbeddingDimension(int embeddingDimension) { + this.embeddingDimension = embeddingDimension; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index d39b769ec..0666cd0e6 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -35,3 +35,4 @@ org.springframework.ai.autoconfigure.vectorstore.elasticsearch.ElasticsearchVect org.springframework.ai.autoconfigure.vectorstore.cassandra.CassandraVectorStoreAutoConfiguration org.springframework.ai.autoconfigure.zhipuai.ZhiPuAiAutoConfiguration org.springframework.ai.autoconfigure.chat.client.ChatClientAutoConfiguration +org.springframework.ai.autoconfigure.vectorstore.typesense.TypesenseVectorStoreAutoConfiguration diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreAutoConfigurationIT.java new file mode 100644 index 000000000..b6c937f96 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/typesense/TypesenseVectorStoreAutoConfigurationIT.java @@ -0,0 +1,107 @@ +package org.springframework.ai.autoconfigure.vectorstore.typesense; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.springframework.ai.ResourceUtils; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +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.util.FileSystemUtils; +import org.testcontainers.containers.BindMode; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.junit.jupiter.Testcontainers; + +import java.io.File; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Pablo Sanchidrian Herrera + */ +@Testcontainers +public class TypesenseVectorStoreAutoConfigurationIT { + + private static GenericContainer typesenseContainer; + + private static final File TEMP_FOLDER = new File("target/test-" + UUID.randomUUID().toString()); + + 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"))); + + @BeforeAll + public static void beforeAll() { + FileSystemUtils.deleteRecursively(TEMP_FOLDER); + TEMP_FOLDER.mkdirs(); + + typesenseContainer = new GenericContainer<>("typesense/typesense:26.0").withExposedPorts(8108) + .withCommand("--data-dir", "/data", "--api-key=xyz", "--enable-cors") + .withFileSystemBind(TEMP_FOLDER.getAbsolutePath(), "/data", BindMode.READ_WRITE) + .withStartupTimeout(Duration.ofSeconds(100)); + + typesenseContainer.start(); + } + + @AfterAll + public static void afterAll() { + typesenseContainer.stop(); + FileSystemUtils.deleteRecursively(TEMP_FOLDER); + } + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(TypesenseVectorStoreAutoConfiguration.class)) + .withUserConfiguration(Config.class); + + @Test + public void addAndSearch() { + contextRunner + .withPropertyValues("spring.ai.vectorstore.typesense.embeddingDimension=384", + "spring.ai.vectorstore.typesense.collectionName=myTestCollection", + "spring.ai.vectorstore.typesense.client.apiKey=xyz", + "spring.ai.vectorstore.typesense.client.protocol=http", + "spring.ai.vectorstore.typesense.client.host=" + typesenseContainer.getHost(), + "spring.ai.vectorstore.typesense.client.port=" + typesenseContainer.getMappedPort(8108).toString()) + .run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + vectorStore.add(documents); + + List results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId()); + assertThat(resultDoc.getContent()).contains( + "Spring AI provides abstractions that serve as the foundation for developing AI applications."); + assertThat(resultDoc.getMetadata()).hasSize(2); + assertThat(resultDoc.getMetadata()).containsKeys("spring", "distance"); + + vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList()); + + results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1)); + assertThat(results).hasSize(0); + }); + } + + @Configuration(proxyBeanMethods = false) + static class Config { + + @Bean + public EmbeddingModel embeddingClient() { + return new TransformersEmbeddingModel(); + } + + } + +} diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-typesense-store/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-typesense-store/pom.xml new file mode 100644 index 000000000..de5170aa8 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-typesense-store/pom.xml @@ -0,0 +1,44 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-typesense-store-spring-boot-starter + jar + Spring AI Starter - Typesense + Spring AI Typesense Auto Configuration + 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.springframework.boot + spring-boot-starter + + + + org.springframework.ai + spring-ai-spring-boot-autoconfigure + ${project.parent.version} + + + + org.springframework.ai + spring-ai-typesense-store + ${project.parent.version} + + + + \ No newline at end of file diff --git a/vector-stores/spring-ai-typesense-store/pom.xml b/vector-stores/spring-ai-typesense-store/pom.xml new file mode 100644 index 000000000..df2484c19 --- /dev/null +++ b/vector-stores/spring-ai-typesense-store/pom.xml @@ -0,0 +1,67 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + + spring-ai-typesense-store + jar + Spring AI Typesense Vector Store + Spring AI Typesense 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.springframework.ai + spring-ai-core + ${parent.version} + + + + org.typesense + typesense-java + ${typesense.version} + + + + + org.springframework.ai + spring-ai-test + ${parent.version} + test + + + + org.springframework.ai + spring-ai-transformers + ${parent.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.testcontainers + junit-jupiter + test + + + + + \ No newline at end of file diff --git a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseFilterExpressionConverter.java b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseFilterExpressionConverter.java new file mode 100644 index 000000000..0f19340d8 --- /dev/null +++ b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseFilterExpressionConverter.java @@ -0,0 +1,60 @@ +package org.springframework.ai.vectorstore; + +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter; + +/** + * Converts {@link Filter.Expression} into Typesense metadata filter expression format. + * (https://typesense.org/docs/0.24.0/api/search.html#filter-parameters) + * + * @author Pablo Sanchidrian + */ +public class TypesenseFilterExpressionConverter extends AbstractFilterExpressionConverter { + + @Override + protected void doExpression(Filter.Expression exp, StringBuilder context) { + this.convertOperand(exp.left(), context); + context.append(getOperationSymbol(exp)); + this.convertOperand(exp.right(), context); + } + + private String getOperationSymbol(Filter.Expression exp) { + switch (exp.type()) { + case AND: + return " && "; + case OR: + return " || "; + case EQ: + return " "; // in typesense "EQ" operator looks like -> country:USA + case NE: + return " != "; + case LT: + return " < "; + case LTE: + return " <= "; + case GT: + return " > "; + case GTE: + return " >= "; + case IN: + return " "; // in typesense "IN" operator looks like -> country: [USA, UK] + case NIN: + return " != "; // in typesense "NIN" operator looks like -> country: + // !=[USA, UK] + default: + throw new RuntimeException("Not supported expression type:" + exp.type()); + } + } + + @Override + protected void doGroup(Filter.Group group, StringBuilder context) { + this.convertOperand(new Filter.Expression(Filter.ExpressionType.AND, group.content(), group.content()), + context); // trick + } + + @Override + protected void doKey(Filter.Key key, StringBuilder context) { + context.append("metadata." + key.key() + ":"); + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java new file mode 100644 index 000000000..45dd233eb --- /dev/null +++ b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java @@ -0,0 +1,337 @@ +package org.springframework.ai.vectorstore; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.util.Assert; +import org.typesense.api.Client; +import org.typesense.api.FieldTypes; +import org.typesense.model.CollectionResponse; +import org.typesense.model.CollectionSchema; +import org.typesense.model.DeleteDocumentsParameters; +import org.typesense.model.Field; +import org.typesense.model.ImportDocumentsParameters; +import org.typesense.model.MultiSearchCollectionParameters; +import org.typesense.model.MultiSearchResult; +import org.typesense.model.MultiSearchSearchesParameter; + +/** + * @author Pablo Sanchidrian Herrera + */ +public class TypesenseVectorStore implements VectorStore, InitializingBean { + + private static final Logger logger = LoggerFactory.getLogger(TypesenseVectorStore.class); + + /** + * The name of the field that contains the document ID. It is mandatory to set "id" as + * the field name because that is the name that typesense is going to look for. + */ + public static final String DOC_ID_FIELD_NAME = "id"; + + public static final String CONTENT_FIELD_NAME = "content"; + + public static final String METADATA_FIELD_NAME = "metadata"; + + public static final String EMBEDDING_FIELD_NAME = "embedding"; + + public static final int OPENAI_EMBEDDING_DIMENSION_SIZE = 1536; + + public static final String DEFAULT_COLLECTION_NAME = "vector_store"; + + public static final int INVALID_EMBEDDING_DIMENSION = -1; + + private final Client client; + + private final EmbeddingModel embeddingClient; + + private final TypesenseVectorStoreConfig config; + + public final FilterExpressionConverter filterExpressionConverter = new TypesenseFilterExpressionConverter(); + + public static class TypesenseVectorStoreConfig { + + private final String collectionName; + + private final int embeddingDimension; + + public TypesenseVectorStoreConfig(String collectionName, int embeddingDimension) { + this.collectionName = collectionName; + this.embeddingDimension = embeddingDimension; + } + + /** + * {@return the default config} + */ + public static TypesenseVectorStoreConfig defaultConfig() { + return builder().build(); + } + + private TypesenseVectorStoreConfig(Builder builder) { + this.collectionName = builder.collectionName; + this.embeddingDimension = builder.embeddingDimension; + } + + /** + * Start building a new configuration. + * @return The entry point for creating a new configuration. + */ + public static Builder builder() { + + return new Builder(); + } + + public static class Builder { + + private String collectionName; + + private int embeddingDimension; + + /** + * Set the collection name. + * @param collectionName The collection name. + * @return The builder. + */ + public Builder withCollectionName(String collectionName) { + this.collectionName = collectionName; + return this; + } + + /** + * Set the embedding dimension. + * @param embeddingDimension The embedding dimension. + * @return The builder. + */ + public Builder withEmbeddingDimension(int embeddingDimension) { + this.embeddingDimension = embeddingDimension; + return this; + } + + /** + * Build the configuration. + * @return The configuration. + */ + public TypesenseVectorStoreConfig build() { + return new TypesenseVectorStoreConfig(this); + } + + } + + } + + public TypesenseVectorStore(Client client, EmbeddingModel embeddingClient) { + this(client, embeddingClient, TypesenseVectorStoreConfig.defaultConfig()); + } + + public TypesenseVectorStore(Client client, EmbeddingModel embeddingClient, TypesenseVectorStoreConfig config) { + Assert.notNull(client, "Typesense must not be null"); + Assert.notNull(embeddingClient, "EmbeddingClient must not be null"); + + this.client = client; + this.embeddingClient = embeddingClient; + this.config = config; + } + + @Override + public void add(List documents) { + Assert.notNull(documents, "Documents must not be null"); + + List> documentList = documents.stream().map(document -> { + HashMap typesenseDoc = new HashMap<>(); + typesenseDoc.put(DOC_ID_FIELD_NAME, document.getId()); + typesenseDoc.put(CONTENT_FIELD_NAME, document.getContent()); + typesenseDoc.put(METADATA_FIELD_NAME, document.getMetadata()); + List embedding = this.embeddingClient.embed(document.getContent()); + typesenseDoc.put(EMBEDDING_FIELD_NAME, embedding); + + return typesenseDoc; + }).toList(); + + ImportDocumentsParameters importDocumentsParameters = new ImportDocumentsParameters(); + importDocumentsParameters.action("upsert"); + + try { + this.client.collections(this.config.collectionName) + .documents() + .import_(documentList, importDocumentsParameters); + + logger.info("Added {} documents", documentList.size()); + } + catch (Exception e) { + logger.error("Failed to add documents", e); + } + } + + @Override + public Optional delete(List idList) { + DeleteDocumentsParameters deleteDocumentsParameters = new DeleteDocumentsParameters(); + deleteDocumentsParameters.filterBy(DOC_ID_FIELD_NAME + ":=[" + String.join(",", idList) + "]"); + + try { + int deletedDocs = (Integer) this.client.collections(this.config.collectionName) + .documents() + .delete(deleteDocumentsParameters) + .getOrDefault("num_deleted", 0); + + if (deletedDocs < idList.size()) { + logger.warn("Failed to delete all documents"); + } + + return Optional.of(deletedDocs > 0); + } + catch (Exception e) { + logger.error("Failed to delete documents", e); + return Optional.of(Boolean.FALSE); + } + } + + @Override + public List similaritySearch(SearchRequest request) { + Assert.notNull(request.getQuery(), "Query string must not be null"); + + String nativeFilterExpressions = (request.getFilterExpression() != null) + ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : ""; + + logger.info("Filter expression: {}", nativeFilterExpressions); + + List embedding = this.embeddingClient.embed(request.getQuery()); + + MultiSearchCollectionParameters multiSearchCollectionParameters = new MultiSearchCollectionParameters(); + multiSearchCollectionParameters.collection(this.config.collectionName); + multiSearchCollectionParameters.q("*"); + + // typesnese uses only cosine similarity + String vectorQuery = EMBEDDING_FIELD_NAME + ":(" + "[" + + String.join(",", embedding.stream().map(String::valueOf).toList()) + "], " + "k: " + request.getTopK() + + ", " + "distance_threshold: " + (1 - request.getSimilarityThreshold()) + ")"; + + multiSearchCollectionParameters.vectorQuery(vectorQuery); + multiSearchCollectionParameters.filterBy(nativeFilterExpressions); + + MultiSearchSearchesParameter multiSearchesParameter = new MultiSearchSearchesParameter() + .addSearchesItem(multiSearchCollectionParameters); + + try { + MultiSearchResult result = this.client.multiSearch.perform(multiSearchesParameter, + Map.of("query_by", EMBEDDING_FIELD_NAME)); + + List documents = result.getResults() + .stream() + .flatMap(searchResult -> searchResult.getHits().stream().map(hit -> { + Map rawDocument = hit.getDocument(); + String docId = rawDocument.get(DOC_ID_FIELD_NAME).toString(); + String content = rawDocument.get(CONTENT_FIELD_NAME).toString(); + Map metadata = rawDocument.get(METADATA_FIELD_NAME) instanceof Map + ? (Map) rawDocument.get(METADATA_FIELD_NAME) : Map.of(); + metadata.put("distance", hit.getVectorDistance()); + return new Document(docId, content, metadata); + })) + .toList(); + + logger.info("Found {} documents", documents.size()); + return documents; + } + catch (Exception e) { + logger.error("Failed to search documents", e); + return List.of(); + } + } + + int embeddingDimensions() { + if (this.config.embeddingDimension != INVALID_EMBEDDING_DIMENSION) { + return this.config.embeddingDimension; + } + try { + int embeddingDimensions = this.embeddingClient.dimensions(); + if (embeddingDimensions > 0) { + return embeddingDimensions; + } + } + catch (Exception e) { + logger.warn("Failed to obtain the embedding dimensions from the embedding client and fall backs to default:" + + this.config.embeddingDimension, e); + } + return OPENAI_EMBEDDING_DIMENSION_SIZE; + } + + // --------------------------------------------------------------------------------- + // Initialization + // --------------------------------------------------------------------------------- + @Override + public void afterPropertiesSet() throws Exception { + this.createCollection(); + } + + private boolean hasCollection() { + try { + this.client.collections(this.config.collectionName).retrieve(); + return true; + } + catch (Exception e) { + return false; + } + } + + void createCollection() { + if (this.hasCollection()) { + logger.info("Collection {} already exists", this.config.collectionName); + return; + } + + CollectionSchema collectionSchema = new CollectionSchema(); + + collectionSchema.name(this.config.collectionName) + .addFieldsItem(new Field().name(DOC_ID_FIELD_NAME).type(FieldTypes.STRING).optional(false)) + .addFieldsItem(new Field().name(CONTENT_FIELD_NAME).type(FieldTypes.STRING).optional(false)) + .addFieldsItem(new Field().name(METADATA_FIELD_NAME).type(FieldTypes.OBJECT).optional(true)) + .addFieldsItem(new Field().name(EMBEDDING_FIELD_NAME) + .type(FieldTypes.FLOAT_ARRAY) + .numDim(this.embeddingDimensions()) + .optional(false)) + .enableNestedFields(true); + + try { + this.client.collections().create(collectionSchema); + logger.info("Collection {} created", this.config.collectionName); + } + catch (Exception e) { + logger.error("Failed to create collection {}", this.config.collectionName, e); + } + } + + void dropCollection() { + if (!this.hasCollection()) { + logger.info("Collection {} does not exist", this.config.collectionName); + return; + } + + try { + this.client.collections(this.config.collectionName).delete(); + logger.info("Collection {} dropped", this.config.collectionName); + } + catch (Exception e) { + logger.error("Failed to drop collection {}", this.config.collectionName, e); + } + } + + Map getCollectionInfo() { + try { + CollectionResponse retrievedCollection = this.client.collections(this.config.collectionName).retrieve(); + return Map.of("name", retrievedCollection.getName(), "num_documents", + retrievedCollection.getNumDocuments()); + } + catch (Exception e) { + logger.error("Failed to retrieve collection info", e); + return null; + } + + } + +} diff --git a/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/TypesenseVectorStoreIT.java b/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/TypesenseVectorStoreIT.java new file mode 100644 index 000000000..952458ea6 --- /dev/null +++ b/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/TypesenseVectorStoreIT.java @@ -0,0 +1,282 @@ +package org.springframework.ai.vectorstore; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Test; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.core.io.DefaultResourceLoader; +import org.testcontainers.containers.BindMode; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.typesense.api.Client; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.util.*; + +import org.springframework.ai.vectorstore.TypesenseVectorStore.TypesenseVectorStoreConfig; +import org.typesense.api.Configuration; +import org.typesense.resources.Node; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Pablo Sanchidrian Herrera + */ +@Testcontainers +public class TypesenseVectorStoreIT { + + private static Path tempDirectory; + + static { + try { + tempDirectory = Files.createTempDirectory("typesense-test"); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Container + private static GenericContainer typesenseContainer = new GenericContainer<>("typesense/typesense:26.0") + .withExposedPorts(8108) + .withCommand("--data-dir", "/data", "--api-key=xyz", "--enable-cors") + .withFileSystemBind(tempDirectory.toString(), "/data", BindMode.READ_WRITE); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class); + + 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); + } + } + + private void resetCollection(VectorStore vectorStore) { + ((TypesenseVectorStore) vectorStore).dropCollection(); + ((TypesenseVectorStore) vectorStore).createCollection(); + } + + @Test + void documentUpdate() { + contextRunner.run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + resetCollection(vectorStore); + + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")); + + vectorStore.add(List.of(document)); + + Map info = ((TypesenseVectorStore) vectorStore).getCollectionInfo(); + assertThat(info.get("num_documents")).isEqualTo(1L); + + List results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).containsKey("meta1"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", + Collections.singletonMap("meta2", "meta2")); + + vectorStore.add(List.of(sameIdDocument)); + + info = ((TypesenseVectorStore) vectorStore).getCollectionInfo(); + assertThat(info.get("num_documents")).isEqualTo(1L); + + results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5)); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + vectorStore.delete(List.of(document.getId())); + + info = ((TypesenseVectorStore) vectorStore).getCollectionInfo(); + assertThat(info.get("num_documents")).isEqualTo(0L); + + }); + } + + @Test + void addAndSearch() { + + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + resetCollection(vectorStore); + + vectorStore.add(documents); + + Map info = ((TypesenseVectorStore) vectorStore).getCollectionInfo(); + + assertThat(info.get("num_documents")).isEqualTo(3L); + + List results = vectorStore.similaritySearch(SearchRequest.query("Spring")); + + assertThat(results).hasSize(3); + }); + } + + @Test + void searchWithFilters() { + + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + resetCollection(vectorStore); + + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2020)); + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "NL")); + var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2023)); + + vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); + + List results = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(5)); + assertThat(results).hasSize(3); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'NL'")); + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country in ['BG']")); + + assertThat(results).hasSize(2); + assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG' && year == 2020")); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("NOT(country == 'BG' && year == 2020)")); + + assertThat(results).hasSize(2); + assertThat(results.get(0).getId()).isIn(nlDocument.getId(), bgDocument2.getId()); + assertThat(results.get(1).getId()).isIn(nlDocument.getId(), bgDocument2.getId()); + + }); + } + + @Test + void searchWithThreshold() { + + contextRunner.run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + resetCollection(vectorStore); + + vectorStore.add(documents); + + List fullResult = vectorStore + .similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThresholdAll()); + + List distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList(); + + assertThat(distances).hasSize(3); + + float threshold = (distances.get(0) + distances.get(1)) / 2; + + List results = vectorStore + .similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThreshold(1 - threshold)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId()); + assertThat(resultDoc.getContent()).contains( + "Spring AI provides abstractions that serve as the foundation for developing AI applications."); + assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance"); + + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Bean + public VectorStore vectorStore(Client client, EmbeddingModel embeddingClient) { + + TypesenseVectorStoreConfig config = TypesenseVectorStoreConfig.builder() + .withCollectionName("test_vector_store") + .withEmbeddingDimension(embeddingClient.dimensions()) + .build(); + + return new TypesenseVectorStore(client, embeddingClient, config); + } + + @Bean + public Client typesenseClient() { + List nodes = new ArrayList<>(); + nodes + .add(new Node("http", typesenseContainer.getHost(), typesenseContainer.getMappedPort(8108).toString())); + + Configuration configuration = new Configuration(nodes, Duration.ofSeconds(5), "xyz"); + return new Client(configuration); + } + + @Bean + public EmbeddingModel embeddingClient() { + return new TransformersEmbeddingModel(); + } + + } + + @AfterAll + static void deleteContainer() { + if (typesenseContainer != null) { + typesenseContainer.stop(); + } + + if (tempDirectory != null) { + tempDirectory.toFile().delete(); + } + } + +}