Azure Cosmos DB

本節將引導您設定 CosmosDBVectorStore,以儲存文件嵌入並執行相似性搜尋。

什麼是 Azure Cosmos DB?

Azure Cosmos DB 是 Microsoft 的全球分散式雲端原生資料庫服務,專為任務關鍵型應用程式設計。它提供高可用性、低延遲以及水平擴展能力,以滿足現代應用程式的需求。它從頭開始建構,核心是全球分佈、細緻的多租戶和水平可擴展性。它是 Azure 中的基礎服務,被 Microsoft 大多數任務關鍵型應用程式在全球範圍內使用,包括 Teams、Skype、Xbox Live、Office 365、Bing、Azure Active Directory、Azure Portal、Microsoft Store 等等。它也被數千家外部客戶使用,包括 OpenAI 的 ChatGPT 和其他任務關鍵型 AI 應用程式,這些應用程式需要彈性擴展、開箱即用的全球分佈,以及遍及全球的低延遲和高可用性。

什麼是 DiskANN?

DiskANN(Disk-based Approximate Nearest Neighbor Search,基於磁碟的近似最近鄰搜尋)是 Azure Cosmos DB 中使用的一種創新技術,用於增強向量搜尋的效能。它透過索引儲存在 Cosmos DB 中的嵌入,實現跨高維度資料的高效且可擴展的相似性搜尋。

DiskANN 提供以下優點

  • 效率:透過利用基於磁碟的結構,DiskANN 顯著減少了尋找最近鄰所需的時間,與傳統方法相比。

  • 可擴展性:它可以處理超出記憶體容量的大型資料集,使其適用於各種應用程式,包括機器學習和 AI 驅動的解決方案。

  • 低延遲:DiskANN 最大程度地減少了搜尋操作期間的延遲,確保應用程式即使在大量資料的情況下也能快速檢索結果。

在 Spring AI for 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.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();
    }
}

自動設定

將以下相依性新增至您的 Maven 專案

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

設定屬性

以下設定屬性可用於 Cosmos DB 向量儲存體

屬性 描述

spring.ai.vectorstore.cosmosdb.databaseName

要使用的 Cosmos DB 資料庫名稱。

spring.ai.vectorstore.cosmosdb.containerName

要使用的 Cosmos DB 容器名稱。

spring.ai.vectorstore.cosmosdb.partitionKeyPath

分割區金鑰的路徑。

spring.ai.vectorstore.cosmosdb.metadataFields

以逗號分隔的中繼資料欄位清單。

spring.ai.vectorstore.cosmosdb.vectorStoreThroughput

向量儲存體的輸送量。

spring.ai.vectorstore.cosmosdb.vectorDimensions

向量的維度數量。

spring.ai.vectorstore.cosmosdb.endpoint

Cosmos DB 的端點。

spring.ai.vectorstore.cosmosdb.key

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

在沒有自動設定的情況下設定 Azure Cosmos DB 向量儲存體

以下程式碼示範如何在不依賴自動設定的情況下設定 CosmosDBVectorStore

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();
    }
}

手動設定相依性

在您的 Maven 專案中新增以下相依性

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