Spring AI ​Azure Cosmos DB

什么是 Azure Cosmos DB？

Azure Cosmos DB 是微软推出的全球分布式云原生数据库服务，专为关键业务应用设计。它具备高可用性、低延迟的特性，支持水平扩展以满足现代应用的需求。该服务从底层构建时就以全球分布式部署、细粒度多租户和水平扩展为核心。它是 Azure 的基础服务之一，被微软大多数全球级关键业务应用使用，包括 Teams、Skype、Xbox Live、Office 365、必应、Azure 活动目录、Azure 门户、微软应用商店等。同时，数千家外部客户也在使用该服务，其中包括 OpenAI 的 ChatGPT 及其他需要弹性扩展、一键式全球部署、全球低延迟和高可用性的关键人工智能应用。

什么是 DiskANN？

DiskANN（基于磁盘的近似最近邻搜索）是 Azure Cosmos DB 中用于提升向量搜索性能的创新技术。它通过为存储在 Cosmos DB 中的嵌入向量建立索引，实现对高维数据高效、可扩展的相似度搜索。

DiskANN 具备以下优势：

• 高效性：与传统方法相比，DiskANN 采用基于磁盘的结构，大幅缩短了查找最近邻的时间。

• 可扩展性：能够处理超出内存容量的大型数据集，适用于机器学习、人工智能驱动解决方案等各类应用。

• 低延迟：DiskANN 最大限度降低搜索操作的延迟，确保应用即使在数据量庞大的情况下也能快速获取结果。

在 Spring AI 集成 Azure Cosmos DB 的场景中，向量搜索会创建并使用 DiskANN 索引，保障相似度查询的最优性能。

通过自动配置设置 Azure Cosmos DB 向量存储

以下代码演示了如何通过自动配置方式设置 CosmosDBVectorStore：

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));
        Listresults = this.vectorStore.similaritySearch(SearchRequest.builder().query("Sample content").topK(1).build());

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

        // 从向量存储中删除文档
		this.vectorStore.delete(List.of(document1.getId(), document2.getId()));
    }

    @Bean
    public ObservationRegistry observationRegistry() {
        return ObservationRegistry.create();
    }
}

自动配置

Spring AI 自动配置、启动器模块的构件名称发生了重大变更。更多信息请参考升级说明。

将以下依赖添加到你的 Maven 项目中：

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

配置属性

Cosmos DB 向量存储支持以下配置属性：

带过滤器的复杂搜索

你可以在 Cosmos DB 向量存储中使用过滤器执行更复杂的搜索。以下示例演示了如何在搜索查询中使用过滤器：

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.builder().query("The World")
    .topK(10)
    .filterExpression((this.builder.in("country", "UK", "NL")).build()).build());

不使用自动配置设置 Azure Cosmos DB 向量存储

以下代码演示了如何不依赖自动配置设置 CosmosDBVectorStore。建议使用 DefaultAzureCredential 进行 Azure Cosmos DB 的身份验证。

@Bean
public VectorStore vectorStore(ObservationRegistry observationRegistry) {
    // 创建 Cosmos DB 客户端
    CosmosAsyncClient cosmosClient = new CosmosClientBuilder()
            .endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))
            .credential(new DefaultAzureCredentialBuilder().build())
            .userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
            .gatewayMode()
            .buildAsyncClient();

    // 创建并配置向量存储
    return CosmosDBVectorStore.builder(cosmosClient, embeddingModel)
            .databaseName("test-database")
            .containerName("test-container")
            // 配置用于过滤的元数据字段
            .metadataFields(List.of("country", "year", "city"))
            // 设置分区键路径（可选）
            .partitionKeyPath("/id")
            // 配置性能参数
            .vectorStoreThroughput(1000)
            .vectorDimensions(1536)  // 与你的嵌入模型维度保持一致
            // 添加自定义批处理策略（可选）
            .batchingStrategy(new TokenCountBatchingStrategy())
            // 添加用于指标监控的观测注册表
            .observationRegistry(observationRegistry)
            .build();
}

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

该配置展示了所有可用的构建器选项：

• databaseName：Cosmos DB 数据库名称

• containerName：数据库内的容器名称

• partitionKeyPath：分区键路径（如 "/id"）

• metadataFields：用于过滤的元数据字段列表

• vectorStoreThroughput：向量存储容器的吞吐量（RU/s）

• vectorDimensions：向量维度（需与嵌入模型匹配）

• batchingStrategy：文档操作批处理策略（可选）

手动依赖配置

在你的 Maven 项目中添加以下依赖：

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

访问原生客户端

Azure Cosmos DB 向量存储实现通过 getNativeClient() 方法提供对底层原生 Azure Cosmos DB 客户端（CosmosClient）的访问：

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

if (nativeClient.isPresent()) {
    CosmosClient client = nativeClient.get();
    // 使用原生客户端执行 Azure Cosmos DB 专属操作
}

原生客户端允许你访问 VectorStore 接口未暴露的 Azure Cosmos DB 专属特性和操作。