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
2024-10-21 18:41:42 +01:00
parent 2c17577f2f
commit 7b06fcf98b
15 changed files with 1526 additions and 0 deletions
--- a/pom.xml
+++ b/pom.xml
@@ -27,6 +27,7 @@
 		<module>document-readers/pdf-reader</module>
 		<module>document-readers/tika-reader</module>

+		<module>vector-stores/spring-ai-azure-cosmos-db-store</module>
 		<module>vector-stores/spring-ai-azure-store</module>
 		<module>vector-stores/spring-ai-cassandra-store</module>
 		<module>vector-stores/spring-ai-chroma-store</module>
--- 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"),
--- 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
@@ -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<Document> 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]
+----
+<dependency>
+    <groupId>org.springframework.ai</groupId>
+    <artifactId>spring-ai-azure-cosmos-db-store-spring-boot-starter</artifactId>
+</dependency>
+----
+
+== 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<String, Object> metadata1 = new HashMap<>();
+metadata1.put("country", "UK");
+metadata1.put("year", 2021);
+metadata1.put("city", "London");
+
+Map<String, Object> 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<Document> 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<Document> 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]
+----
+<dependency>
+    <groupId>org.springframework.ai</groupId>
+    <artifactId>spring-ai-azure-cosmos-db-store</artifactId>
+</dependency>
+----
--- a/spring-ai-spring-boot-autoconfigure/pom.xml
+++ b/spring-ai-spring-boot-autoconfigure/pom.xml
@@ -387,6 +387,14 @@
 			<optional>true</optional>
 		</dependency>

+		<!-- Azure Cosmos DB vector store -->
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-azure-cosmos-db-store</artifactId>
+			<version>${project.parent.version}</version>
+			<optional>true</optional>
+		</dependency>
+
 		<!-- test dependencies -->

 		<dependency>
--- 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
@@ -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<VectorStoreObservationConvention> 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);
+	}
+
+}
--- 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
@@ -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;
+	}
+
+}
--- 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
@@ -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<Document> 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<Document> 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<String, Object> metadata1;
+		metadata1 = new HashMap<>();
+		metadata1.put("country", "UK");
+		metadata1.put("year", 2021);
+		metadata1.put("city", "London");
+
+		Map<String, Object> metadata2;
+		metadata2 = new HashMap<>();
+		metadata2.put("country", "NL");
+		metadata2.put("year", 2022);
+		metadata2.put("city", "Amsterdam");
+
+		Map<String, Object> metadata3;
+		metadata3 = new HashMap<>();
+		metadata3.put("country", "US");
+		metadata3.put("year", 2019);
+		metadata3.put("city", "Sofia");
+
+		Map<String, Object> 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<Document> 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<Document> 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<Document> 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<Document> 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();
+		}
+
+	}
+
+}
--- 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
@@ -0,0 +1,42 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <groupId>org.springframework.ai</groupId>
+        <artifactId>spring-ai</artifactId>
+        <version>1.0.0-SNAPSHOT</version>
+        <relativePath>../../pom.xml</relativePath>
+    </parent>
+    <artifactId>spring-ai-azure-cosmos-db-store-spring-boot-starter</artifactId>
+    <packaging>jar</packaging>
+    <name>Spring AI Starter - Azure Cosmos DB Vector Store</name>
+    <description>Spring AI Azure Cosmos DB Vector Store Auto Configuration</description>
+    <url>https://github.com/spring-projects/spring-ai</url>
+
+    <scm>
+        <url>https://github.com/spring-projects/spring-ai</url>
+        <connection>git://github.com/spring-projects/spring-ai.git</connection>
+        <developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
+    </scm>
+
+    <dependencies>
+
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-azure-cosmos-db-store</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+    </dependencies>
+
+</project>
--- a/vector-stores/spring-ai-azure-cosmos-db-store/README.md
+++ b/vector-stores/spring-ai-azure-cosmos-db-store/README.md
@@ -0,0 +1 @@
+[Azure Cosmos DB Vector Store Documentation]()
--- a/vector-stores/spring-ai-azure-cosmos-db-store/pom.xml
+++ b/vector-stores/spring-ai-azure-cosmos-db-store/pom.xml
@@ -0,0 +1,81 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+	xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+	<modelVersion>4.0.0</modelVersion>
+	<parent>
+		<groupId>org.springframework.ai</groupId>
+		<artifactId>spring-ai</artifactId>
+		<version>1.0.0-SNAPSHOT</version>
+		<relativePath>../../pom.xml</relativePath>
+	</parent>
+	<artifactId>spring-ai-azure-cosmos-db-store</artifactId>
+	<packaging>jar</packaging>
+	<name>Spring AI Vector Store – Azure Cosmos DB</name>
+	<description>Spring AI Vector Store for Azure Cosmos DB</description>
+    <url>https://github.com/spring-projects/spring-ai</url>
+
+	<scm>
+		<url>https://github.com/spring-projects/spring-ai</url>
+		<connection>git://github.com/spring-projects/spring-ai.git</connection>
+		<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
+	</scm>
+
+	<properties>
+        <maven.compiler.target>17</maven.compiler.target>
+        <maven.compiler.source>17</maven.compiler.source>
+	</properties>
+
+	<dependencies>
+		<dependency>
+			<groupId>com.azure</groupId>
+			<artifactId>azure-spring-data-cosmos</artifactId>
+			<version>LATEST</version>
+		</dependency>
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-core</artifactId>
+			<version>${project.parent.version}</version>
+		</dependency>
+
+		<!-- TESTING -->
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-transformers</artifactId>
+			<version>${project.parent.version}</version>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.springframework.ai</groupId>
+			<artifactId>spring-ai-test</artifactId>
+			<version>${project.parent.version}</version>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-starter-test</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.testcontainers</groupId>
+			<artifactId>junit-jupiter</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.testcontainers</groupId>
+			<artifactId>azure</artifactId>
+			<version>1.20.1</version>
+		</dependency>
+
+		<dependency>
+			<groupId>io.micrometer</groupId>
+			<artifactId>micrometer-observation-test</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+	</dependencies>
+
+</project>
--- 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
@@ -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<String, String> metadataFields;
+
+	public CosmosDBFilterExpressionConverter(Collection<String> 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<String> 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<String> 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());
+		}
+	}
+
+}
--- 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
@@ -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<String> pathsfromCommaSeparatedList = new ArrayList<String>();
+		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<Document> 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<ImmutablePair<String, CosmosItemOperation>> 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<CosmosItemOperation> 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<Boolean> doDelete(List<String> idList) {
+		try {
+			// Convert the list of IDs into bulk delete operations
+			List<CosmosItemOperation> 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<Document> similaritySearch(String query) {
+		return similaritySearch(SearchRequest.query(query));
+	}
+
+	@Override
+	public List<Document> 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<Float> 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<SqlParameter> 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<JsonNode> pagedFlux = container.queryItems(sqlQuerySpec, options, JsonNode.class);
+
+		logger.info("Executing similarity search query: {}", query);
+		try {
+			// Collect documents from the paged flux
+			List<JsonNode> documents = pagedFlux.byPage()
+				.flatMap(page -> Flux.fromIterable(page.getResults()))
+				.collectList()
+				.block();
+			// Convert JsonNode to Document
+			List<Document> 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");
+	}
+
+}
--- 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
@@ -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<String> 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<String> 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;
+	}
+
+}
--- 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
@@ -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<Document> 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<Document> 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<String, Object> metadata1;
+		metadata1 = new HashMap<>();
+		metadata1.put("country", "UK");
+		metadata1.put("year", 2021);
+		metadata1.put("city", "London");
+
+		Map<String, Object> metadata2;
+		metadata2 = new HashMap<>();
+		metadata2.put("country", "NL");
+		metadata2.put("year", 2022);
+		metadata2.put("city", "Amsterdam");
+
+		Map<String, Object> metadata3;
+		metadata3 = new HashMap<>();
+		metadata3.put("country", "US");
+		metadata3.put("year", 2019);
+		metadata3.put("city", "Sofia");
+
+		Map<String, Object> 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<Document> 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<Document> 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<Document> 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<Document> 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() {
+			};
+		}
+
+	}
+
+}
--- 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
@@ -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}
				`@@ -0,0 +1 @@`
				`[Azure Cosmos DB Vector Store Documentation]()`