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 を設定します。

注: 設定フローを進める際は、 Data Connect の Vertex AI インテグレーションが特定の Cloud SQL for PostgreSQL のロケーションでのみサポートされていることに注意してください。ロケーションのリストをご覧ください。
Vertex AI API を有効にします。

注: 別の Google Cloud プロジェクトを作成したり、 Google Cloud CLI をインストールしたりする必要はありません。

設定

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

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

ローカル PostgreSQL と統合する

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

IDX と統合する

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

スキーマを設計する

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

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($title: String!, $description: String!) {
  movie_insert(data: {
    title: $title,
    description: $description,
    descriptionEmbedding_embed: {model: "textembedding-gecko@003", text: $description}
  })
}

場合によっては、映画の説明とエンベディングを更新することが必要になることがあります。

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

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

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

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

// Use the response

クエリ

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

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

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

類似性検索を実行する

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

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

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

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

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

類似性クエリを調整する

method や within などの検索パラメータのデフォルト値は、ほとんどのユースケースで適切に機能します。ただし、クエリが類似性の低い結果を返す場合や、含めるべき結果が欠落している場合は、これらのパラメータを調整してみてください。

within の適切な値を見つけるには、選択したフィールドに _metadata.distance を追加して、各結果がクエリベクトルからどの程度離れているかを確認します。返された distance 値に基づいて、within パラメータを設定できます。within の値よりも距離が小さい結果のみが含まれます。

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

method パラメータを設定して、さまざまな距離関数を試すこともできます。

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

メソッドによって距離の値が大きく異なることに注意してください。within を設定した場合は、method を変更した後でその値を再度調整する必要があります。

類似性クエリを呼び出す

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

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

const response = await searchMovieDescriptionUsingL2similarity({ query });

// Use the response

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

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

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

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 イテレーションの最後のステップは、アセットを本番環境にデプロイすることです。

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

firebase deploy --only dataconnect

CloudSQL インスタンスでエンベディングのサポートを手動で有効にする場合や、 CLI エラーが発生した場合は、こちらの手順に沿って操作してください。

ベクトル検索の構文

スキーマの拡張

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

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

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

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

size は Vector 型でのみサポートされています。類似性検索などの 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

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

類似性検索のパラメータ

method: COSINE|INNER_PRODUCT|L2

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

within: float

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

where: FDC filter condition

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

limit: int

返される結果の数。