Azure Cosmos DB Vector Store

This guide walks you through setting up CosmosDBVectorStore to store document embeddings and perform similarity searches.

What is Azure 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.

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.

💡 Note: 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.

💡 Tip: For zero-config Spring Boot setup, see the Auto-Configuration page.

Dependency

<dependency>
    <groupId>com.azure.spring.ai</groupId>
    <artifactId>spring-ai-azure-cosmos-db-store</artifactId>
    <version>1.0.0-SNAPSHOT</version>
</dependency>

Example

DefaultAzureCredential is recommended for authentication.

import com.azure.cosmos.CosmosAsyncClient;
import com.azure.cosmos.CosmosClientBuilder;
import com.azure.identity.DefaultAzureCredentialBuilder;
import com.azure.spring.ai.vectorstore.cosmosdb.CosmosDBVectorStore;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class CosmosVectorStoreConfig {

    @Bean
    public VectorStore vectorStore(CosmosAsyncClient cosmosClient,
                                   EmbeddingModel embeddingModel) {
        return CosmosDBVectorStore.builder(cosmosClient, embeddingModel)
                .databaseName("my-database")
                .containerName("my-vectors")
                .metadataFields(List.of("country", "year", "city"))
                .partitionKeyPath("/id")
                .vectorStoreThroughput(1000)
                .vectorDimensions(1536)
                .batchingStrategy(new TokenCountBatchingStrategy())
                .build();
    }

    @Bean
    public CosmosAsyncClient cosmosClient() {
        return new CosmosClientBuilder()
                .endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))
                .credential(new DefaultAzureCredentialBuilder().build())
                .userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
                .gatewayMode()
                .buildAsyncClient();
    }
}

Builder Options

Option	Description
`databaseName`	The name of your Cosmos DB database
`containerName`	The name of your container within the database
`partitionKeyPath`	The path for the partition key (e.g., "/id")
`metadataFields`	List of metadata fields used for filtering
`vectorStoreThroughput`	The throughput (RU/s) for the vector store container
`vectorDimensions`	The number of dimensions for your vectors (should match your embedding model)
`batchingStrategy`	Strategy for batching document operations (optional)

Complex Searches with Filters

You can perform more complex searches using filters:

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;

Map<String, Object> metadata1 = Map.of("country", "UK", "year", 2021, "city", "London");
Map<String, Object> metadata2 = Map.of("country", "NL", "year", 2022, "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.builder()
        .query("The World")
        .topK(10)
        .filterExpression(builder.in("country", "UK", "NL").build())
        .build());

Accessing the Native Client

The vector store provides access to the underlying native Azure Cosmos DB client through getNativeClient():

import com.azure.spring.ai.vectorstore.cosmosdb.CosmosDBVectorStore;
import com.azure.cosmos.CosmosAsyncContainer;

CosmosDBVectorStore vectorStore = context.getBean(CosmosDBVectorStore.class);
Optional<CosmosAsyncContainer> nativeClient = vectorStore.getNativeClient();

if (nativeClient.isPresent()) {
    CosmosAsyncContainer container = nativeClient.get();
    // Use the native client for Cosmos DB-specific operations
}