tl;dr
Python で提供されている discoveryengine ライブラリを使って Vertex AI Search のデータストアロードを自動化した。大量のファイルを読み込みたいなら GcsSource の設定のリミットに気をつけろ
やりたきこと
- Google Cloud Storage をデータソースにしている Vertex AI Search の Data Store を読み込みたい。
前準備
- サービスアカウント発行しておく
- サービスアカウントに権限ロール "ディスカバリーエンジン編集者" を付与しておく (discoveryengine.documents.import 権限が必要になります。)
実装
from google.cloud import discoveryengine
KEY_JSON = "{key json file path}"
# @see: https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1.services.document_service.DocumentServiceClient#google_cloud_discoveryengine_v1_services_document_service_DocumentServiceClient_import_documents
# @see: https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1.services.document_service.DocumentServiceClient#google_cloud_discoveryengine_v1_services_document_service_DocumentServiceClient_update_document
project = "{ project id }"
location = "global" # ここは固定
collection = "default_collection" # ここは固定
data_store = "{ data store's id }" # Vertex AI Search のデータストアの ID を指定しておく
branch = "0" # ここは固定
parent_value = f"projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/branches/{branch}"
def import_documents():
client = discoveryengine.DocumentServiceClient.from_service_account_json(
KEY_JSON)
gcs = discoveryengine.GcsSource(
input_uris=["gs://hogehoge/*"], data_schema="content")
request = discoveryengine.ImportDocumentsRequest(
parent=parent_value, gcs_source=gcs)
response = client.import_documents(request=request)
print(response)
import_documents()
注意点
GCS をソースとして指定する場合、そのまま使うとリミットが100しかない
ライブラリのコードを覗きにいくと、data_schema を content にしておくことで100,000ファイルまで読み込みが可能になると書いてある。。
ちゃんとエラーも出ます。
google.api_core.exceptions.InvalidArgument: 400 The request contained 342 files which exceeds the maximum number of files allowed (100).
ただし一応制約もあるので注意。
- content で指定可能なデータ形式に限りあり。 (PDF や HTML などを限定としている。)
- 1ファイルの上限が 100MB までになってしまう。
import_config.py
"""
input_uris (MutableSequence[str]):
Required. Cloud Storage URIs to input files. URI can be up
to 2000 characters long. URIs can match the full object path
(for example, ``gs://bucket/directory/object.json``) or a
pattern matching one or more files, such as
``gs://bucket/directory/*.json``.
A request can contain at most 100 files (or 100,000 files if
``data_schema`` is ``content``). Each file can be up to 2 GB
(or 100 MB if ``data_schema`` is ``content``).
"""
さいごに
Vertex AI Search めちゃくちゃ便利なのでどんどん使いたい。
たくさんの需要に応えるために環境構築の自動化にもどんどん取り組む。
まだドキュメントが少ないがどんどんトライする。
Discussion