This version is still in development and is not considered stable yet. For the latest snapshot version, please use Spring AI 1.0.0-SNAPSHOT!

Azure Cosmos DB

This section 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. 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:

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"));
		this.vectorStore.add(List.of(document1, document2));
        List<Document> results = this.vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1));

        log.info("Search results: {}", results);

        // Remove the documents from the vector store
		this.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:

<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.vectorStoreThroughput

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.

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", this.metadata1);
Document document2 = new Document("2", "A document about the Netherlands", this.metadata2);

vectorStore.add(List.of(document1, document2));

FilterExpressionBuilder builder = new FilterExpressionBuilder();
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("The World")
    .withTopK(10)
    .withFilterExpression((this.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:

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"));
		this.vectorStore.add(List.of(document1, document2));

        List<Document> results = this.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.setVectorStoreThroughput(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, this.embeddingModel);
    }

    @Bean
    public EmbeddingModel embeddingModel() {
        return new TransformersEmbeddingModel();
    }
}

Manual Dependency Setup

Add the following dependency in your Maven project:

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-azure-cosmos-db-store</artifactId>
</dependency>