momentoのベクターストア機能（Vector Index）を触ってみた

この記事は「Momento Advent Calendar 2023」10日目の記事です。

LLMを利用したアプリケーションやRAG（Retrieval Augmented Generation）アプリを作ってみたいと思っているため、ここ最近さまざまなサービスのベクターストア機能を触っています。今回はMomentoでインデックスの作成とAPIからのアクセス（Node）を試してみました。

GUIからインデックスを作る

ベクターストアのインデックス作成は、momentoのコンソール画面からも作れます。この場合は、[Create an Index]を選びましょう。ベクターストアでは、DBのデータベースに該当するリソースのことを「インデックス」とよびます。

インデックスの管理画面を見た感じ、複数のインデックスが作れそうなUIになっています。[インデックスの作成]を押すと作成ウィザードに移動します。

ベクターストアを初めて作るときのハードルが、「次元の数」でしょう。これは保存するベクトルデータの次元（Dimension）を指していて、OpenAIやClaudeなどでEmbeddingを生成する場合は、これらのモデルの次元と数字を合わせる必要があります。たぶん。

手軽にベクターストアの機能だけを試したい場合は、2とかでよいでしょう。OpenAIを使う場合は、1536を入れましょう。以前Pineconeでインデックスを作った際に書いた記事で参考記事などを紹介しましたので、こちらもぜひご覧ください。

MomentoではAWSかGCPどちらにリソースを作るかが選べることがあります。ですがベクターストアについては、2023/12時点ではAWS固定の様子です。Opensearch Serviceでも内部で利用しているのでしょうか・・・？ベクターストアを提供するだけならオーバースペックなので、シンプルにアプリをホストするインフラの構成上の問題だとは思いますが。

作成に必要な情報はここまでです。成功すれば、インデックス一覧ページに追加されます。

SDKからインデックスをつくる方法

SDKからもインデックスの作成や管理ができます。

APIキーを取得する

SDKから操作するには、APIキーが必要です。こちらもコンソールから生成できますが、2023年時点ではSupser user keyしかベクターストアにアクセスできる権限を付与できない様子です。

トークンを生成すると、APIキーやトークンなどが取得できます。有効期限をつけている場合は、リフレッシュトークンを使って更新する流れになる・・・のでしょうか。その場合は、AWSのSecrets Managerのようなキーのローテーションを自動化できる仕組みがあると良さそうです。

環境変数にAPIキーを登録します。.envに先ほど取得したAPIキーを入れておきましょう。

MOMENT_API_KEY=eyJlbmRwb2ludCxxxxx

SDKのインストールと、クライアントのセットアップ

SDKはさまざまな言語で提供されています。今回はNodeを利用するので、npmからインストールしましょう。

npm i @gomomento/sdk

APIを呼び出すためのクライアントは、次のように初期化します。

import { PreviewVectorIndexClient, Configurations, CredentialProvider} from '@gomomento/sdk';


    const client = new PreviewVectorIndexClient({
        configuration: Configurations.Laptop.latest(),
        credentialProvider: CredentialProvider.fromString({
            apiKey: process.env.MOMENT_API_KEY as string
        }),
    });

CredentialProvider.fromfromStringに環境変数の値をそのまま入れています。

SDKからインデックスを作成する

ベクターストアのインデックス作成は、次のような処理で行います。インデックス名を第一引数に、第二引数には次元の数（Dimension）を入れましょう。

await client.createIndex('test-index', 2);

成功すると、is_successを含むJSONが返ってきます。

 {
    "is_success": true
  }

SDKからベクターストアのインデックスを操作する

インデックスが作れたので、操作してみましょう

データの投入と更新は、upserItemBatch

データの投入と更新は同じメソッドで行えます。

      await client.upsertItemBatch('test-index', [
        {
          id: 'example_item_1',
          vector: [1.0, 2.0],
          metadata: {key1: 'value1'},
        },
        {
          id: 'example_item_2',
          vector: [3.0, 4.0],
          metadata: {key2: 'value2'},
        },
      ])

vectorに渡す配列の長さは、インデックス作成時に指定したdimensionの値と同じになるようにしましょう。もし上のコードを、OpenAI向けに1536で設定したインデックスにリクエストすると、エラーが発生します。

"3 INVALID_ARGUMENT: invalid parameter: vector, vector dimension has to match the index's dimension"

こちらも成功した場合のレスポンスはis_successで判断できます。

{
    "is_success": true
  }

検索はsearchで実施する

データの検索はsearchで行います。ベクトルの類似性で検索するため、検索クエリとして渡す値はベクトルデータになることに注意しましょう。取得件数などを第三引数のオプションで設定できます。

 await client.search('test-index', [1.0, 2.0], {
        topK: 3,
        metadataFields: ALL_VECTOR_METADATA,
      })

取得結果はこのようになります。scoresでどれくらい類似しているかを見ることができます。そのため「類似性はあまり高くないよ」とUIにラベルをつけたりするような実装も可能です。

{
    "is_success": true,
    "_hits": [
      {
        "id": "example_item_1",
        "score": 1,
        "metadata": {
          "key1": "value1"
        }
      },
      {
        "id": "example_item_2",
        "score": 0.9838699102401733,
        "metadata": {
          "key2": "value2"
        }
      }
    ]
  }

ベクトルデータも取得したい場合は、searchAndFetchVectors

検索のレスポンスデータそれぞれのベクトルデータも取得する場合は、searchAndFetchVectorsを使いましょう。

 await client.searchAndFetchVectors('test-index', [1.0, 2.0], {
        topK: 3,
        metadataFields: ALL_VECTOR_METADATA,
      })

こちらは、取得したそれぞれのアイテムが持つベクトルデータも取得できます。

{
    "is_success": true,
    "_hits": [
      {
        "id": "example_item_1",
        "score": 1,
        "vector": [
          1,
          2
        ],
        "metadata": {
          "key1": "value1"
        }
      },
      {
        "id": "example_item_2",
        "score": 0.9838699102401733,
        "vector": [
          3,
          4
        ],
        "metadata": {
          "key2": "value2"
        }
      }
    ]
  }

「検索結果と関連性の高いデータを見たい」のようなすこし踏み込んだレコメンドや、「別のサービスにベクトルデータをまとめて渡したい場合」といったバッチ方面にも使えそうです。

おわりに

momentoのベクターストアを軽く触ってみました。とてもシンプルな仕組みで使えるのと、APIキーやトークンを使った認証が細かく行える点がなかなか良い感じに見えますので、AWS上にシステムを作る際に、ベクトルデータだけmomentoに入れるなどの使い方もできそうです。あとは・・・Amazon Bedrock Knowledge Baseと連携できれば最高だなぁというところでしょうか。