ついに登場！？Firestoreの全文検索（ベクトル検索）

今までサードパーティーを利用するしかなかった、
Firestoreでついに全文検索ができるようになりました！！

ベクトル検索による全文検索ですが、今回はその実施手順をご紹介します。

前提
大まかな手順
Google Cloud Functionの作成
Googleが提供するVertex AIの設定
ベクトル検索で使用するベクトル算出
ベクトル検索のためのインデックス作成
Firestoreでのベクトル検索の実施コード
最後に

前提

今回紹介する機能はプレビュー版です。
正式リリースに合わせ、下記のコードは実行できなくなる可能性があります。
参考）ベクトル埋め込みによる検索：https://firebase.google.com/docs/firestore/vector-search

今回使用するFunctionやVertex AIは使用料が発生します。
ご利用の際は事前に確認ください。
参考）Vertex AI：https://cloud.google.com/vertex-ai/generative-ai/pricing?hl=ja

大まかな手順

ベクトル検索については別途記事を作成予定ですが、
今回はベクトル検索に使用するベクトルの計算にGoogleが提供する「Vertex AI」を使います。
また、ベクトル検索は現在、Python又はJavaScript（Node.js）でしか行えないため、
今回はGoogle Cloud Functionを使って検索を行います。

そのため、検索を行うために以下の準備が必要です。

Google Cloud Functionの作成
Vertex AIの設定
ベクトルの取得
インデックスの作成
Firestoreでのベクトル検索の実施コード

それぞれ詳しく手順を説明します。

Google Cloud Functionの作成

今回作成したCloud Functionsの環境は以下です。

2nd gen（第2世代）
httpsトリガー
Python3.12

またランタイム環境変数に以下を設定しています。

　名前：GOOGLE_CLOUD_PROJECT
　値　：<プロジェクトID>（プロジェクト名ではないことに注意）

Vertex AIはCloud Runと紐づける必要があるため、GCF作成時にCloud Runが作成される、
2nd gen（第2世代）を使っています。

Googleが提供するVertex AIの設定

FunctionのCloud RunにVertex AIを紐づけます。
（参考公式：https://cloud.google.com/run/docs/integrate/vertex-ai?authuser=3&hl=ja）

参考までにここにも手順を載せておきます。

右側にある「Powered by Cloud Run」の下のリンクをクリックし、Cloud Runに移動
「統合」タブをクリック
「インテグレーションを追加」をクリック
「Vertex AI – 生成 AI」をクリックし、任意の名前を設定して「submit」をクリック
※ただし名前はある程度規則に沿った名前ではないとエラーとなります。
　特にこだわりが無い場合は初期値のままでOKです。
権限等の追加を求められる場合は承認

ベクトル検索で使用するベクトル算出

今回はオーナー権限で実施しているため、権限の追加等は行っておりませんが、
アプリ開発時などはFunctionを実行するアカウントにVertex AIやFirestoreの権限を
割り振る必要があります。

以下はVertex AIを使ってベクトルを算出し、Firestoreへデータを格納するコードです。
（参考公式：https://firebase.google.com/docs/firestore/vector-search）

functions-framework==3.*
google-cloud-firestore
google-cloud-aiplatform

import functions_framework
import os
# firestore
from google.cloud import firestore
from google.cloud.firestore_v1.vector import Vector
# Vertex AI
import vertexai
from vertexai.language_models import TextEmbeddingModel


# プロジェクト名（環境編酢より取得）
MY_PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") 

# 渡された文字列からベクトル値を算出する
def text_embedding(text: str) -> list:

    # locationは各自のロケーションを設定する
    vertexai.init(project=MY_PROJECT_ID, location="asia-northeast1") 

　　# 現在のベクトル算出のための最新AIが「textembedding-gecko@003」のためこちらを利用
    model = TextEmbeddingModel.from_pretrained("textembedding-gecko@003")
    embeddings = model.get_embeddings([text])
    for embedding in embeddings:
        vector = embedding.values

    return Vector(vector)


# メイン処理
# （関数名は適当です）
@functions_framework.http
def hello_http(request):

    # リクエストから記事の要約（description）を取得しています
    request_json = request.get_json(silent=True)
    request_args = request.args

    if request_json and 'description' in request_json:
        description = request_json['description']
    elif request_args and 'description' in request_args:
        description = request_args['description']
    else:
        description = 'World'

    # Firestoreクライアントの初期化
    firestore_client = firestore.Client(project=MY_PROJECT_ID)
    # コレクションの参照（コレクション名は任意（コレクションが無い場合は事前に作成してください））
    collection = firestore_client.collection("article_collection")

    # embeddingを計算
    embedding_vector = text_embedding(description)

    # Firestoreに追加するドキュメントを準備
    doc = {
        "description": description,
        "embedding_field": embedding_vector
    }
    # ドキュメントの追加
    collection.add(doc)

    return 'OK!'

今回は面倒だったのでターミナルより、CLI テストコマンドを実行して動作を確認しました。

curl -m 70 -X POST https://asia-northeast1-python-tool-001.cloudfunctions.net/vector_chenge \
  -H "Authorization: bearer $(gcloud auth print-identity-token)" \
  -H "Content-Type: application/json" \
  -d '{ "description": "<任意の文字列>"}'

実行が成功すると、Firestoreに以下の様にデータが格納されていると思います。

ベクトル検索のためのインデックス作成

ベクトル検索にはインデックスの作成が必須のようです。
今回はコンソールから以下コマンドを実行しインデックスを作成しました。
（参考公式：https://firebase.google.com/docs/firestore/vector-search）

gcloud alpha firestore indexes composite create \
  --collection-group=article_collection \
  --query-scope=COLLECTION \
  --field-config field-path=embedding_field,vector-config='{"dimension":"768", "flat": "{}"}' \
  --database=<データベースID>

collection-group：インデクスを作成するコレクション名
query-scope：こちらは分かりませんがインデックスを作成するスコープ
　　　　　　　複数のコレクション（コレクショングループ）などを範囲に指定できるようです。
field-path：ベクトルを格納しているフィールド名
vector-config：dimensionにベクトルの次元数を設定（今回は768次元だった）
database：対象となるデータベースのIDを指定。defaultの場合はこの指定は不要

実行すると以下のようにFirestoreにインデクスが作成されます。

Firestoreでのベクトル検索の実施コード

データの準備が整ったので、実際に検索を行います。

今回は私のブログ記事の要約内容を検索対象データとして準備しました。
全文表示すると多いので一部省略しています。

No.	タイトル
1	freezed.dartが作成されない時の対処方法 freezedを使用してイミュータブルなクラスを設計する際に、ターミナルで「…
2	Flutterのpubspec.yamlとは？意味や書き方をご紹介！！ YAMLはYAML Ain’t Markup Languageの略で、データを簡潔に表…
3	アプリ開発でよく聞くMVVMとは何か？ MVVM（Model-View-ViewModel）とは、アプリのロジックとUI（ユーザーインターフ…
4	Flutterとは？？Flutterの概要を説明「Flutter」について、”モバイルアプリを開発する場合に便利！”という認識はあ…
5	Riverpodとは？Flutterの一番メジャーな状態管理を紹介！！以前紹介した「StatefulWidget」も状態管理を行う機能の１…

実行したコードは以下です。

import functions_framework
import os
# firestore
from google.cloud import firestore
from google.cloud.firestore_v1.vector import Vector
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
# Vertex AI
import vertexai
from vertexai.language_models import TextEmbeddingModel


# プロジェクト名（環境編酢より取得）
MY_PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") 

# 渡された文字列からベクトル値を算出する
def text_embedding(text: str) -> list:

    # locationは各自のロケーションを設定する
    vertexai.init(project=MY_PROJECT_ID, location="asia-northeast1") 

　　# 現在のベクトル算出のための最新AIが「textembedding-gecko@003」のためこちらを利用
    model = TextEmbeddingModel.from_pretrained("textembedding-gecko@003")
    embeddings = model.get_embeddings([text])
    for embedding in embeddings:
        vector = embedding.values

    return Vector(vector)


# メイン処理
# （関数名は適当です）
@functions_framework.http
def hello_http(request):

    # リクエストから記事の要約（description）を取得しています
    request_json = request.get_json(silent=True)
    request_args = request.args

    if request_json and 'target' in request_json:
        target = request_json['target']
    elif request_args and 'target' in request_args:
        target = request_args['target']
    else:
        target = 'World'

    # Firestoreクライアントの初期化
    firestore_client = firestore.Client(project=MY_PROJECT_ID)
    # コレクションの参照
    collection = firestore_client.collection("article_collection")

    # embeddingを計算
    embedding_vector = text_embedding(target)

    # ベクトル検索の実施
    docs = collection.find_nearest(
        vector_field="embedding_field",
        query_vector=embedding_vector,
        distance_measure=DistanceMeasure.COSINE,
        limit=3
    ).get()

    # 表形式（ここでは文字列形式）での出力用
    output = "Description \n"
    output += "-" * 50 + "\n"
    
    # ベクトル検索で取得したドキュメントの内容を出力
    for doc in docs:
        doc_data = doc.to_dict()
        description = doc_data.get("description", "No description")
        # ドキュメントの内容を文字列に追加
        output += f"{description[:100]} \n"
    
    return output

こちらもターミナルより、CLI テストコマンドを実行して動作を確認しました。
「Riverpodについて」で検索を実行してみます！

curl -m 70 -X POST https://asia-northeast1-python-tool-001.cloudfunctions.net/vector_search \
-H "Authorization: bearer $(gcloud auth print-identity-token)" \
-H "Content-Type: application/json" \
-d '{
  "target": "Riverpodについて"
}'

実行結果

Description 
--------------------------------------------------
Riverpodとは？Flutterの一番メジャーな状態管理を紹介！！ 以前紹介した「StatefulWidget」も状態管理を行う機能の１つですが、 複数の画面や機能を持つアプリを実装する場合、管理 
Flutterとは？？Flutterの概要を説明 「Flutter」について、"モバイルアプリを開発する場合に便利！"という認識は あるかもしれませんが、何が便利で何が人気になっているのでしょうか？ 
アプリ開発でよく聞くMVVMとは何か？ MVVM（Model-View-ViewModel）とは、 アプリのロジックとUI（ユーザーインターフェース）を分けて、 開発の効率化と保守性の向上を目指したソ

特に並べ替えは行っていませんが、１件目にRiverpodの記事がきました！！
今回のデータでは２件目３件目について、近いものが選ばれたのか判断が付きませんでした。

最後に

もう少しデータ量を増やして検証した方が検索精度は検証できそうです。
ベクトルのデータサイズはfloatを４バイトとすると、今回の場合768次元のため768✕4≒3KBです。
ドキュメントの制限が１MBのため少し大きくも感じます。

Firestoreで全文検索できるようになったことは、まだプレビュー版とは言え吉報です。
今後の動向も期待したいです。