Azure Cosmos DB

このセクションでは、ドキュメントの埋め込みを保存し、類似性検索を実行するための CosmosDBVectorStore のセットアップについて説明します。

Azure Cosmos DB とは何ですか ?

Azure Cosmos DB (英語) は、ミッションクリティカルなアプリケーション向けに設計された、Microsoft のグローバル分散型クラウドネイティブデータベースサービスです。高可用性、低レイテンシ、最新のアプリケーション要求を満たすための水平方向のスケーリング機能を提供します。グローバル分散、きめ細かいマルチテナント、水平方向のスケーラビリティを中核としてゼロから構築されました。これは Azure の基盤サービスであり、Teams、Skype、Xbox Live、Office 365、Bing、Azure Active Directory、Azure Portal、Microsoft Store など、Microsoft のほとんどのミッションクリティカルなアプリケーションでグローバル規模で使用されています。また、OpenAI の ChatGPT や、弾力的なスケール、ターンキーのグローバル分散、地球全体での低レイテンシと高可用性を必要とするその他のミッションクリティカルな AI アプリケーションなど、何千もの外部顧客によって使用されています。

DiskANN とは何ですか ?

DiskANN (ディスクベースの近似最近傍検索) は、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.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();
    }
}

自動構成

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

自動構成なしで Azure Cosmos DB ベクトルストアを設定する

次のコードは、自動構成に依存せずに CosmosDBVectorStore を設定する方法を示しています。

@Bean
public VectorStore vectorStore(ObservationRegistry observationRegistry) {
    // Create the Cosmos DB client
    CosmosAsyncClient cosmosClient = new CosmosClientBuilder()
            .endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))
            .key(System.getenv("COSMOSDB_AI_KEY"))
            .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/ 秒)

  • 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
}

ネイティブクライアントを使用すると、VectorStore インターフェースでは公開されない可能性のある Azure Cosmos DB 固有の機能と操作にアクセスできます。