From 7b06fcf98b2daf34f72870e058cb1030eb91717d Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Mon, 21 Oct 2024 18:41:42 +0100 Subject: [PATCH] Add Azure CosmosDB vector store support - Implement core vector store module for CosmosDB integration - Add Spring Boot auto-configuration capabilities - Integrate batch processing strategy for optimized operations - Include comprehensive tests for core and auto-config modules - Add reference docs for the CosmosDB vector store support --- pom.xml | 1 + .../conventions/VectorStoreProvider.java | 1 + .../pages/api/vectordbs/azure-cosmos-db.adoc | 237 +++++++++++++ spring-ai-spring-boot-autoconfigure/pom.xml | 8 + .../CosmosDBVectorStoreAutoConfiguration.java | 73 ++++ .../CosmosDBVectorStoreProperties.java | 112 ++++++ ...osmosDBVectorStoreAutoConfigurationIT.java | 177 ++++++++++ .../pom.xml | 42 +++ .../spring-ai-azure-cosmos-db-store/README.md | 1 + .../spring-ai-azure-cosmos-db-store/pom.xml | 81 +++++ .../CosmosDBFilterExpressionConverter.java | 142 ++++++++ .../ai/vectorstore/CosmosDBVectorStore.java | 327 ++++++++++++++++++ .../CosmosDBVectorStoreConfig.java | 120 +++++++ .../ai/vectorstore/CosmosDBVectorStoreIT.java | 200 +++++++++++ .../src/test/resources/application.properties | 4 + 15 files changed, 1526 insertions(+) create mode 100644 spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/azure-cosmos-db.adoc create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfiguration.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreProperties.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfigurationIT.java create mode 100644 spring-ai-spring-boot-starters/spring-ai-starter-azure-cosmos-db-store/pom.xml create mode 100644 vector-stores/spring-ai-azure-cosmos-db-store/README.md create mode 100644 vector-stores/spring-ai-azure-cosmos-db-store/pom.xml create mode 100644 vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBFilterExpressionConverter.java create mode 100644 vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java create mode 100644 vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreConfig.java create mode 100644 vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreIT.java create mode 100644 vector-stores/spring-ai-azure-cosmos-db-store/src/test/resources/application.properties diff --git a/pom.xml b/pom.xml index 4ab5e3fcb..463665082 100644 --- a/pom.xml +++ b/pom.xml @@ -27,6 +27,7 @@ document-readers/pdf-reader document-readers/tika-reader + vector-stores/spring-ai-azure-cosmos-db-store vector-stores/spring-ai-azure-store vector-stores/spring-ai-cassandra-store vector-stores/spring-ai-chroma-store diff --git a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java index 3d3d76c15..2ceaf2f54 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java @@ -34,6 +34,7 @@ public enum VectorStoreProvider { AZURE("azure"), CASSANDRA("cassandra"), CHROMA("chroma"), + COSMOSDB("cosmosdb"), ELASTICSEARCH("elasticsearch"), GEMFIRE("gemfire"), HANA("hana"), diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/azure-cosmos-db.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/azure-cosmos-db.adoc new file mode 100644 index 000000000..89c2cc150 --- /dev/null +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/azure-cosmos-db.adoc @@ -0,0 +1,237 @@ += Azure Cosmos DB + +This section walks you through setting up `CosmosDBVectorStore` to store document embeddings and perform similarity searches. + +== What is Azure Cosmos DB? + +link:https://azure.microsoft.com/en-us/services/cosmos-db/[Azure Cosmos DB] is Microsoft's globally distributed cloud-native database service designed for mission-critical applications. +It offers high availability, low latency, and the ability to scale horizontally to meet modern application demands. +It was built from the ground up with global distribution, fine-grained multi-tenancy, and horizontal scalability at its core. +It is a foundational service in Azure, used by most of Microsoft’s mission critical applications at global scale, including Teams, Skype, Xbox Live, Office 365, Bing, Azure Active Directory, Azure Portal, Microsoft Store, and many others. +It is also used by thousands of external customers including OpenAI for ChatGPT and other mission-critical AI applications that require elastic scale, turnkey global distribution, and low latency and high availability across the planet. + +== What is DiskANN? + +DiskANN (Disk-based Approximate Nearest Neighbor Search) is an innovative technology used in Azure Cosmos DB to enhance the performance of vector searches. +It enables efficient and scalable similarity searches across high-dimensional data by indexing embeddings stored in Cosmos DB. + +DiskANN provides the following benefits: + +* **Efficiency**: By utilizing disk-based structures, DiskANN significantly reduces the time required to find nearest neighbors compared to traditional methods. +* **Scalability**: It can handle large datasets that exceed memory capacity, making it suitable for various applications, including machine learning and AI-driven solutions. +* **Low Latency**: DiskANN minimizes latency during search operations, ensuring that applications can retrieve results quickly even with substantial data volumes. + +In the context of Spring AI for Azure Cosmos DB, vector searches will create and leverage DiskANN indexes to ensure optimal performance for similarity queries. + +== Setting up Azure Cosmos DB Vector Store with Auto Configuration + +The following code demonstrates how to set up the `CosmosDBVectorStore` with auto-configuration: + +```java +package com.example.demo; + +import io.micrometer.observation.ObservationRegistry; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.CommandLineRunner; +import org.springframework.boot.SpringApplication; +import org.springframework.boot.autoconfigure.SpringBootApplication; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Lazy; + +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; + +@SpringBootApplication +@EnableAutoConfiguration +public class DemoApplication implements CommandLineRunner { + + private static final Logger log = LoggerFactory.getLogger(DemoApplication.class); + + @Lazy + @Autowired + private VectorStore vectorStore; + + public static void main(String[] args) { + SpringApplication.run(DemoApplication.class, args); + } + + @Override + public void run(String... args) throws Exception { + Document document1 = new Document(UUID.randomUUID().toString(), "Sample content1", Map.of("key1", "value1")); + Document document2 = new Document(UUID.randomUUID().toString(), "Sample content2", Map.of("key2", "value2")); + vectorStore.add(List.of(document1, document2)); + List results = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1)); + + log.info("Search results: {}", results); + + // Remove the documents from the vector store + vectorStore.delete(List.of(document1.getId(), document2.getId())); + } + + @Bean + public ObservationRegistry observationRegistry() { + return ObservationRegistry.create(); + } +} +``` + + +== Auto Configuration + +Add the following dependency to your Maven project: + +[source,xml] +---- + + org.springframework.ai + spring-ai-azure-cosmos-db-store-spring-boot-starter + +---- + +== Configuration Properties + +The following configuration properties are available for the Cosmos DB vector store: + +|=========================== +| Property | Description +| spring.ai.vectorstore.cosmosdb.databaseName | The name of the Cosmos DB database to use. +| spring.ai.vectorstore.cosmosdb.containerName | The name of the Cosmos DB container to use. +| spring.ai.vectorstore.cosmosdb.partitionKeyPath | The path for the partition key. +| spring.ai.vectorstore.cosmosdb.metadataFields | Comma-separated list of metadata fields. +| spring.ai.vectorstore.cosmosdb.vectorStoreThoughput | The throughput for the vector store. +| spring.ai.vectorstore.cosmosdb.vectorDimensions | The number of dimensions for the vectors. +| spring.ai.vectorstore.cosmosdb.endpoint | The endpoint for the Cosmos DB. +| spring.ai.vectorstore.cosmosdb.key | The key for the Cosmos DB. +|=========================== + + +== Complex Searches with Filters + +You can perform more complex searches using filters in the Cosmos DB vector store. +Below is a sample demonstrating how to use filters in your search queries. + +[source,java] +---- +Map metadata1 = new HashMap<>(); +metadata1.put("country", "UK"); +metadata1.put("year", 2021); +metadata1.put("city", "London"); + +Map metadata2 = new HashMap<>(); +metadata2.put("country", "NL"); +metadata2.put("year", 2022); +metadata2.put("city", "Amsterdam"); + +Document document1 = new Document("1", "A document about the UK", metadata1); +Document document2 = new Document("2", "A document about the Netherlands", metadata2); + +vectorStore.add(List.of(document1, document2)); + +FilterExpressionBuilder builder = new FilterExpressionBuilder(); +List results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(10) + .withFilterExpression((builder.in("country", "UK", "NL")).build())); +---- + +== Setting up Azure Cosmos DB Vector Store without Auto Configuration + +The following code demonstrates how to set up the `CosmosDBVectorStore` without relying on auto-configuration: + +```java +package com.example.demo; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosClientBuilder; +import io.micrometer.observation.ObservationRegistry; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.CosmosDBVectorStore; +import org.springframework.ai.vectorstore.CosmosDBVectorStoreConfig; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.CommandLineRunner; +import org.springframework.boot.SpringApplication; +import org.springframework.boot.autoconfigure.SpringBootApplication; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Lazy; + +import java.util.List; +import java.util.Map; +import java.util.UUID; + +@SpringBootApplication +public class DemoApplication implements CommandLineRunner { + + @Lazy + @Autowired + private VectorStore vectorStore; + + @Lazy + @Autowired + private EmbeddingModel embeddingModel; + + public static void main(String[] args) { + SpringApplication.run(DemoApplication.class, args); + } + + @Override + public void run(String... args) throws Exception { + Document document1 = new Document(UUID.randomUUID().toString(), "Sample content1", Map.of("key1", "value1")); + Document document2 = new Document(UUID.randomUUID().toString(), "Sample content2", Map.of("key2", "value2")); + vectorStore.add(List.of(document1, document2)); + + List results = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1)); + log.info("Search results: {}", results); + } + + @Bean + public ObservationRegistry observationRegistry() { + return ObservationRegistry.create(); + } + + @Bean + public VectorStore vectorStore(ObservationRegistry observationRegistry) { + CosmosDBVectorStoreConfig config = new CosmosDBVectorStoreConfig(); + config.setDatabaseName("spring-ai-sample"); + config.setContainerName("container"); + config.setMetadataFields("country,city"); + config.setVectorStoreThoughput(400); + + CosmosAsyncClient cosmosClient = new CosmosClientBuilder() + .endpoint(System.getenv("COSMOSDB_AI_ENDPOINT")) + .userAgentSuffix("SpringAI-CDBNoSQL-VectorStore") + .key(System.getenv("COSMOSDB_AI_KEY")) + .gatewayMode() + .buildAsyncClient(); + + return new CosmosDBVectorStore(observationRegistry, null, cosmosClient, config, embeddingModel); + } + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } +} +``` + +== Manual Dependency Setup + +Add the following dependency in your Maven project: + +[source,xml] +---- + + org.springframework.ai + spring-ai-azure-cosmos-db-store + +---- \ No newline at end of file diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index c025947f0..2a738a8d4 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -387,6 +387,14 @@ true + + + org.springframework.ai + spring-ai-azure-cosmos-db-store + ${project.parent.version} + true + + diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfiguration.java new file mode 100644 index 000000000..6498d233b --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfiguration.java @@ -0,0 +1,73 @@ +/* + * Copyright 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.autoconfigure.vectorstore.cosmosdb; + +import com.azure.cosmos.CosmosClientBuilder; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.CosmosDBVectorStore; +import org.springframework.ai.vectorstore.CosmosDBVectorStoreConfig; +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.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; +import com.azure.cosmos.CosmosAsyncClient; +import io.micrometer.observation.ObservationRegistry; + +/** + * @author Theo van Kraay + * @since 1.0.0 + */ + +@AutoConfiguration +@ConditionalOnClass({ CosmosDBVectorStore.class, EmbeddingModel.class, CosmosAsyncClient.class }) +@EnableConfigurationProperties(CosmosDBVectorStoreProperties.class) +public class CosmosDBVectorStoreAutoConfiguration { + + String endpoint; + + String key; + + @Bean + public CosmosAsyncClient cosmosClient(CosmosDBVectorStoreProperties properties) { + return new CosmosClientBuilder().endpoint(properties.getEndpoint()) + .userAgentSuffix("SpringAI-CDBNoSQL-VectorStore") + .key(properties.getKey()) + .gatewayMode() + .buildAsyncClient(); + } + + @Bean + @ConditionalOnMissingBean + public CosmosDBVectorStore cosmosDBVectorStore(ObservationRegistry observationRegistry, + ObjectProvider customObservationConvention, + CosmosDBVectorStoreProperties properties, CosmosAsyncClient cosmosAsyncClient, + EmbeddingModel embeddingModel) { + + CosmosDBVectorStoreConfig config = new CosmosDBVectorStoreConfig(); + config.setDatabaseName(properties.getDatabaseName()); + config.setContainerName(properties.getContainerName()); + config.setMetadataFields(properties.getMetadataFields()); + config.setVectorStoreThoughput(properties.getVectorStoreThoughput()); + config.setVectorDimensions(properties.getVectorDimensions()); + return new CosmosDBVectorStore(observationRegistry, customObservationConvention.getIfAvailable(), + cosmosAsyncClient, config, embeddingModel); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreProperties.java new file mode 100644 index 000000000..732573de2 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreProperties.java @@ -0,0 +1,112 @@ +/* + * Copyright 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.autoconfigure.vectorstore.cosmosdb; + +import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * @author Theo van Kraay + * @since 1.0.0 + */ + +@ConfigurationProperties(CosmosDBVectorStoreProperties.CONFIG_PREFIX) +public class CosmosDBVectorStoreProperties extends CommonVectorStoreProperties { + + public static final String CONFIG_PREFIX = "spring.ai.vectorstore.cosmosdb"; + + private String containerName; + + private String databaseName; + + private String metadataFields; + + private int vectorStoreThoughput = 400; + + private long vectorDimensions = 1536; + + private String partitionKeyPath; + + private String endpoint; + + private String key; + + public int getVectorStoreThoughput() { + return vectorStoreThoughput; + } + + public void setVectorStoreThoughput(int vectorStoreThoughput) { + this.vectorStoreThoughput = vectorStoreThoughput; + } + + public String getMetadataFields() { + return metadataFields; + } + + public void setMetadataFields(String metadataFields) { + this.metadataFields = metadataFields; + } + + public String getEndpoint() { + return endpoint; + } + + public void setEndpoint(String endpoint) { + this.endpoint = endpoint; + } + + public String getKey() { + return key; + } + + public void setKey(String key) { + this.key = key; + } + + public String getDatabaseName() { + return databaseName; + } + + public void setDatabaseName(String databaseName) { + this.databaseName = databaseName; + } + + public String getContainerName() { + return containerName; + } + + public void setContainerName(String containerName) { + this.containerName = containerName; + } + + public String getPartitionKeyPath() { + return partitionKeyPath; + } + + public void setPartitionKeyPath(String partitionKeyPath) { + this.partitionKeyPath = partitionKeyPath; + } + + public long getVectorDimensions() { + return vectorDimensions; + } + + public void setVectorDimensions(long vectorDimensions) { + this.vectorDimensions = vectorDimensions; + } + +} \ No newline at end of file diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfigurationIT.java new file mode 100644 index 000000000..c6c4ea01e --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfigurationIT.java @@ -0,0 +1,177 @@ +/* + * 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.autoconfigure.vectorstore.cosmosdb; + +import io.micrometer.observation.tck.TestObservationRegistry; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +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.ai.vectorstore.filter.FilterExpressionBuilder; +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 java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Theo van Kraay + * @since 1.0.0 + */ + +@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_ENDPOINT", matches = ".+") +@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_KEY", matches = ".+") +public class CosmosDBVectorStoreAutoConfigurationIT { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(CosmosDBVectorStoreAutoConfiguration.class)) + .withPropertyValues("spring.ai.vectorstore.cosmosdb.databaseName=test-database") + .withPropertyValues("spring.ai.vectorstore.cosmosdb.containerName=test-container") + .withPropertyValues("spring.ai.vectorstore.cosmosdb.partitionKeyPath=/id") + .withPropertyValues("spring.ai.vectorstore.cosmosdb.metadataFields=country,year,city") + .withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThoughput=1000") + .withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorDimensions=384") + .withPropertyValues("spring.ai.vectorstore.cosmosdb.endpoint=" + System.getenv("AZURE_COSMOSDB_ENDPOINT")) + .withPropertyValues("spring.ai.vectorstore.cosmosdb.key=" + System.getenv("AZURE_COSMOSDB_KEY")) + .withUserConfiguration(Config.class); + + private VectorStore vectorStore; + + @BeforeEach + public void setup() { + contextRunner.run(context -> { + vectorStore = context.getBean(VectorStore.class); + }); + } + + @Test + public void testAddSearchAndDeleteDocuments() { + + // Create a sample document + Document document1 = new Document(UUID.randomUUID().toString(), "Sample content1", Map.of("key1", "value1")); + Document document2 = new Document(UUID.randomUUID().toString(), "Sample content2", Map.of("key2", "value2")); + + // Add the document to the vector store + vectorStore.add(List.of(document1, document2)); + + // Perform a similarity search + List results = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1)); + + // Verify the search results + assertThat(results).isNotEmpty(); + assertThat(results.get(0).getId()).isEqualTo(document1.getId()); + + // Remove the documents from the vector store + vectorStore.delete(List.of(document1.getId(), document2.getId())); + + // Perform a similarity search again + List results2 = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1)); + + // Verify the search results + assertThat(results2).isEmpty(); + } + + @Test + void testSimilaritySearchWithFilter() { + + // Insert documents using vectorStore.add + Map metadata1; + metadata1 = new HashMap<>(); + metadata1.put("country", "UK"); + metadata1.put("year", 2021); + metadata1.put("city", "London"); + + Map metadata2; + metadata2 = new HashMap<>(); + metadata2.put("country", "NL"); + metadata2.put("year", 2022); + metadata2.put("city", "Amsterdam"); + + Map metadata3; + metadata3 = new HashMap<>(); + metadata3.put("country", "US"); + metadata3.put("year", 2019); + metadata3.put("city", "Sofia"); + + Map metadata4; + metadata4 = new HashMap<>(); + metadata4.put("country", "US"); + metadata4.put("year", 2020); + metadata4.put("city", "Sofia"); + + Document document1 = new Document("1", "A document about the UK", metadata1); + Document document2 = new Document("2", "A document about the Netherlands", metadata2); + Document document3 = new Document("3", "A document about the US", metadata3); + Document document4 = new Document("4", "A document about the US", metadata4); + + vectorStore.add(List.of(document1, document2, document3, document4)); + FilterExpressionBuilder b = new FilterExpressionBuilder(); + List results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(10) + .withFilterExpression((b.in("country", "UK", "NL")).build())); + + assertThat(results).hasSize(2); + assertThat(results).extracting(Document::getId).containsExactlyInAnyOrder("1", "2"); + + List results2 = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(10) + .withFilterExpression( + b.and(b.or(b.gte("year", 2021), b.eq("country", "NL")), b.ne("city", "Amsterdam")).build())); + + assertThat(results2).hasSize(1); + assertThat(results2).extracting(Document::getId).containsExactlyInAnyOrder("1"); + + List results3 = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(10) + .withFilterExpression(b.and(b.eq("country", "US"), b.eq("year", 2020)).build())); + + assertThat(results3).hasSize(1); + assertThat(results3).extracting(Document::getId).containsExactlyInAnyOrder("4"); + + vectorStore.delete(List.of(document1.getId(), document2.getId(), document3.getId(), document4.getId())); + + // Perform a similarity search again + List results4 = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(1)); + + // Verify the search results + assertThat(results4).isEmpty(); + } + + @Configuration(proxyBeanMethods = false) + static class Config { + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } + + @Bean + public TestObservationRegistry observationRegistry() { + return TestObservationRegistry.create(); + } + + } + +} \ No newline at end of file diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-azure-cosmos-db-store/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-azure-cosmos-db-store/pom.xml new file mode 100644 index 000000000..ac8c2ffb1 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-azure-cosmos-db-store/pom.xml @@ -0,0 +1,42 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-azure-cosmos-db-store-spring-boot-starter + jar + Spring AI Starter - Azure Cosmos DB Vector Store + Spring AI Azure Cosmos DB Vector Store 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-azure-cosmos-db-store + ${project.parent.version} + + + + diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/README.md b/vector-stores/spring-ai-azure-cosmos-db-store/README.md new file mode 100644 index 000000000..8e610a268 --- /dev/null +++ b/vector-stores/spring-ai-azure-cosmos-db-store/README.md @@ -0,0 +1 @@ +[Azure Cosmos DB Vector Store Documentation]() \ No newline at end of file diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/pom.xml b/vector-stores/spring-ai-azure-cosmos-db-store/pom.xml new file mode 100644 index 000000000..0410a18bd --- /dev/null +++ b/vector-stores/spring-ai-azure-cosmos-db-store/pom.xml @@ -0,0 +1,81 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-azure-cosmos-db-store + jar + Spring AI Vector Store – Azure Cosmos DB + Spring AI Vector Store for Azure Cosmos DB + 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 + + + + + com.azure + azure-spring-data-cosmos + LATEST + + + org.springframework.ai + spring-ai-core + ${project.parent.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.testcontainers + junit-jupiter + test + + + + org.testcontainers + azure + 1.20.1 + + + + io.micrometer + micrometer-observation-test + test + + + + + diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBFilterExpressionConverter.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBFilterExpressionConverter.java new file mode 100644 index 000000000..7424b25e5 --- /dev/null +++ b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBFilterExpressionConverter.java @@ -0,0 +1,142 @@ +/* + * Copyright 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.vectorstore; + +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.Filter.ExpressionType; +import org.springframework.ai.vectorstore.filter.Filter.Key; +import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter; + +import java.util.Collection; +import java.util.Map; +import java.util.Optional; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR; + +/** + * Converts {@link org.springframework.ai.vectorstore.filter.Filter.Expression} into + * Cosmos DB NoSQL API where clauses. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +class CosmosDBFilterExpressionConverter extends AbstractFilterExpressionConverter { + + private Map metadataFields; + + public CosmosDBFilterExpressionConverter(Collection columns) { + this.metadataFields = columns.stream().collect(Collectors.toMap(Function.identity(), Function.identity())); + } + + /** + * Gets the metadata field from the Cosmos DB document. + * @param name The name of the metadata field. + * @return The name of the metadata field as it should appear in the query. + */ + private Optional getMetadataField(String name) { + String metadataField = name; + return Optional.ofNullable(metadataFields.get(metadataField)); + } + + @Override + protected void doKey(Key key, StringBuilder context) { + String keyName = key.key(); + Optional metadataField = getMetadataField(keyName); + if (metadataField.isPresent()) { + context.append("c.metadata." + metadataField.get()); + } + else { + throw new IllegalArgumentException(String.format("No metadata field %s has been configured", keyName)); + } + } + + @Override + protected void doExpression(Filter.Expression expression, StringBuilder context) { + // Handling AND/OR + if (AND.equals(expression.type()) || OR.equals(expression.type())) { + doCompoundExpressionType(expression, context); + } + else { + doSingleExpressionType(expression, context); + } + } + + private void doCompoundExpressionType(Filter.Expression expression, StringBuilder context) { + context.append(" ("); + this.convertOperand(expression.left(), context); + context.append(getOperationSymbol(expression)); + context.append(" ("); + this.convertOperand(expression.right(), context); + int start = context.indexOf("["); + if (start != -1) { + context.replace(start, start + 1, ""); + } + int end = context.indexOf("]"); + if (end != -1) { + context.replace(end, end + 1, ""); + } + context.append(")"); + context.append(")"); + } + + private void doSingleExpressionType(Filter.Expression expression, StringBuilder context) { + this.convertOperand(expression.left(), context); + context.append(getOperationSymbol(expression)); + context.append(" ("); + this.convertOperand(expression.right(), context); + int start = context.indexOf("["); + if (start != -1) { + context.replace(start, start + 1, ""); + } + int end = context.indexOf("]"); + if (end != -1) { + context.replace(end, end + 1, ""); + } + context.append(")"); + } + + private String getOperationSymbol(Filter.Expression exp) { + switch (exp.type()) { + case AND: + return " AND "; + case OR: + return " OR "; + case EQ: + return " = "; + case NE: + return " != "; + case LT: + return " < "; + case LTE: + return " <= "; + case GT: + return " > "; + case GTE: + return " >= "; + case IN: + return " IN "; + case NIN: + return " !IN "; + default: + throw new RuntimeException("Not supported expression type:" + exp.type()); + } + } + +} diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java new file mode 100644 index 000000000..fbbcb1f72 --- /dev/null +++ b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java @@ -0,0 +1,327 @@ +/* + * Copyright 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.vectorstore; + +import com.azure.cosmos.*; +import com.azure.cosmos.implementation.guava25.collect.ImmutableList; +import com.azure.cosmos.models.*; +import com.azure.cosmos.util.CosmosPagedFlux; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.micrometer.observation.ObservationRegistry; +import org.apache.commons.lang3.tuple.ImmutablePair; +import org.apache.commons.lang3.tuple.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.BatchingStrategy; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.embedding.EmbeddingOptionsBuilder; +import org.springframework.ai.embedding.TokenCountBatchingStrategy; +import org.springframework.ai.observation.conventions.VectorStoreProvider; +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; +import reactor.core.publisher.Flux; +import java.util.*; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +/** + * @author Theo van Kraay + * @since 1.0.0 + */ + +public class CosmosDBVectorStore extends AbstractObservationVectorStore implements AutoCloseable { + + private static final Logger logger = LoggerFactory.getLogger(CosmosDBVectorStore.class); + + private final CosmosAsyncClient cosmosClient; + + private CosmosAsyncContainer container; + + private final EmbeddingModel embeddingModel; + + private final CosmosDBVectorStoreConfig properties; + + private final BatchingStrategy batchingStrategy; + + public CosmosDBVectorStore(ObservationRegistry observationRegistry, + VectorStoreObservationConvention customObservationConvention, CosmosAsyncClient cosmosClient, + CosmosDBVectorStoreConfig properties, EmbeddingModel embeddingModel) { + super(observationRegistry, customObservationConvention); + this.cosmosClient = cosmosClient; + this.properties = properties; + this.batchingStrategy = new TokenCountBatchingStrategy(); + cosmosClient.createDatabaseIfNotExists(properties.getDatabaseName()).block(); + + initializeContainer(properties.getContainerName(), properties.getDatabaseName(), + properties.getVectorStoreThoughput(), properties.getVectorDimensions(), + properties.getPartitionKeyPath()); + + this.embeddingModel = embeddingModel; + + } + + private void initializeContainer(String containerName, String databaseName, int vectorStoreThoughput, + long vectorDimensions, String partitionKeyPath) { + + // Set defaults if not provided + if (vectorStoreThoughput == 0) { + vectorStoreThoughput = 400; + } + if (partitionKeyPath == null) { + partitionKeyPath = "/id"; + } + + // handle hierarchical partition key + PartitionKeyDefinition subpartitionKeyDefinition = new PartitionKeyDefinition(); + List pathsfromCommaSeparatedList = new ArrayList(); + String[] subpartitionKeyPaths = partitionKeyPath.split(","); + for (String path : subpartitionKeyPaths) { + pathsfromCommaSeparatedList.add(path); + } + if (subpartitionKeyPaths.length > 1) { + subpartitionKeyDefinition.setPaths(pathsfromCommaSeparatedList); + subpartitionKeyDefinition.setKind(PartitionKind.MULTI_HASH); + } + else { + subpartitionKeyDefinition.setPaths(Collections.singletonList(partitionKeyPath)); + subpartitionKeyDefinition.setKind(PartitionKind.HASH); + } + CosmosContainerProperties collectionDefinition = new CosmosContainerProperties(containerName, + subpartitionKeyDefinition); + // Set vector embedding policy + CosmosVectorEmbeddingPolicy embeddingPolicy = new CosmosVectorEmbeddingPolicy(); + CosmosVectorEmbedding embedding = new CosmosVectorEmbedding(); + embedding.setPath("/embedding"); + embedding.setDataType(CosmosVectorDataType.FLOAT32); + embedding.setDimensions(vectorDimensions); + embedding.setDistanceFunction(CosmosVectorDistanceFunction.COSINE); + embeddingPolicy.setCosmosVectorEmbeddings(Collections.singletonList(embedding)); + collectionDefinition.setVectorEmbeddingPolicy(embeddingPolicy); + + // set vector indexing policy + IndexingPolicy indexingPolicy = new IndexingPolicy(); + indexingPolicy.setIndexingMode(IndexingMode.CONSISTENT); + ExcludedPath excludedPath = new ExcludedPath("/*"); + indexingPolicy.setExcludedPaths(Collections.singletonList(excludedPath)); + IncludedPath includedPath1 = new IncludedPath("/metadata/?"); + IncludedPath includedPath2 = new IncludedPath("/content/?"); + indexingPolicy.setIncludedPaths(ImmutableList.of(includedPath1, includedPath2)); + CosmosVectorIndexSpec cosmosVectorIndexSpec = new CosmosVectorIndexSpec(); + cosmosVectorIndexSpec.setPath("/embedding"); + cosmosVectorIndexSpec.setType(CosmosVectorIndexType.DISK_ANN.toString()); + indexingPolicy.setVectorIndexes(List.of(cosmosVectorIndexSpec)); + collectionDefinition.setIndexingPolicy(indexingPolicy); + + ThroughputProperties throughputProperties = ThroughputProperties.createManualThroughput(vectorStoreThoughput); + CosmosAsyncDatabase cosmosAsyncDatabase = cosmosClient.getDatabase(databaseName); + cosmosAsyncDatabase.createContainerIfNotExists(collectionDefinition, throughputProperties).block(); + this.container = cosmosAsyncDatabase.getContainer(containerName); + } + + @Override + public void close() { + if (cosmosClient != null) { + cosmosClient.close(); + logger.info("Cosmos DB client closed successfully."); + } + } + + private JsonNode mapCosmosDocument(Document document, float[] queryEmbedding) { + ObjectMapper objectMapper = new ObjectMapper(); + + String id = document.getId(); + String content = document.getContent(); + + // Convert metadata and embedding directly to JsonNode + JsonNode metadataNode = objectMapper.valueToTree(document.getMetadata()); + JsonNode embeddingNode = objectMapper.valueToTree(queryEmbedding); + + // Create an ObjectNode specifically + ObjectNode objectNode = objectMapper.createObjectNode(); + + // Use put for simple values and set for JsonNode values + objectNode.put("id", id); + objectNode.put("content", content); + objectNode.set("metadata", metadataNode); // Use set to add JsonNode directly + objectNode.set("embedding", embeddingNode); // Use set to add JsonNode directly + + return objectNode; + } + + @Override + public void doAdd(List documents) { + + // Batch the documents based on the batching strategy + this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy); + + // Create a list to hold both the CosmosItemOperation and the corresponding + // document ID + List> itemOperationsWithIds = documents.stream().map(doc -> { + CosmosItemOperation operation = CosmosBulkOperations + .getCreateItemOperation(mapCosmosDocument(doc, doc.getEmbedding()), new PartitionKey(doc.getId())); + return new ImmutablePair<>(doc.getId(), operation); // Pair the document ID + // with the operation + }).collect(Collectors.toList()); + + try { + // Extract just the CosmosItemOperations from the pairs + List itemOperations = itemOperationsWithIds.stream() + .map(ImmutablePair::getValue) + .collect(Collectors.toList()); + + container.executeBulkOperations(Flux.fromIterable(itemOperations)).doOnNext(response -> { + if (response != null && response.getResponse() != null) { + int statusCode = response.getResponse().getStatusCode(); + if (statusCode == 409) { + // Retrieve the ID associated with the failed operation + String documentId = itemOperationsWithIds.stream() + .filter(pair -> pair.getValue().equals(response.getOperation())) + .findFirst() + .map(ImmutablePair::getKey) + .orElse("Unknown ID"); // Fallback if the ID can't be found + + String errorMessage = String.format("Duplicate document id: %s", documentId); + logger.error(errorMessage); + throw new RuntimeException(errorMessage); // Throw an exception + // for status code 409 + } + else { + logger.info("Document added with status: {}", statusCode); + } + } + else { + logger.warn("Received a null response or null status code for a document operation."); + } + }) + .doOnError(error -> logger.error("Error adding document: {}", error.getMessage())) + .doOnComplete(() -> logger.info("Bulk operation completed successfully.")) + .blockLast(); // Block until the last item of the Flux is processed + } + catch (Exception e) { + logger.error("Exception occurred during bulk add operation: {}", e.getMessage(), e); + throw e; // Rethrow the exception after logging + } + } + + @Override + public Optional doDelete(List idList) { + try { + // Convert the list of IDs into bulk delete operations + List itemOperations = idList.stream() + .map(id -> CosmosBulkOperations.getDeleteItemOperation(id, new PartitionKey(id))) + .collect(Collectors.toList()); + + // Execute bulk delete operations synchronously by using blockLast() on the + // Flux + container.executeBulkOperations(Flux.fromIterable(itemOperations)) + .doOnNext(response -> logger.info("Document deleted with status: {}", + response.getResponse().getStatusCode())) + .doOnError(error -> logger.error("Error deleting document: {}", error.getMessage())) + .blockLast(); // This will block until all operations have finished + + return Optional.of(true); + } + catch (Exception e) { + logger.error("Exception while deleting documents: {}", e.getMessage()); + return Optional.of(false); + } + } + + @Override + public List similaritySearch(String query) { + return similaritySearch(SearchRequest.query(query)); + } + + @Override + public List doSimilaritySearch(SearchRequest request) { + // Ensure topK is within acceptable limits + if (request.getTopK() > 1000) { + throw new IllegalArgumentException("Top K must be 1000 or less."); + } + + // Convert query into vector embedding + float[] embedding = this.embeddingModel.embed(request.getQuery()); + + logger.info("similarity threshold: {}", request.getSimilarityThreshold()); + + List embeddingList = IntStream.range(0, embedding.length) + .mapToObj(i -> embedding[i]) + .collect(Collectors.toList()); + + // Start building query for similarity search + StringBuilder queryBuilder = new StringBuilder("SELECT TOP @topK * FROM c WHERE "); + queryBuilder.append("VectorDistance(c.embedding, @embedding) > @similarityThreshold"); + + // Handle filter expression if it's set + Filter.Expression filterExpression = request.getFilterExpression(); + if (filterExpression != null) { + CosmosDBFilterExpressionConverter filterExpressionConverter = new CosmosDBFilterExpressionConverter( + properties.getMetadataFieldsList()); // Use the expression directly as + // it handles the "metadata" + // fields internally + String filterQuery = filterExpressionConverter.convertExpression(filterExpression); + queryBuilder.append(" AND ").append(filterQuery); + } + + queryBuilder.append(" ORDER BY VectorDistance(c.embedding, @embedding)"); + + String query = queryBuilder.toString(); + List parameters = new ArrayList<>(); + parameters.add(new SqlParameter("@embedding", embeddingList)); + parameters.add(new SqlParameter("@topK", request.getTopK())); + parameters.add(new SqlParameter("@similarityThreshold", request.getSimilarityThreshold())); + + SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query, parameters); + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions(); + + CosmosPagedFlux pagedFlux = container.queryItems(sqlQuerySpec, options, JsonNode.class); + + logger.info("Executing similarity search query: {}", query); + try { + // Collect documents from the paged flux + List documents = pagedFlux.byPage() + .flatMap(page -> Flux.fromIterable(page.getResults())) + .collectList() + .block(); + // Convert JsonNode to Document + List docs = documents.stream() + .map(doc -> new Document(doc.get("id").asText(), doc.get("content").asText(), new HashMap<>())) + .collect(Collectors.toList()); + + return docs != null ? docs : List.of(); + } + catch (Exception e) { + logger.error("Error during similarity search: {}", e.getMessage()); + return List.of(); + } + } + + @Override + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { + return VectorStoreObservationContext.builder(VectorStoreProvider.COSMOSDB.value(), operationName) + .withCollectionName(container.getId()) + .withDimensions(this.embeddingModel.dimensions()) + .withNamespace(container.getDatabase().getId()) + .withSimilarityMetric("cosine"); + } + +} diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreConfig.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreConfig.java new file mode 100644 index 000000000..f40f85556 --- /dev/null +++ b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreConfig.java @@ -0,0 +1,120 @@ +/* + * Copyright 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.vectorstore; + +import java.util.List; + +/** + * @author Theo van Kraay + * @since 1.0.0 + */ + +public class CosmosDBVectorStoreConfig implements AutoCloseable { + + private String containerName; + + private String databaseName; + + private String partitionKeyPath; + + private String endpoint; + + private String key; + + private String metadataFields; + + private int vectorStoreThoughput = 400; + + private long vectorDimensions = 1536; + + private List metadataFieldsList; + + public int getVectorStoreThoughput() { + return vectorStoreThoughput; + } + + public void setVectorStoreThoughput(int vectorStoreThoughput) { + this.vectorStoreThoughput = vectorStoreThoughput; + } + + public void setMetadataFields(String metadataFields) { + this.metadataFields = metadataFields; + this.metadataFieldsList = List.of(metadataFields.split(",")); + } + + public String getMetadataFields() { + return metadataFields; + } + + public List getMetadataFieldsList() { + return metadataFieldsList; + } + + public String getEndpoint() { + return endpoint; + } + + public void setEndpoint(String endpoint) { + this.endpoint = endpoint; + } + + public String getKey() { + return key; + } + + public void setKey(String key) { + this.key = key; + } + + public String getContainerName() { + return containerName; + } + + public void setContainerName(String containerName) { + this.containerName = containerName; + } + + public String getDatabaseName() { + return databaseName; + } + + public void setDatabaseName(String databaseName) { + this.databaseName = databaseName; + } + + public String getPartitionKeyPath() { + return partitionKeyPath; + } + + public void setPartitionKeyPath(String partitionKeyPath) { + this.partitionKeyPath = partitionKeyPath; + } + + @Override + public void close() throws Exception { + + } + + public long getVectorDimensions() { + return vectorDimensions; + } + + public void setVectorDimensions(long vectorDimensions) { + this.vectorDimensions = vectorDimensions; + } + +} diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreIT.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreIT.java new file mode 100644 index 000000000..4269a0d72 --- /dev/null +++ b/vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreIT.java @@ -0,0 +1,200 @@ +/* + * Copyright 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.vectorstore; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosClientBuilder; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * @author Theo van Kraay + * @since 1.0.0 + */ + +@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_ENDPOINT", matches = ".+") +@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_KEY", matches = ".+") +public class CosmosDBVectorStoreIT { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class); + + private VectorStore vectorStore; + + @BeforeEach + public void setup() { + contextRunner.run(context -> { + vectorStore = context.getBean(VectorStore.class); + }); + } + + @Test + public void testAddSearchAndDeleteDocuments() { + + // Create a sample document + Document document1 = new Document(UUID.randomUUID().toString(), "Sample content1", Map.of("key1", "value1")); + Document document2 = new Document(UUID.randomUUID().toString(), "Sample content2", Map.of("key2", "value2")); + + // Add the document to the vector store + vectorStore.add(List.of(document1, document2)); + + // create duplicate docs and assert that second one throws exception + Document document3 = new Document(document1.getId(), "Sample content3", Map.of("key3", "value3")); + assertThatThrownBy(() -> vectorStore.add(List.of(document3))).isInstanceOf(Exception.class) + .hasMessageContaining("Duplicate document id: " + document1.getId()); + + // Perform a similarity search + List results = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1)); + + // Verify the search results + assertThat(results).isNotEmpty(); + assertThat(results.get(0).getId()).isEqualTo(document1.getId()); + + // Remove the documents from the vector store + vectorStore.delete(List.of(document1.getId(), document2.getId())); + + // Perform a similarity search again + List results2 = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1)); + + // Verify the search results + assertThat(results2).isEmpty(); + + } + + @Test + void testSimilaritySearchWithFilter() { + + // Insert documents using vectorStore.add + Map metadata1; + metadata1 = new HashMap<>(); + metadata1.put("country", "UK"); + metadata1.put("year", 2021); + metadata1.put("city", "London"); + + Map metadata2; + metadata2 = new HashMap<>(); + metadata2.put("country", "NL"); + metadata2.put("year", 2022); + metadata2.put("city", "Amsterdam"); + + Map metadata3; + metadata3 = new HashMap<>(); + metadata3.put("country", "US"); + metadata3.put("year", 2019); + metadata3.put("city", "Sofia"); + + Map metadata4; + metadata4 = new HashMap<>(); + metadata4.put("country", "US"); + metadata4.put("year", 2020); + metadata4.put("city", "Sofia"); + + Document document1 = new Document("1", "A document about the UK", metadata1); + Document document2 = new Document("2", "A document about the Netherlands", metadata2); + Document document3 = new Document("3", "A document about the US", metadata3); + Document document4 = new Document("4", "A document about the US", metadata4); + + vectorStore.add(List.of(document1, document2, document3, document4)); + FilterExpressionBuilder b = new FilterExpressionBuilder(); + List results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(10) + .withFilterExpression((b.in("country", "UK", "NL")).build())); + + assertThat(results).hasSize(2); + assertThat(results).extracting(Document::getId).containsExactlyInAnyOrder("1", "2"); + + List results2 = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(10) + .withFilterExpression( + b.and(b.or(b.gte("year", 2021), b.eq("country", "NL")), b.ne("city", "Amsterdam")).build())); + + assertThat(results2).hasSize(1); + assertThat(results2).extracting(Document::getId).containsExactlyInAnyOrder("1"); + + List results3 = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(10) + .withFilterExpression(b.and(b.eq("country", "US"), b.eq("year", 2020)).build())); + + assertThat(results3).hasSize(1); + assertThat(results3).extracting(Document::getId).containsExactlyInAnyOrder("4"); + + vectorStore.delete(List.of(document1.getId(), document2.getId(), document3.getId(), document4.getId())); + + // Perform a similarity search again + List results4 = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(1)); + + // Verify the search results + assertThat(results4).isEmpty(); + } + + @SpringBootConfiguration + @EnableAutoConfiguration + public static class TestApplication { + + @Bean + public VectorStore vectorStore(CosmosAsyncClient cosmosClient, EmbeddingModel embeddingModel, + VectorStoreObservationConvention convention) { + CosmosDBVectorStoreConfig config = new CosmosDBVectorStoreConfig(); + config.setDatabaseName("test-database"); + config.setContainerName("test-container"); + config.setMetadataFields("country,year,city"); + config.setVectorStoreThoughput(1000); + return new CosmosDBVectorStore(null, convention, cosmosClient, config, embeddingModel); + + } + + @Bean + public CosmosAsyncClient cosmosClient() { + return new CosmosClientBuilder().endpoint(System.getenv("AZURE_COSMOSDB_ENDPOINT")) + .key(System.getenv("AZURE_COSMOSDB_KEY")) + .userAgentSuffix("SpringAI-CDBNoSQL-VectorStore") + .gatewayMode() + .buildAsyncClient(); + } + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } + + @Bean + public VectorStoreObservationConvention observationConvention() { + // Replace with an actual observation convention or a mock if needed + return new VectorStoreObservationConvention() { + }; + } + + } + +} diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/test/resources/application.properties b/vector-stores/spring-ai-azure-cosmos-db-store/src/test/resources/application.properties new file mode 100644 index 000000000..20c6c6220 --- /dev/null +++ b/vector-stores/spring-ai-azure-cosmos-db-store/src/test/resources/application.properties @@ -0,0 +1,4 @@ +spring.ai.vectorstore.cosmosdb.databaseName=db +spring.ai.vectorstore.cosmosdb.containerName=container +spring.ai.vectorstore.cosmosdb.key=${COSMOSDB_AI_ENDPOINT} +spring.ai.vectorstore.cosmosdb.uri=${COSMOSDB_AI_KEY} \ No newline at end of file