Azure Functions の Azure OpenAI Extension を使ってベクトル検索を構築する

Azure Functions では、Azure OpenAI 向けの拡張機能（Preview）が提供されています。
learn.microsoft.com

今回は、拡張機能の Embeddings binding を使って、ベクトル検索を構築してみました。
その他の拡張機能については、別記事を参照してください。
gooner.hateblo.jp
gooner.hateblo.jp
gooner.hateblo.jp

Embeddings binding とは

ベクトル検索を構築するための３種類の binding が提供されています。ベクトルストアは AI Search や Cosmos DB に対応しており、今回は AI Search を使います。

Embeddings input binding

Embeddings input binding を使うことで、テキストやファイルを Embedding することができます。
learn.microsoft.com

Embeddings store output binding

Embeddings store output binding を使うことで、テキストやファイルを Embedding してベクトルストアに書き込むことができます。
learn.microsoft.com

Semantic search input binding

Semantic search input binding を使うことで、ベクトルストアに対してセマンティック検索を実行できます。
learn.microsoft.com

前準備

Azure Functions を Isolated worker model で作成し、NGet ライブラリをインストールします。インストールする前に、Visual Studio などのテンプレートで作成されたプロジェクトのライブラリ群を最新版に更新しておくことをお勧めします。

$ dotnet add package Microsoft.Azure.Functions.Worker.Extensions.OpenAI --version 0.16.0-alpha
$ dotnet add package Microsoft.Azure.Functions.Worker.Extensions.OpenAI.AzureAISearch --version 0.3.0-alpha
$ dotnet add package Azure.Search.Documents --version 11.6.0

local.settings.json において、Azure OpenAI Service と AI Search のエンドポイントとキーを定義しておきます。

{
  "IsEncrypted": false,
  "Values": {
    "AzureWebJobsStorage": "UseDevelopmentStorage=true",
    "FUNCTIONS_WORKER_RUNTIME": "dotnet-isolated",
    "AZURE_OPENAI_ENDPOINT": "https://xxx.openai.azure.com/",
    "AZURE_OPENAI_KEY": "xxx",
    "CHAT_MODEL_DEPLOYMENT_NAME": "gpt-4o",
    "EMBEDDING_MODEL_DEPLOYMENT_NAME": "text-embedding-ada-002",
    "AISearchEndpoint": "https://xxx.search.windows.net",
    "SearchAPIKey": "xxx"
  }
}

host.json において、セマンティック検索の構成を定義しておきます。

{
  "extensions": {
    "openai": {
      "searchProvider": {
        "type": "azureAiSearch",
        "isSemanticSearchEnabled": true,
        "useSemanticCaptions": true,
        "vectorSearchDimensions": 1536,
        "searchAPIKeySetting": "SearchAPIKey"
      }
    }
  }
}

テキストを Embedding する

EmbeddingsInput のバインディングを使うことで、AOAI の Embedding Model を使ってテキストを Embedding できます。Embedding した結果は、EmbeddingsContext で受け取ることができます。

public class Embeddings
{
    private readonly ILogger<Embeddings> _logger;

    public Embeddings(ILogger<Embeddings> logger)
    {
        _logger = logger;
    }

    [Function(nameof(GenerateEmbeddings))]
    public async Task GenerateEmbeddings(
        [HttpTrigger(AuthorizationLevel.Function, "post", Route = "embeddings/generate")] HttpRequestData req,
        [EmbeddingsInput("{rawText}", InputType.RawText, Model = "%EMBEDDING_MODEL_DEPLOYMENT_NAME%")] EmbeddingsContext embeddings)
    {
        _logger.LogInformation("Embeddings input binding function processed a request.");

        var requestBody = await req.ReadFromJsonAsync<GenerateEmbeddingsRequest>();
        if (requestBody?.RawText == null)
        {
            throw new ArgumentException("Invalid request body. Make sure that you pass in {\"RawText\": value } as the request body.");
        }

        _logger.LogInformation(
            "Received {count} embedding(s) for input text containing {length} characters.",
            embeddings.Count,
            requestBody?.RawText?.Length);

        // TODO: Store the embeddings into a database or other storage.
    }

    internal class GenerateEmbeddingsRequest
    {
        public string? RawText { get; set; }
    }
}

なお、EmbeddingsInput のバインディングでは、InputType で FilePathや Url も指定できます。

テキストを Embedding してベクトルストアに書き込む

EmbeddingsStoreOutput のバインディングを使うことで、AOAI の Embedding Model を使ってテキストを Embedding した結果をベクトルストアに書き込むことができます。AI Search に openai-index という名前のインデックスを作成するように指定しています。

[Function(nameof(IngestEmbeddings))]
public async Task<EmbeddingsStoreOutputResponse> IngestEmbeddings(
    [HttpTrigger(AuthorizationLevel.Function, "post", Route = "embeddings/ingest")] HttpRequestData req)
{
    _logger.LogInformation("Embeddings store output binding function processed a request.");

    var requestBody = await req.ReadFromJsonAsync<IngestEmbeddingsRequest>();
    if (requestBody?.RawText == null || requestBody?.Title == null)
    {
        throw new ArgumentException("Invalid request body. Make sure that you pass in {\"RawText\": value, \"Title\": value} as the request body.");
    }

    IActionResult result = new OkObjectResult(new { status = HttpStatusCode.OK });
    return new EmbeddingsStoreOutputResponse
    {
        HttpResponse = result,
        SearchableDocument = new SearchableDocument(requestBody.Title)
    };
}

internal class IngestEmbeddingsRequest
{
    public string? RawText { get; set; }

    public string? Title { get; set; }
}

public class EmbeddingsStoreOutputResponse
{
    [EmbeddingsStoreOutput("{RawText}", InputType.RawText, "AISearchEndpoint", "openai-index", Model = "%EMBEDDING_MODEL_DEPLOYMENT_NAME%")]
    public required SearchableDocument SearchableDocument { get; init; }

    public IActionResult? HttpResponse { get; set; }
}

なお、EmbeddingsStoreOutput のバインディングでは、InputType で FilePathや Url も指定できます。

ベクトルストアに対してセマンティック検索を実行する

SemanticSearchInput のバインディングを使うことで、ベクトルストアに対してセマンティック検索を実行できます。AI Search の openai-index という名前のインデックスから検索します。検索結果は、SemanticSearchContext で受け取ることができます。

[Function(nameof(SemanticSearch))]
public IActionResult SemanticSearch(
    [HttpTrigger(AuthorizationLevel.Function, "post", Route = "embeddings/search")] SemanticSearchRequest unused,
    [SemanticSearchInput("AISearchEndpoint", "openai-index", Query = "{Prompt}", ChatModel = "%CHAT_MODEL_DEPLOYMENT_NAME%", EmbeddingsModel = "%EMBEDDING_MODEL_DEPLOYMENT_NAME%")] SemanticSearchContext result)
{
    _logger.LogInformation("Semantic search input binding function processed a request.");

    return new ContentResult { Content = result.Response, ContentType = "text/plain" };
}

なお、SemanticSearchInput のバインディングでは、検索オプションを指定できます。

動作確認

VS Code 拡張機能の REST Client を使って、ベクトル検索を実行してみます。

テキストを Embedding する

@function_app_HostAddress = http://localhost:7074

### Embeddings input binding
POST {{function_app_HostAddress}}/api/embeddings/generate HTTP/1.1
content-type: application/json

{
  "rawText": "暑い日に食べられている京菓子は？"
}

テキストを Embedding してベクトルストアに書き込む

### Embeddings store output binding
POST {{function_app_HostAddress}}/api/embeddings/ingest HTTP/1.1
content-type: application/json

< ./Data/request-body01.json

ベクトル化する JSON は、別ファイル（request-body01.json）に切り出しています。

{
  "rawText": "花びら餅は、京都の伝統的な和菓子で、美しい花びら模様の薄い皮と甘さ控えめのあんこが特徴です。この和菓子は季節ごとにバリエーションが楽しめ、茶道やお茶うけに使われ、京都のお土産として人気があります。見た目も美しく、食べるだけでなく、芸術的な価値も持っています。また、花びら餅は季節感あふれる和菓子で、特に春に楽しまれます。春の季節になると、桜の花びらを使った花びら餅が登場し、桜の風味を楽しむことができます。この季節になると、花見のお供としても愛され、桜の美しさと和の風情を同時に楽しむことができます。花びら餅は、京都の伝統を受け継ぐ和菓子として、その美しさと風味によって多くの人に愛されています。",
  "title": "花びら餅"
}

花びら餅のほかに、若鮎、栗饅頭、八つ橋のデータをベクトル化しました。

ベクトルストアに対してセマンティック検索を実行する

### Semantic search input binding
POST {{function_app_HostAddress}}/api/embeddings/search HTTP/1.1
content-type: application/json

{
  "prompt": "暑い日に食べられている京菓子は？"
}

HTTP/1.1 200 OK
Content-Length: 175
Content-Type: text/plain
Date: Sat, 07 Sep 2024 11:02:22 GMT
Server: Kestrel

暑い日に食べられている京菓子の一例として、「若鮎」があります。若鮎は初夏を涼しげに彩る京の季節菓子です。

Reference: 若鮎

ベクトル検索で「若鮎」が返却されたことが分かります。

まとめ

Azure Functions 拡張機能の Embeddings binding を使って、Azure AI Search によるベクトル検索を構築してみました。
AOAI の SDK を使った定型のコードが不要となり、シンプルな実装で Embedding Model を使ったベクトルストアへのデータ投入やセマンティック検索ができて便利だと感じました。

今回のソースコードは、こちらのリポジトリで公開しています。
github.com

ROMANCE DAWN for the new world

Microsoft Azure を中心とした技術情報を書いています。

Azure Functions の Azure OpenAI Extension を使ってベクトル検索を構築する

Embeddings binding とは

Embeddings input binding

Embeddings store output binding

Semantic search input binding

前準備

テキストを Embedding する

テキストを Embedding してベクトルストアに書き込む

ベクトルストアに対してセマンティック検索を実行する

動作確認

テキストを Embedding する

テキストを Embedding してベクトルストアに書き込む

ベクトルストアに対してセマンティック検索を実行する

まとめ