テックブログ

from __future__ import annotations
import os
import functions_framework

from google.cloud import discoveryengine, storage

project_id = os.getenv("PROJECT_ID")
search_engine_id = os.getenv("SEARCH_ENGINE_ID")
gcs_bucket_name = os.getenv("GCS_BUCKET_NAME")


# バケット名とprefix を指定すると、ディレクトリ一覧を取得できます
def list_directories(bucket_name, prefix=""):
    client = storage.Client()
    blobs = client.list_blobs(bucket_name, prefix=prefix)
    files = []
    for b in blobs:
        files.append(b.name)
    dirs = set()
    for file in files:
        tmp = "/".join(file.split("/")[0:-1]) + "/**"
        dirs.add("gs://{}/{}".format(bucket_name, tmp))
    
    return list(dirs)


def import_documents_sample(
    project_id: str,
    location: str,
    search_engine_id: str,
    data_schema: str = "custom",
    gcs_bucket_name: str | None = None,
    gcs_prefix: str = "",
    bigquery_dataset: str | None = None,
    bigquery_table: str | None = None,
) -> str:
    # Create a client
    client = discoveryengine.DocumentServiceClient()
    # The full resource name of the search engine branch.
    # e.g. projects/{project}/locations/{location}/dataStores/{data_store}/branches/{branch}
    parent = client.branch_path(
        project=project_id,
        location=location,
        data_store=search_engine_id,
        branch="default_branch",
    )
    
    uris = list_directories(gcs_bucket_name, gcs_prefix)
    uris_from_root_splitted = [uri[5:].split("/")[1:] for uri in uris]  # remove "gs://bucket_name/"
    max_depth_from_root = max([len(uri_from_root) for uri_from_root in uris_from_root_splitted])

    depth_prefix = len([folder for folder in gcs_prefix.split("/") if folder != ""])
    max_depth_from_prefix = max_depth_from_root - depth_prefix
       
    uris_wildcard = []
    for i in range(max_depth_from_prefix - 1):
        uri_depth_i_wildcard = f"gs://{gcs_bucket_name}{gcs_prefix}" + "/**" * (i + 1)
        uri_depth_i_wildcard = uri_depth_i_wildcard[:-1]
        uris_wildcard.append(uri_depth_i_wildcard)
    request = discoveryengine.ImportDocumentsRequest(
        parent=parent,
        gcs_source=discoveryengine.GcsSource(
            input_uris=uris_wildcard,
            data_schema=data_schema,
        ),
        # Options: `FULL`, `INCREMENTAL`
        reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
    )

    # Make the request
    operation = client.import_documents(request=request)
    print(f"Waiting for operation to complete: {operation.operation.name}")
    response = operation.result()
    # Once the operation is complete,
    # get information from operation metadata
    metadata = discoveryengine.ImportDocumentsMetadata(operation.metadata)
    # Handle the response
    print(response)
    print(metadata)
    return operation.operation.name

@functions_framework.cloud_event
def refresh(cloud_event):
    response = import_documents_sample(
        project_id=project_id,
        location="global",
        search_engine_id=search_engine_id,
        data_schema="content",
        gcs_bucket_name=gcs_bucket_name,
        # 特定のディレクトリ以下を対象としたい場合、ディレクトリ名に変更してください。
        # バケットに直接ファイルが入っている場合はこのまま
        gcs_prefix="",
    )

functions-framework==3.*
google-cloud-storage==2.10.0
google-cloud-discoveryengine==0.11.2