wehub-resource-sync
112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
"""Function to update the Vertex AI Search and Conversion index"""
|
|
|
|
import os
|
|
|
|
import functions_framework
|
|
from google.api_core.client_options import ClientOptions
|
|
from google.cloud import discoveryengine
|
|
|
|
|
|
def import_documents_sample(
|
|
project_id: str,
|
|
location: str,
|
|
data_store_id: str,
|
|
gcs_uri: str | None = None,
|
|
bigquery_dataset: str | None = None,
|
|
bigquery_table: str | None = None,
|
|
) -> str:
|
|
"""Function to import documents"""
|
|
# For more information, refer to:
|
|
# https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store
|
|
client_options = (
|
|
ClientOptions(api_endpoint=f"{location}-discoveryengine.googleapis.com")
|
|
if location != "global"
|
|
else None
|
|
)
|
|
|
|
# Create a client
|
|
client = discoveryengine.DocumentServiceClient(client_options=client_options)
|
|
|
|
# The full resource name of the search engine branch.
|
|
# e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}
|
|
parent = client.branch_path(
|
|
project=project_id,
|
|
location=location,
|
|
data_store=data_store_id,
|
|
branch="default_branch",
|
|
)
|
|
|
|
if gcs_uri:
|
|
request = discoveryengine.ImportDocumentsRequest(
|
|
parent=parent,
|
|
gcs_source=discoveryengine.GcsSource(
|
|
input_uris=[gcs_uri], data_schema="document"
|
|
),
|
|
# Options: `FULL`, `INCREMENTAL`
|
|
reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
|
|
)
|
|
else:
|
|
request = discoveryengine.ImportDocumentsRequest(
|
|
parent=parent,
|
|
bigquery_source=discoveryengine.BigQuerySource(
|
|
project_id=project_id,
|
|
dataset_id=bigquery_dataset,
|
|
table_id=bigquery_table,
|
|
data_schema="custom",
|
|
),
|
|
# 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 update_search_index(cloud_event):
|
|
"""Main function"""
|
|
# Set event vars
|
|
data = cloud_event.data
|
|
event_id = cloud_event["id"]
|
|
event_type = cloud_event["type"]
|
|
bucket = data["bucket"]
|
|
name = data["name"]
|
|
metageneration = data["metageneration"]
|
|
timeCreated = data["timeCreated"]
|
|
updated = data["updated"]
|
|
|
|
# Print event vars to log
|
|
print(f"Event ID: {event_id}")
|
|
print(f"Event type: {event_type}")
|
|
print(f"Bucket: {bucket}")
|
|
print(f"File: {name}")
|
|
print(f"Metageneration: {metageneration}")
|
|
print(f"Created: {timeCreated}")
|
|
print(f"Updated: {updated}")
|
|
|
|
project_id = os.environ["PROJECT_ID"]
|
|
location = "global"
|
|
data_store_id = os.environ["DATASTORE_ID"]
|
|
docs_metadata_bucket = os.environ["DOCS_METADATA_BUCKET"]
|
|
gcs_uri = f"gs://{docs_metadata_bucket}/*.jsonl"
|
|
|
|
import_documents_sample(
|
|
project_id=project_id,
|
|
location=location,
|
|
data_store_id=data_store_id,
|
|
gcs_uri=gcs_uri,
|
|
)
|