Azure Cosmos DB

Azure Cosmos DB

本节将引导您完成设置 CosmosDBVectorStore，以存储文档嵌入并执行相似性搜索。

什么是 Azure Cosmos DB？

Azure Cosmos DB 是微软全球分布的云原生数据库服务，专为关键任务应用设计。 它提供高可用性、低延迟以及水平扩展能力，以满足现代应用的需求。 其核心设计从一开始就融入了全球分布、细粒度多租户和水平可扩展性。 作为 Azure 的一项基础服务，它被微软大多数全球规模的关键任务应用所采用，包括 Teams、Skype、Xbox Live、Office 365、Bing、Azure Active Directory、Azure Portal、Microsoft Store 等众多产品。 此外，成千上万的外部客户也在使用它，如 OpenAI 用于 ChatGPT 及其他需要弹性扩展、即开即用的全球分布以及全球范围内低延迟和高可用性的关键任务 AI 应用。

什么是 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));
        List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder().query("Sample content").topK(1).build());

        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）

将以下依赖项添加到您的 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) {
    // Create the Cosmos DB client
    CosmosAsyncClient cosmosClient = new CosmosClientBuilder()
            .endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))
            .credential(new DefaultAzureCredentialBuilder().build())
            .userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
            .gatewayMode()
            .buildAsyncClient();

    // Create and configure the vector store
    return CosmosDBVectorStore.builder(cosmosClient, embeddingModel)
            .databaseName("test-database")
            .containerName("test-container")
            // Configure metadata fields for filtering
            .metadataFields(List.of("country", "year", "city"))
            // Set the partition key path (optional)
            .partitionKeyPath("/id")
            // Configure performance settings
            .vectorStoreThroughput(1000)
            .vectorDimensions(1536)  // Match your embedding model's dimensions
            // Add custom batching strategy (optional)
            .batchingStrategy(new TokenCountBatchingStrategy())
            // Add observation registry for metrics
            .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();
    // Use the native client for Azure Cosmos DB-specific operations
}

原生客户端让您能够访问 Azure Cosmos DB 特有的功能和操作，这些可能未通过 VectorStore 接口公开。