🔍
Workforce Identity + Auth0 で Vertex AI Search の ACL 制御を行う

2024/12/23に公開
Google Cloud
Auth0
Vertex AI
記事投稿コンテスト「今年の最も大きなチャレンジ」
workforceidentity
tech
3-shake AdventCalendar 第2シーズン 23日目の記事になります。2回目の登場です。

今回は真面目な(?)技術記事になります。私としては前回書いた記事も大真面目でしたが。

 概要今回やりたいこととしては、ウェブアプリケーション上で Id Provider(以後 IdP) 認証をして、その結果を利用して Vertex AI Agent Builder の Search 機能(以後めんどいので旧称の Vertex AI Search として説明) の ACL による検索の権限管理を行うというものです。

今回 IdP として Auth0 を利用します。そのため、少し Auth0 独自の設定が多くなりますが、OIDC の設定ができれば問題ないと思いますので、他の IdP をご利用の方は、適宜読み替えて見ていただければと思います。

今回のシーケンスとしては以下のような図になります。

 Workforce Identity とはそもそも Google Cloud の Workforce Identity Federation についてご存知でしょうか？
この機能は以下の説明のように IdP の認証情報を使って Google Cloud のリソースを触れるようにする機能です。[1]
外部 ID プロバイダ（IdP）を使用して、IAM を使用して従業員（従業員、パートナー、請負業者などのユーザー グループ）を認証および認可し、ユーザーが Google Cloud サービスにアクセスできるようにします
Cloud Identity でいいじゃない？を思う方も多いと思いますが、例えば組織で、すでに他の IdP を使っていた場合、Cloud Identity を使うと従業員のデータの2重管理になってしまいます。

また、Cloud Identity で進めた場合に一定人数以上になった場合には、人数によって金銭コストがかかってきます。

Workforce Identity はそのようなことを避けるため、IdP のデータを使って Google Cloud でリソースを作るなどができるソリューションになります。
また、今回の Vertex AI Search の ACL は呼び出し時に認証者の認証情報を利用し、その認証情報に含まれるユーザ情報に従って制御します。

この認証情報としては、現在 Cloud Identity と Workforce Identity の2種類の認証情報を使って、制御できます。
ちなみに、こちらの記事[2]でも触れられていますが、Google Cloud の似たような機能に Workload Identity があります。

このあたりの違いや、今回の設定について、別で登壇した際に簡単ではありますが、資料にまとめております。ご参照ください。

 実装
 前提まず、Workforce Identity を設定するには Google Cloud の組織設定が必要です。すでに組織が設定されている前提で話を進めます。
今回の Google Cloud の構成として以下の図をイメージしてもらえればよいです。
上記の Silde に設定手順は記載していますので、画面などはそちらを参考にしてください。

 Workforce Identity と Auth0 接続Auth0 の設定
Auth0 にて Application を作成し、今回の設定に必要な 発行元 Client ID クライアントシークレット の情報を取得します
Application に Allowed Callback URLs を https://${Cloud Runのドメイン}/callback に設定します
接続テストのためにも、https://auth.cloud.google/signin-callback/locations/global/workforcePools/{poolId}/providers/{providerId} も入れておくことをおすすめします。
ここの poolId, prividerId は次のWorkfoce Identity で設定するものになります。



Workforce Identity の設定
組織上で Pool を作成し、その中に Provider を設定する
Provider の設定時に属性マッピングとして、assertion.email を入れる


ここまでで Provider まで作成できたはずです。Workforce Identity ログインのテストをして、接続してください。

接続でうまくいかない場合は、組織上で Cloud Logging にて以下クエリを利用し調査を行ってください
resource.type="audited_resource"
resource.labels.service="sts.googleapis.com"
protoPayload.resourceName:"workforceIdentityPools/{poolId}"

 Vertex AI Search にて Workforce Identity を使うように設定を追加これで Workforce Identity と Auth0 の接続ができました。

対象のプロジェクトの Vertex AI Search に認証を使うように設定を追加します。

具体的には以下のように Agent Builder の設定を行います。
次に、Vertex AI Search にて認証を使うようなアプリと、Store を作成します。

BigQuery でも ACL を扱うことはできますが、今回は非構造化データ(Google Cloud Storage)を利用します。
リンクされた非構造化ドキュメント（JSONL とメタデータ）
1回限りの同期（定期的の場合は ACL を遵守しないと書いてあります[3])

このデータストアにはアクセス制御に関する情報が含まれています のチェックを忘れずにつけます。忘れると作り直しです。


JSONL のデータは以下の形式になっていれば問題ありません。
{"id": "doc-000001", "structData": {"description": "desc1", "title": "title1", "filename": "name1"}, "content": {"mimeType": "application/pdf", "uri": "${ファイルの位置}"}, "acl_info": {"readers": [{"principals": [{ "user_id": "${IdPにおける email の値}" },{ "user_id": "${IdPにおける email の値}" }]}]}}
{"id": "doc-000002", "structData": {"description": "desc2", "title": "title2", "filename": "name2"}, "content": {"mimeType": "application/pdf", "uri": "${ファイルの位置}"}, "acl_info": {"readers": [{"principals": [{ "user_id": "${IdPにおける email の値}" }]}]}}
ここで user_id は email 形式でなければ import 時にエラーになるので注意してください。
データとしては以下の PDF が見れるように入れてみています。
https://3-shake.com/announcement-20220921/

 Project への権限付与この状態だと、Project へのアクセス権限がないため次に、プロジェクトへの Role の設定をします。
まず以下の権限を持つ Custom Role を作成します。Vertex AI Search に対して検索を実行する最低限の権限になりますdiscoveryengine.answers.get
discoveryengine.servingConfigs.answer
discoveryengine.servingConfigs.search
discoveryengine.sessions.get
Pool に入っているひと 全体に上記の Custom Role を与えます
IAM に設定するのは principalSet://iam.googleapis.com/locations/global/workforcePools/{POOL_ID}/* とします
もし特定の人のみに設定する場合は、principal://iam.googleapis.com/locations/global/workforcePools/{POOL_ID}/subject/{IdPにおける email の値} とします。
Group の設定などを設定することもできます


ここまで作成しおえたら、Agent Builder にてテストを実施してください。
Agent Builder の統合タブの以下ページの 開く を押すと以下のような画面になるので検索してみて、ちゃんと必要なファイルだけが取得できているかを確認します。

 Cloud Run の方から利用するようやく、本題の web app です。

Cloud Run から Vertex AI Search の API をたたけるようにするためにアプリケーションを作成します。

今回は Python + Flask を使用し、Auth0 を利用します。認証のライブラリには authlib を使います。

以下のページを参考にするとスムーズかと思います。
https://auth0.com/docs/quickstart/webapp/python/interactive
ちなみに認証部分の完成版は以下に公開していただけておりますので、それを使って頂くのが簡単で良いと思います。

https://github.com/auth0-samples/auth0-python-web-app
これを使い、認証が通ることが確認できたと思います。
次に Workforce Identity Federation を実行するため idToken を使って sts.googleapis.com から credential を取得してきます。

以下ライブラリを追加します
google-cloud-discoveryengine == 0.12.3
google-auth == 2.36.0
google-cloud-core == 2.4.1
そして、実装として、以下のようにします。現状の library として、idToken を直接受け取って動く interface がないため、無理矢理に突っ込んでいます。[4][5]
from google.auth import identity_pool

def get_google_cloud_credentials(auth0_id_token: str):

    class Auth0SubjectTokenSupplier(identity_pool.SubjectTokenSupplier):
        def get_subject_token(self, context, re):
            return auth0_id_token

    return identity_pool.Credentials(
        audience=AUTH_AUDIENCE,
        token_url="https://sts.googleapis.com/v1/token",
        subject_token_type="urn:ietf:params:oauth:token-type:jwt",
        subject_token_supplier=Auth0SubjectTokenSupplier(),
    )
そして、この関数から取得できる credential を使って Vertex AI Search の API を叩きます。環境変数部分は適宜読み替えてください。
from google.cloud import discoveryengine_v1alpha as discoveryengine
from google.cloud import discoveryengine_v1beta as discoveryengine_beta
from google.api_core.client_options import ClientOptions

def search(query: str, credentials):
    client = discoveryengine.SearchServiceClient(
            client_options=ClientOptions(api_endpoint=f'${LOCATION}-discoveryengine.googleapis.com'),
            credentials=credentials
        )
    
    request = discoveryengine.SearchRequest(
            serving_config=client.serving_config_path(
                project=${PROJECT_ID},
                location=${LOCATION},
                data_store=${DATASTORE_ID},
                serving_config='default_search:search',
            ),
            query_expansion_spec=discoveryengine.SearchRequest.QueryExpansionSpec(
                condition=discoveryengine.SearchRequest.QueryExpansionSpec.Condition.AUTO
            ),
            spell_correction_spec=discoveryengine.SearchRequest.SpellCorrectionSpec(
                mode=discoveryengine.SearchRequest.SpellCorrectionSpec.Mode.AUTO
            ),
            content_search_spec=discoveryengine.SearchRequest.ContentSearchSpec(
                search_result_mode='DOCUMENTS',
                summary_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec(
                    summary_result_count=5,
                    include_citations=True,
                    language_code='ja',
                    model_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec.ModelSpec(
                        version='stable',
                    ),
                ),
                snippet_spec=discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(
                    max_snippet_count=1,
                    reference_only=True,
                    return_snippet=True,
                )
            ),
            query=query,
        )

    return client.search(request)
deploy 部分等は本題ではないので省きます。

こちらのアプリケーションを Cloud Run に deploy し動かすと、結果としては、Agent Builder でのテストと同じ動きを実現することができます。

 注意点今回は Workforce Identity の権限を使って Vertex AI Search を動かしましたが、Workforce Identity には制限があります。

 Workforce Identity でたたけない API があるWorkforce Identity による認証でできないことはそれなりにあります。例えば、Workforce Identity の認証情報では均一なバケットレベルのアクセスが許可されたバケットにしかアクセスができません。

https://cloud.google.com/iam/docs/federated-identity-supported-services?hl=ja#cloud-storage

そのため、Workfoce Identity の認証では Cloud Storage への ACL の制御がきかないので、検索結果ででてきた Cloud Storage 自体へのアクセス管理は別に実装を行う必要があります。

自前で実装ですかね。めんどい。

 ACL として設定できる数の上限ACL に使える user_id/group_id は最大 250 人までです。

https://cloud.google.com/generative-ai-app-builder/docs/data-source-access-control#limitations

とはいえ、基本的には group_id で 250 も利用することは想定しにくいので、あまり気にしなくても良いかもです。

 Session の認証期限があるWorkfoce Identity のセッションの時間として基本 15分から 12時間まで選べます。

その時間が切れると Workforce Identity のセッションとして切れますので注意が必要です。

今回の実装だと毎回検索時に credential を作成するのでそんなに認識しなくてもいいですが、覚えておくとよいです。

これとは別ですが Auth0 における idToken の期限もありますのでそこは別途注意が必要です。

 まとめVertex AI Search で IdP の認証結果を用いた検索について書かせてもらいました。

正直単純な検索アプリであれば Vertex AI Search の出来合いのものでも十分かと思いますが、機能や UI を修正していきたいということになれば、今回のような実装が必要になります。

もしもしで、これが皆様の役立つことあれば幸いです！
更にちなみに、今回副次的ではありますが、Auth0 何もわからん、から 完全に理解した、まではいきました。

脚注
https://cloud.google.com/workforce-identity-federation?hl=en ↩︎
https://zenn.dev/google_cloud_jp/articles/3f06bac31aa013 ↩︎
https://cloud.google.com/generative-ai-app-builder/docs/create-data-store-es#cloud-storage ↩︎
ここの実装は session store のようなものを作成できればいいのかもしれないですが、web application でそれを作成するのはかなり手間なのでまたの機会に考えてみます ↩︎
ちなみに、今回 id_token をどう保存するかなどは書いてません、が基本的には DB に保存しておきそれを随時読み込むようにします。これはこれで実装が大変ですね。。。 ↩︎