chore: import upstream snapshot with attribution

2026-07-13 13:30:30 +08:00
commit 914fea506e
2793 changed files with 802106 additions and 0 deletions
@@ -0,0 +1,111 @@
+"""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,
+    )
@@ -0,0 +1,2 @@
+functions-framework==3.*
+google-cloud-discoveryengine==0.12.2