Vertex AI Agent Builder の検索システムを Python SDK から試してみる

はじめに

Vertex AI Agent Builder で作る検索システム

Vertex AI Agent Builder（旧 Vertex AI Search & Conversation）を使用すると文書検索システムが簡単に構築できて、コンソール上のデモ用検索ポータルから検索処理が体験できます。検索キーワードの「意味」を理解して検索するセマンティックサーチを行うので、次のように微妙にタイプミスをしても、こちらの意図を汲み取って検索結果を返してくれます。また、検索結果のサマリーテキストも表示されます。

コンソールの検索ポータルで検索する例

Vertex AI Agent Builder による検索システムは、次のような構成になります。「データストア」と「検索アプリ」の2つのコンポーネントを作成して利用します。

Agent Builder による検索システムの構成図

データストアは、ドキュメント検索の基本機能を提供します。Cloud Storage などのデータソースからドキュメントをインポートすると、ドキュメントの内容を分析して、検索に必要な情報を抽出・保存します。「RAG 用のチャンク取得」で説明する機能は、データストアを用意するだけで利用できます。そして、検索アプリは、サマリー生成などの追加機能を提供します。「ドキュメント検索」、および、「ドキュメントサマリーの表示」で説明する機能を使用する際は、検索アプリの作成が必要です。

コンソール上の検索ポータルは、「社内ポータルサイト・外部公開 Web サイトなど」の部分にあたります。本来は、検索機能を組み込みたい Web サイトやアプリケーションなど、利用者が独自に構築する部分です。その際は、検索ウィジェットを Web ページに組み込んだり^[1]、クライアント SDK を用いて直接 API を呼び出すと言った使い方が必要です。

この記事で試す事

この記事では、Python SDK を用いて、Python のコードで「API リクエスト」の部分を実行する例を紹介します。次のように、Vertex AI Workbench のノートブックからコードを実行していきます。

この記事で作成する環境の構成図

この後の手順では、特に次の機能を試します。

1. ドキュメントごとに ACL を設定して、検索できるユーザーを制限する。
2. ドキュメントにメタデータを追加して登録する。
3. ドキュメントのリストに加えて、検索文に関連した箇所のテキストチャンクを返却する。

2 のメタデータは、検索対象とするキーワードを追加したり、検索システムを利用するアプリケーションにドキュメントの補足情報を渡す際に使用します。3 のテキストチャンクは RAG システムを構築する際に利用します。エンドユーザーの質問文を検索システムに渡して、得られたテキストチャンクを情報源として任意の LLM で質問の回答を生成すると言った使い方になります。

事前準備

新しいプロジェクトを作成した後に、Cloud Shell を開いて次のコマンドを実行していきます。これ以降に登場するスクリーンショットでは etsuji-search-demo という名前のプロジェクトを使用しています。

API 有効化

Vertex AI Workbench を使用するのに必要な API を有効化します。

gcloud services enable \
  aiplatform.googleapis.com \
  notebooks.googleapis.com

ストレージバケット作成

検索対象のドキュメントは、事前に Cloud Storage に保存した上で、Agent Builder のデータストアにインポートすることで検索可能になります。ドキュメントを事前に保存するバケットを作成します。バケット名は [Project ID]-sample とします^[2]。

PROJECT_ID=$(gcloud config list --format 'value(core.project)')
BUCKET="gs://${PROJECT_ID}-sample"
gsutil mb -b on -l us-central1 $BUCKET

gsutil コマンドでバケットを作成する際は、オプション -b on を忘れずにつけてください。これにより、データストアにドキュメントをインポートする際に必要なアクセス権がストレージバケットに設定されます。

検索利用者の IAM 登録

検索機能を使用するユーザーのアカウントに discoveryengine.viewer ロールを付与します。ここでは、[e-mail] と [e-mail] の二人のユーザーを用いて、それぞれ、自身に許可されたドキュメントだけが検索対象になることを確認します。[e-mail] は日本語のドキュメント、[e-mail] は英語のドキュメントを扱うという想定です。それぞれのアカウント名は、実際に使用するアカウント名に置き換えてください。

japanese_account='[e-mail]'
english_account='[e-mail]'

gcloud projects add-iam-policy-binding $PROJECT_ID \
  --member user:$japanese_account \
  --role roles/discoveryengine.viewer

gcloud projects add-iam-policy-binding $PROJECT_ID \
  --member user:$english_account \
  --role roles/discoveryengine.viewer

Workbench インスタンス作成

この後、Workbench のノートブックから作業を行うので、Workbench インスタンスを作成しておきます。

gcloud workbench instances create search-development \
  --project=$PROJECT_ID \
  --location=us-central1-a \
  --machine-type=e2-standard-2

ダミードキュメントの作成

検索に使用するダミーのドキュメントを作成します。

作成手順

クラウドコンソールのナビゲーションメニューから「Vertex AI」→「ワークベンチ」を選択すると、先ほど作成したインスタンス search-develpment があります。インスタンスの起動が完了するのを待って、「JUPYTERLAB を開く」をクリックしたら、「Python 3(ipykernel)」の新規ノートブックを作成します。

この後は、ノートブックのセルで次のコマンドを実行していきます。

はじめに、必要なパッケージをインストールします。

!pip install --user \
  google-cloud-discoveryengine pydata-google-auth Faker

インストールしたパッケージを使用するために、一度、カーネルをリスタートします。

import IPython
app = IPython.Application.instance()
_ = app.kernel.do_shutdown(True)

必要なモジュールをインポートします。

import os, shutil, json
from faker import Faker

ディレクトリ sample_data を作成して、日本語と英語のテキストファイルをそれぞれ50個ずつ作成します。ファイルの内容は、Faker パッケージで作成したランダムな文章です。

tmp_dir = 'sample_data'
shutil.rmtree(tmp_dir, ignore_errors=True)
os.makedirs(tmp_dir)

for lang in ['ja-JP', 'en-US']:
    fake = Faker(lang)
    for c in range(50):
        filename = f'{tmp_dir}/{lang}-{c:03d}.txt'
        with open(filename, 'wt') as f:
            f.write(fake.text(max_nb_chars=1000))

作成したドキュメントに対するメタデータを用意します。次のコマンドを実行すると、1つのドキュメントに対して、1行のJSON文字列でメタデータを記述したファイル metadata.jsonl が用意されます。最初に設定しているユーザーアカウント [e-mail] と [e-mail] は、実際に使用するものに置き換えてください。

japanese_account = '[e-mail]'
english_account = '[e-mail]'

PROJECT_ID, = !gcloud config list --format 'value(core.project)'
BUCKET = f'gs://{PROJECT_ID}-sample'

jsonl_text = ''
for lang in ['ja-JP', 'en-US']:
    for c in range(50):
        if lang == 'ja-JP':
            account = japanese_account
            title = f'サンプルドキュメント:{c:03d}'
        else:
            account = english_account
            title = f'Sample document-{c:03d}'
        fake = Faker(lang)
        filename = f'{lang}-{c:03d}.txt'
        doc_info = {
            'id': filename.split('.')[0],
            'structData': {
                'title': title,
                'company': fake.company()
            },
            'content': {
                'mimeType': 'text/plain',
                'uri': f'{BUCKET}/{filename}'
             },
            'acl_info': {
                'readers': [{'principals': [{'user_id': f'{account}'}]}]
            }
        }
        jsonl_text += json.dumps(doc_info) + '\n'

filename = f'{tmp_dir}/metadata.jsonl'
with open(filename, 'wt') as f:
    f.write(jsonl_text)

先に作成したストレージバケットに、ここで用意したファイルをまとめてコピーします。

!gsutil cp {tmp_dir}/* {BUCKET}/

メタデータの構成について

先ほど実行したコマンドの下記の部分が、1つのドキュメントに対するメタデータの内容です。

        doc_info = {
            'id': filename.split('.')[0],
            'structData': {
                'title': title,
                'company': fake.company()
            },
            'content': {
                'mimeType': 'text/plain',
                'uri': f'{BUCKET}/{filename}'
             },
            'acl_info': {
                'readers': [{'principals': [{'user_id': f'{account}'}]}]
            }
        }

各要素の説明は、次のとおりです。

要素	説明
id	ドキュメント ID
structData	任意の追加情報
content	ドキュメントファイル
acl_info	ドキュメントに対するアクセス権限

ドキュメント ID は、それぞれのドキュメントにユニークな ID を割り当てます。content は、ドキュメントファイルの情報で、mimeType に MIME タイプ、uri にストレージパスを指定します。acl_info はアクセス権限の設定で、readers の principals リストに、このドキュメントの検索を許可するユーザー、および、グループを並べます。一般には、次のような形式になります。

'principals': [
    { 'user_id': 'user_1' },
    { 'user_id': 'user_2' },
    { 'group_id': 'group_1' },
    { 'group_id': 'group_2' }
]

今回の設定では、日本語のドキュメント ja-JP-XXX.txt には [e-mail] 、英語のドキュメント en-US-XXX.txt には [e-mail] を readers に指定しています。

最後に structData には、ディクショナリ形式で任意のメタデータを付与することができます。ここでは、例として、ドキュメントのタイトル（title）とドキュメントを作成した会社名（company）を追加しています。

データストアの準備

データストアを作成して、先ほどバケットに保存したダミードキュメントをインポートします。

この作業は、クラウドコンソールから行います。

はじめに、ナビゲーションメニューから「Agent Builder」を選択すると、利用開始の確認画面が表示されるので、「CONTINUE AND ACTIVATE THE API」をクリックします。

認証設定

Agent Builderのメニューから「設定」を選ぶと、次の認証設定画面が表示されます。

Agent Builder の認証設定画面

ここでは、認証に使用する ID プロバイダーを設定します。今回は、Google 標準のアカウント認証を用いるので、3箇所の鉛筆アイコンをクリックして、すべて「Google Identity」を選択します。

データストア作成

Agent Builderのメニューから「データストア」を選択して、「データストアを作成」をクリックします。次のデータソースの選択画面が表示されるので、「Cloud Storage」を選択します。

データソースに Cloud Storage を選択

次の画面では、「ファイル」を選択して、「参照」をクリックした後に、先ほどアップロードしたメタデータファイル metadata.jsonl を選択します。「メタデータを含む非構造化ドキュメントの JSONL」を選択して、「続行」をクリックします。

データストアの作成画面 (1)

次の画面では、「データストアのロケーション」に「global（グローバル）」、データストア名に「sample-datastore」を入力し、さらに、「このデータストアにはアクセス制御に関する情報が含まれています」をチェックします。また、「DOCUMENT PROCESSING OPTIONS」を開いて「チャンクモードを有効にする」をチェックします。最後に「作成」をクリックします。

データストアの作成画面 (2)

データストアの一覧画面に切り替わるので、「sample-datastore」をクリックして、「アクティビティ」タブを選択すると、インポートの進行状況が表示されます。次のように「インポートが完了しました」と表示されるまで、10〜20分程度待ちます。また、この画面に表示されている「データストアの ID」の値をメモしておきます。

インポートが完了した状態

Agent Builderのメニューから「データストア」を選択して、再度、データストアの一覧画面から「sample-datastore」をクリックすると、次のように「スキーマ」のタブが追加されているので、これを選択します。

スキーマの情報

先ほど structData で定義した、「title」と「company」のメタデータが付与されていることがわかります。どちらも「検索可能」がチェックされているので、これらの内容も検索対象となります。「キープロパティ」は、「title」「description」「category」など、事前に定義された項目にマッピングするもので、これらの項目の「意味（役割）」を理解した検索処理が行われます。「EDIT」をクリックして、これらの設定を変更することができます。

検索アプリの作成

Agent Builder のメニュー「アプリ」を選んで、「新しいアプリを作成」をクリックします。アプリの種類を選択する次の画面が表示されるので、「検索」を選択します。

アプリの種類に「検索」を選択

アプリケーションの構成を設定する次の画面が表示されるので、次のように設定して、「作成」をクリックします。

• Enterprise エディションの機能をオフにする^[3]
• アプリ名に「sample-application」を入力する
• 会社名または組織名に「example company」を入力する
• アプリのロケーションに「global（グローバル）」を選択する

アプリケーションの構成画面 (1)

データストアを選択する次の画面が表示されるので、「sample-datastore」をチェックして、「作成」をクリックします。

アプリケーションの構成画面 (2)

これで、検索システムの構築は完了です。次は、Python SDK を用いて、検索処理をテストします。

検索処理のテスト

Python SDK を用いて、検索処理を実行してみます。ダミードキュメントの作成に使用したノートブックに戻って、ノートブックのセルで次のコマンドを実行していきます。

事前準備

必要なモジュールをインポートします。

import pydata_google_auth
from google.cloud import discoveryengine_v1alpha as discoveryengine
from google.api_core.client_options import ClientOptions

プロジェクト ID、データストア ID、そして、データストアのロケーションを設定します。DATASTORE_ID には、先ほど確認したデータストアの ID を指定します。

PROJECT_ID, = !gcloud config list --format 'value(core.project)'
DATASTORE_ID = 'sample-datastore_1714456642133'
LOCATION = 'global'

検索を実行するユーザーの認証情報を取得します。

credentials = pydata_google_auth.get_user_credentials(
    ['https://www.googleapis.com/auth/cloud-platform'],
    use_local_webserver=False,
    credentials_cache=pydata_google_auth.cache.NOOP
)

このコマンドを実行すると、次のように認証処理の URL が表示されるので、これをコピーしてブラウザーで開きます。

Please visit this URL to authorize this application: https://accounts.google.com/o/oauth2/auth?response_type=code&client_id=...
Enter the authorization code: 

ユーザーアカウントを選択して認証すると、次の画面が表示されます^[4]。

ユーザー認証が完了した画面

画面の下部に表示された認証コードをコピーして、ノートブックの Enter the authorization code: に入力すると、このユーザーの認証情報（アクセストークンなど）が credentials に保存されます^[5]。ここでは、英語のドキュメントを使用するユーザー [e-mail] を選択したものとして説明を続けます。

ドキュメント検索

ドキュメントを検索する関数 get_documents() を定義します。

def get_documents(query, 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=3,
                include_citations=True,
                ignore_adversarial_query=True,
                ignore_non_summary_seeking_query=True,
                model_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec.ModelSpec(
                    version='stable'
                )
            ),
        ),
        query=query,
        page_size=10,
    )

    response = client.search(request)
    return response

コード内のオプション summary_result_count=3 は、検索結果からトップ3のドキュメントを用いてサマリーを作成するという指定です。1〜5の値が指定できます。

適当なキーワードで検索してみます。

response = get_documents('立候補者の情報', credentials)

response には検索されたドキュメントの情報をまとめたオブジェクトが格納されており、for 文で順番にドキュメント情報を取り出すことができます。次のコマンドを実行すると、ドキュメント情報のデータ構造が確認できます。

for item in response:
    print(item)
    break

ここでは、先頭の3つのドキュメントについて、タイトル、会社名、リンクを表示してみます。タイトルと会社名は、メタデータとして埋め込んだ情報です。

c = 0
for item in response:
    print(f"\
{item.document.struct_data['title']}, \
{item.document.struct_data['company']}, \
{item.document.derived_struct_data['link']}")
    c += 1
    if c == 3:
        break

結果は次のようになります。

Sample document-027, Willis, Chavez and Moore, gs://etsuji-search-demo-sample/en-US-027.txt
Sample document-023, Valdez LLC, gs://etsuji-search-demo-sample/en-US-023.txt
Sample document-036, Molina-Pena, gs://etsuji-search-demo-sample/en-US-036.txt

先ほど認証したユーザー [e-mail] に許可された、英語のドキュメントのみが検索されています。ユーザー [e-mail] で認証すれば、日本語のドキュメントのみが検索されることも確認できるでしょう。

それでは、これらのドキュメントには、検索キーワードである「立候補者の情報」に関連する内容は本当に含まれているでしょうか？　英語のドキュメントなので、「立候補者」に対応する単語をチェックしてみます。

!grep '[Cc]andidate' sample_data/en-US-027.txt

次の結果が得られました。確かに candidate（立候補者）という単語が含まれています。

Inside age standard candidate contain catch physical development. Yes think alone left life this time. Analysis American everyone blood work.

場合によっては、candidate という単語そのものは含まれていない事もありますが、テキスト全体をみると、「election」などの選挙に関連した言葉が含まれているはずです。

メタデータとして登録したドキュメントのタイトルでも検索できるか確認してみます。

response = get_documents('Sample document-015', credentials)
c = 0
for item in response:
    print(f"\
{item.document.struct_data['title']}, \
{item.document.struct_data['company']}, \
{item.document.derived_struct_data['link']}")
    c += 1
    if c == 3:
        break

次の結果が得られました。指定したタイトルのドキュメントが検索結果のトップになっています。

Sample document-015, Keller-Mathis, gs://etsuji-search-demo-sample/en-US-015.txt
Sample document-008, Gould, Haynes and Brown, gs://etsuji-search-demo-sample/en-US-008.txt
Sample document-021, Nelson, Stevenson and Keller, gs://etsuji-search-demo-sample/en-US-021.txt

ドキュメントサマリーの表示

先ほどの検索結果には、コンソールの検索ポータルで表示されたサマリーテキストも含まれており、次のコマンドで内容が確認できます。

print(response.summary.summary_text)
for c, i in enumerate(response.summary.summary_with_metadata.references):
    print(f'[{c+1}] {i.title}')

「立候補者の情報」で検索した後に実行すると、次の結果になりました。今回は、内容が出鱈目なダミードキュメントなので、意味のあるサマリーは生成できなかったようです。

立候補者の情報はありません [1, 2, 3]。
[1] Sample document-027
[2] Sample document-023
[3] Sample document-036

きちんとしたドキュメントであれば、検索ポータルと同様の次のような結果が得られます。

アジャイル開発は、要件が明確で、変更が少なく、開発期間が短いプロジェクトに適しています[1]。ウォーターフォール型の開発手法は、要件が明確で、変更が少なく、開発期間が短いプロジェクトに適しています[2]。
[1] アジャイル開発実践ガイドブック 
[2] デジタル社会推進標準ガイドライン DS-100 

RAG 用のチャンク取得

RAG に使用するテキストチャンクを取得する関数 get_chunks() を定義します。

def get_chunks(query, 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='CHUNKS'
        ),
        query=query,
        page_size=10,
    )

    response = client.search(request)
    return response

適当なキーワードで検索してみます。

response = get_chunks('立候補者の情報', credentials)

先ほどと同様に、response から for 文で順番にテキストチャンクを取り出すことができます。ここでは、先頭の3つのチャンクを表示してみます。

c = 0
for item in response:
    print(f'[Chunk from {item.chunk.document_metadata.uri}]')
    print(item.chunk.content)
    print()
    c += 1
    if c == 3:
        break

結果は次のようになります。

[Chunk from gs://etsuji-search-demo-sample/en-US-027.txt]
Outside late ago. Since more newspaper five approach her soon rock. Office agree there current guess provide beat.
Page door future seat with economic matter. Bad exist common suffer. Product think subject cup teacher population.
Inside social else soon summer more artist. Learn threat decision fight either. Election data others natural minute.
（以下省略）

得られたテキストチャンクを情報源として、任意の LLM で質問の回答を生成することができます。

まとめ

この記事では、Python SDK を用いて、Vertex AI Agent Builder（旧 Vertex AI Search & Conversation）の文書検索システムを Python のコードから利用する例を紹介しました。単純なドキュメント検索だけではなく、テキストチャンクを取得して RAG のシステムを構築する際のパーツとして利用できることもわかりました。

生成 AI というと、「チャットのインターフェースで調べ物をする道具」というイメージを持つ方も多いようですが、さまざまなパーツを組み合わせることで、「生成 AI を活用したアプリケーション」が構築できます^[6]。Vertex AI Agent Builder は、その名の通り、アプリケーション構築のさまざまなパーツを提供するサービスです。検索以外にどんな機能があるのか、興味を持った方は、ぜひ調べてみてください。

脚注

1. 検索ウィジェットを Web ページに組み込む方法は、公式ドキュメント Add the search widget to a web page を参照してください。

2. 本文内の [Project ID] の部分は、実際に使用するプロジェクトのプロジェクト ID に読み替えてください。

3. 「Enterprise エディションの機能」と「高度な LLM 機能」が提供する機能については、公式ドキュメント About advanced features を参照してください。

4. 画面に「Sign in to BigQuery」と表示されていますが、これは、BigQuery 用のライブラリパッケージ pydata_google_auth の認証機能を利用しているためです。

5. Web アプリケーションから検索機能を利用する際は、同等の認証機能を Web アプリケーションに組み込む必要があります。「Google認証とOAuth2.0で、Googleサービスと連携する例のアレを作る」などの具体例を参考にしてください。

6. 生成 AI を利用したアプリケーション構築については、書籍「Google Cloudで学ぶ生成AIアプリ開発入門」が参考になります。