Vertex AI でベクトル類似度検索を行う

Firebase Data Connect のベクトル類似度検索について説明します。これは、Google Vertex AI と統合された Firebase のセマンティック検索の実装です。

この機能の中心となるのはベクトル エンベディングです。これは、テキストやメディアのセマンティックな意味を表す浮動小数点数の配列です。入力ベクトル エンベディングを使用して最近傍検索を実行すると、意味的に類似するすべてのコンテンツを見つけることができます。Data Connect は、この機能に PostgreSQL の pgvector 拡張機能を使用します。

この強力なセマンティック検索は、レコメンデーション エンジンや検索エンジンなどのユースケースを推進できます。また、生成 AI フローの検索拡張生成でも重要な要素です。詳細については、Vertex AI のドキュメントをご覧ください。

Data Connect の組み込みサポートを使用して、Vertex AI の Embeddings API を使用してベクトル エンベディングを自動的に生成することも、その API を使用して手動で生成することもできます。

前提条件

  • プロジェクト用に Data Connect を設定します。

  • Vertex AI API を有効にします。

設定

ローカル開発フロー(ウェブ、Kotlin Android、iOS デベロッパーの場合)または IDX フロー(ウェブ デベロッパーの場合)を選択できます。開発には、ローカル データベースまたは本番環境の Data Connect プロジェクトとその Cloud SQL for PostgreSQL インスタンスを使用できます。

以下の手順は、クイックスタート ガイドに沿って Data Connect プロジェクトを作成していることを前提としています。

ローカル PostgreSQL と統合する

  1. ローカル PostgreSQL インスタンスを設定します。
  2. 自分自身に Vertex AI ユーザーの IAM ロール付与します。
  3. 環境で Google Cloud アプリケーションのデフォルト認証情報を設定します。
  4. ローカル PostgreSQL インスタンスに pgvector 拡張機能をインストールします。
  5. pgvector リポジトリの手順に沿って、CREATE EXTENSION vector を使用して拡張機能を有効にします。

IDX と統合する

  1. Data Connect テンプレートを使用して IDX ワークスペースを設定します。
  2. 自分に Vertex AI ユーザーの IAM ロール付与します。
  3. pgvector リポジトリの手順に沿って、CREATE EXTENSION vector を使用して拡張機能を有効にします。

スキーマを設計する

ベクトル検索を実行するには、スキーマに Vector タイプの新しいフィールドを追加します。たとえば、映画の説明を使用してセマンティック検索を行う場合は、映画の説明に関連付けられたベクトル エンベディングを保持するフィールドを追加します。このスキーマでは、description フィールドのベクトル エンベディングを保存するために descriptionEmbedding が追加されています。

type Movie @table {
 id: ID! @col(name: "movie_id") @default(id: ID! @col(name: "movie_id") @default(expr: "uuidV4()")
 title: String!
 description: String
 descriptionEmbedding: Vector! @col(size:768)
 // ...
}

エンベディングの生成と取得

Data Connect は、_embed サーバー値でベクトル エンベディングを統合的にサポートします。これにより、Data Connect は Vertex AI の Embedding API を内部で呼び出してベクトル エンベディングを生成するように指示されます。_embed サーバー値は、ミューテーションとクエリの両方で使用できます。

ミューテーション

Data Connect を使用してエンベディングを生成して保存する

ベクトル検索アプリでは、データベースにレコードを追加するとすぐにエンベディングの生成をリクエストすることをおすすめします。次の createMovie ミューテーションは、映画レコードを Movie テーブルに追加し、指定されたエンベディング model を含む映画の説明を渡します。

mutation createMovie($movieData: Movie_Data! @pick(fields: ["title", "genre", "description"])) {
  movie_insert(data: {
    ...movieData,
    descriptionEmbedding_embed: {model: "textembedding-gecko@003", text: $movieData.description}
  })
}

映画の説明や埋め込みを更新したい場合もあります。

mutation updateDescription($id: String!, $description: String!) {
  movie_update(id: $id, data: {
    description: $description,
    descriptionEmbedding_embed: {model: "textembedding-gecko@003", text: $description}
  })
}

クライアントから後者のミューテーションを呼び出すには:

import { updateMovieDescriptionWithEmbed } from 'lib/dataconnect-sdk/';

await updateMovieDescriptionWithEmbed({ id: movieId, description: description});

// Use the response

クエリ

次のようなクエリを使用してベクトル エンベディングを取得します。クエリによって返される descriptionEmbedding は浮動小数点数のアレイであり、通常は人間が読み取れません。したがって、Data Connect で生成された SDK は、直接返すことをサポートしていません。

返されたベクトル エンベディングを使用して、次のセクションで説明するように類似度検索を行うことができます。

query getMovieDescription($id: String!) @auth(is: PUBLIC) {
 movie(id: $id)
   id
   description
   descriptionEmbedding
}

類似性検索を行う

これで類似性検索を実行できます。

Vector フィールドごとに、Data Connect は類似性検索を実装する GraphQL 関数を生成します。この生成された関数の名前は ${pluralType}_${vectorFieldName}_similarity です。以下の例とリファレンス リストに示されているように、いくつかのパラメータがサポートされています。

類似性検索を呼び出す GraphQL 関数を定義できます。前述のように、_embed サーバー値は、Vertex AI の Embedding API を使用してベクトル エンベディングを生成するように Data Connect に指示します。この場合は、映画の説明エンベディングとの比較に使用される検索文字列のエンベディングを作成します。

この例では、類似性検索によって、入力クエリに意味的に最も近い説明を持つ映画が最大 5 本返されます。結果セットは、距離の昇順(最も近い場所から最も遠い場所)で並べ替えられます。

query searchMovieDescriptionUsingL2Similarity ($query: String!) @auth(level: PUBLIC) {
    movies_descriptionEmbedding_similarity(
      compare_embed: {model: "textembedding-gecko@003", text: $query},
      method: L2,
      within: 2,
      where: {content: {ne: "No info available for this movie."}}, limit: 5)
      {
        id
        title
        description
      }
  }

類似度クエリを呼び出す

クライアント コードから類似性検索を呼び出すには:

import { searchMovieDescriptionUsingL2similarity} from 'lib/dataconnect-sdk';

const response = await searchMovieDescriptionUsingL2similarity({ query });

// Use the response

カスタム エンベディングを使用する

Data Connect を使用すると、_embed サーバー値を使用してエンベディングを生成するのではなく、エンベディングを Vector として直接操作することもできます。

カスタム エンベディングを保存する

Vertex Embeddings API を使用して、一致するモデルを指定し、正しいディメンションのエンベディング結果をリクエストします。

次に、返された浮動小数点数配列を Vector にキャストして、ストレージの更新オペレーションに渡します。

mutation updateDescription($id: String!, $description: String!, $descriptionEmbedding: Vector!) {
  movie_update(id: $id, data: {
    // title, genre...
    description: $description,
    descriptionEmbedding: $descriptionEmbedding
  })
}

カスタム エンベディングを使用して類似性検索を行う

同じオペレーションを実行して、検索語句のエンベディングを取得し、Vectors にキャストします。

次に、_similarity クエリを呼び出して各検索を実行します。

query searchMovieDescriptionUsingL2Similarity($compare: Vector!, $within: Float, $excludesContent: String, $limit: Int) @auth(level: PUBLIC) {
    movies_descriptionEmbedding_similarity(
      compare: $compare,
      method: L2,
      within: $within,
      where: {content: {ne: $excludesContent}}, limit: $limit)
      {
        id
        title
        description
      }
  }

本番環境にデプロイ

スキーマとコネクタをデプロイする

一般的な Data Connect イテレーションの最後のステップは、アセットを本番環境にデプロイすることです。

firebase deploy コマンドを使用して Vector 型を含むスキーマを CloudSQL にデプロイすると、Firebase CLI は CloudSQL インスタンスで Vertex AI ベースのエンベディング生成を有効にするために必要な手順を実行します。

firebase deploy --only dataconnect

CloudSQL インスタンスで埋め込みサポートを手動で有効にする場合は、こちらの手順に沿って操作します。CLI エラーが発生した場合は、こちらの手順に沿って操作します。

ベクトル検索の構文

スキーマ拡張機能

Data ConnectVector データ型は、pgvector 拡張機能で定義されている PostgreSQL の vector 型にマッピングされます。pgvector の vector 型は、単精度浮動小数点数の配列として PostgreSQL に保存されます。

Data Connect では、Vector 型は JSON number の配列として表されます。入力は、float32 値の配列に変換されます。型変換に失敗すると、エラーが発生します。

@col ディレクティブの size パラメータを使用して、ベクトルのディメンションを設定します。

type Question @table {
    text: String!
    category: String!
    textEmbedding: Vector! @col(size: 768)
}

sizeVector 型でのみサポートされます。類似度検索などの Vector オペレーションでは、すべての Vector が同じ次元数を持つ必要があります。

directive @col(
  # … existing args
  """
  Defines a fixed column size for certain scalar types.

  - For Vector, size is required.
  - For all other types, size is currently unsupported and hence supplying it will result in a schema error.
  """
  size: Int
) on FIELD_DEFINITION

クエリとミューテーションの _embed サーバー値

_embed

このサーバー値は、Vertex AI の Embedding API を使用してエンベディングを生成して保存するように Data Connect サービスに指示します。このサーバー値は、クエリとミューテーションの両方で使用できます。

類似性検索のパラメータ

method: COSINE|INNER_PRODUCT|L2

近傍を検索するために使用する距離関数。現在サポートされているアルゴリズムは、pgvector 検索アルゴリズムのサブセットです。

within: float

最近傍探索が実行される距離の制約。

where: FDC filter condition

スキーマ、クエリ、ミューテーションのガイドをご覧ください。

limit: int

返す結果の数。