Azure Cosmos DBのベクトル検索を使ってみよう

Microsoft Igniteにて、Cosmos DBにおけるベクトル検索の機能がGAされました。

devblogs.microsoft.com

プレビューの時から存在は知っていたのですが使ったことが無かったので、この機会に検証してみようと思います。

開発環境
Cosmos DB側の設定
ベクトル化する方法
コンソールアプリを実装する
終わりに
参考資料

開発環境

Consoleアプリ（.NET 8）
Microsoft.Azure.Cosmos 3.46.0

Cosmos DB側の設定

まず、ベクトル検索を使うためには、Azure Cosmos DB アカウントで機能を有効にする必要があります。

左側メニューの設定の中にある「機能」という項目をクリックして「Vector Search for NoSQL API」をクリックすると、ベクトル検索を有効にする画面が出てくるので「有効」を押します。10分ほどで機能が有効化されます。

次にコンテナを作ります。ベクトル検索をするためには、コンテナにベクトル検索のポリシーを設定する必要があるのですが、これはコンテナを作成する時のみ設定することが可能です。

Containers Policiesを見ると設定されていることが分かります。これによりvector1というパスが、ベクトル検索の対象として設定されました。

ベクトル化する方法

ベクトル検索をするためには文字列をベクトル化したデータが必要であり、自分で実装する必要があります。

そのため、Azure OpenAIでベクトル化するためのモデルをデプロイしておきます。今回は「text-embedding-3-large」を使用します。

ベクトル化については以下のドキュメントを参考にしました。

learn.microsoft.com

コンソールアプリを実装する

実際に、コードを使ってベクトル検索をしてみましょう。今回は以下のような実装にしました。

using Microsoft.Azure.Cosmos;
using System.Net.Http.Json;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;

string cosmosConnectionString = "[CosmosDBの接続文字列]";
string databaseId = "[CosmosDBのデータベース名]";
string containerId = "[CosmosDBのコンテナ名]";

string openAiApiKey = "[Azure OpenAIのAPIキー]";
Uri oaiEndpoint = new("[Azure OpenAIのEmbeddingsのエンドポイントURL]");


CosmosClient client = new(cosmosConnectionString);
Container container = client.GetContainer(databaseId, containerId);

//CosmosDBにアイテムを作成
await CreateCosmosItemAsync("Blazorは、Microsoftが提供するWebアプリケーション開発フレームワークです。");
await CreateCosmosItemAsync("自転車は、人力で走る乗り物です。");
await CreateCosmosItemAsync("うさぎは、生き物の一種です。");

//ベクトル検索を行う
await VectorSearchAsync("アニマル");


//ベクトル検索を行う
async Task VectorSearchAsync(string searchString)
{

    var embedding = await CreateEmbeddingAsync(searchString);
    var queryDef = new QueryDefinition(
        query: $"SELECT c.content, VectorDistance(c.vector1,@embedding) AS SimilarityScore FROM c ORDER BY VectorDistance(c.vector1,@embedding)"
        ).WithParameter("@embedding", embedding);
    using FeedIterator<Object> feed = container.GetItemQueryIterator<Object>(
        queryDefinition: queryDef
    );
    while (feed.HasMoreResults)
    {
        FeedResponse<Object> response = await feed.ReadNextAsync();
        foreach (Object item in response)
        {
            Console.WriteLine($"Found item:\t{item}");
        }
    }
}

//CosmosDBにアイテムを作成する
async Task CreateCosmosItemAsync(string testString)
{
    //文字列をベクトル化する
    var embeddingResult = await CreateEmbeddingAsync(testString);

    //CosmosDBにアイテムを作成する
    var result = await container.CreateItemAsync(new { id = Guid.NewGuid().ToString(), content = testString, vector1 = embeddingResult });
}

//Azure OpenAIにリクエストを送信して、埋め込みを取得する
async Task<List<float>> CreateEmbeddingAsync(string input)
{
    HttpClient httpClient = new();

    var json = new
    {
        input = new List<string>() { input }
    };
    using var content = new StringContent(JsonSerializer.Serialize(json), Encoding.UTF8, "application/json");
    HttpRequestMessage request = new(HttpMethod.Post, oaiEndpoint);
    request.Content = content;
    request.Headers.Add("api-key", openAiApiKey);
    var response = await httpClient.SendAsync(request);
    var responseContent = await response.Content.ReadFromJsonAsync<Response>();

    return responseContent?.Data[0].Embedding ?? new List<float>();
}


#nullable disable

/// <summary>
/// CosmosDBに保存するアイテム
/// </summary>
public class Item
{
    public string id { get; set; }
    public string content { get; set; }
    public List<float> vector1 { get; set; }
}


//以下はAzure OpenAIのレスポンスをデシリアライズするためのクラス

/// <summary>
/// Azure OpenAIからのレスポンス
/// </summary>
public class Response
{
    public string Object { get; set; }
    public List<Data> Data { get; set; }
    public string Model { get; set; }
    public Usage Usage { get; set; }
}

public class Data
{
    public string Object { get; set; }
    public int Index { get; set; }
    public List<float> Embedding { get; set; }
}

public class Usage
{
    [JsonPropertyName("prompt_tokens")]
    public int PromptTokens { get; set; }

    [JsonPropertyName("total_tokens")]
    public int TotalTokens { get; set; }
}

これを実行すると、CosmosDBのアイテム内の「content」が「アニマル」という単語の意味に近い順に表示されるはずです。

実行結果は以下のようになりました。想定通り「うさぎは～」が一番最初にヒットしました！

Found item:     {
  "content": "うさぎは、生き物の一種です。",
  "SimilarityScore": 0.3583362145311595
}
Found item:     {
  "content": "自転車は、人力で走る乗り物です。",
  "SimilarityScore": 0.18994718770566332
}
Found item:     {
  "content": "Blazorは、Microsoftが提供するWebアプリケーション開発フレームワークです。",
  "SimilarityScore": 0.13536014742330527
}