samples(discoveryengine): Add Create/Delete Data Store Samples (GoogleCloudPlatform#11870)

Author: holtskinner

discoveryengine/create_data_store_sample.py

@@ -0,0 +1,86 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# [START genappbuilder_create_data_store]
+
+from google.api_core.client_options import ClientOptions
+from google.cloud import discoveryengine
+
+# TODO(developer): Uncomment these variables before running the sample.
+# project_id = "YOUR_PROJECT_ID"
+# location = "YOUR_LOCATION" # Values: "global"
+# data_store_id = "YOUR_DATA_STORE_ID"
+
+
+def create_data_store_sample(
+    project_id: str,
+    location: str,
+    data_store_id: str,
+) -> str:
+    #  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.DataStoreServiceClient(client_options=client_options)
+
+    # The full resource name of the collection
+    # e.g. projects/{project}/locations/{location}/collections/default_collection
+    parent = client.collection_path(
+        project=project_id,
+        location=location,
+        collection="default_collection",
+    )
+
+    data_store = discoveryengine.DataStore(
+        display_name="My Data Store",
+        # Options: GENERIC, MEDIA, HEALTHCARE_FHIR
+        industry_vertical=discoveryengine.IndustryVertical.GENERIC,
+        # Options: SOLUTION_TYPE_RECOMMENDATION, SOLUTION_TYPE_SEARCH, SOLUTION_TYPE_CHAT, SOLUTION_TYPE_GENERATIVE_CHAT
+        solution_types=[discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH],
+        # Options: NO_CONTENT, CONTENT_REQUIRED, PUBLIC_WEBSITE
+        content_config=discoveryengine.DataStore.ContentConfig.PUBLIC_WEBSITE,
+    )
+
+    request = discoveryengine.CreateDataStoreRequest(
+        parent=parent,
+        data_store_id=data_store_id,
+        data_store=data_store,
+        # Optional: For Advanced Site Search Only
+        # create_advanced_site_search=True,
+    )
+
+    # Make the request
+    operation = client.create_data_store(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.CreateDataStoreMetadata(operation.metadata)
+
+    # Handle the response
+    print(response)
+    print(metadata)
+
+    return operation.operation.name
+
+
+# [END genappbuilder_create_data_store]

discoveryengine/data_store_sample_test.py

@@ -0,0 +1,37 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+from uuid import uuid4
+
+from discoveryengine import create_data_store_sample, delete_data_store_sample
+
+project_id = os.environ["GOOGLE_CLOUD_PROJECT"]
+location = "global"
+data_store_id = f"test-data-store-{str(uuid4())}"
+
+
+def test_create_data_store():
+    operation_name = create_data_store_sample.create_data_store_sample(
+        project_id, location, data_store_id
+    )
+    assert operation_name
+
+
+def test_delete_data_store():
+    operation_name = delete_data_store_sample.delete_data_store_sample(
+        project_id, location, data_store_id
+    )
+    assert operation_name

discoveryengine/delete_data_store_sample.py

@@ -0,0 +1,56 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# [START genappbuilder_delete_data_store]
+
+from google.api_core.client_options import ClientOptions
+from google.cloud import discoveryengine
+
+# TODO(developer): Uncomment these variables before running the sample.
+# project_id = "YOUR_PROJECT_ID"
+# location = "YOUR_LOCATION" # Values: "global"
+# data_store_id = "YOUR_DATA_STORE_ID"
+
+
+def delete_data_store_sample(
+    project_id: str,
+    location: str,
+    data_store_id: str,
+) -> str:
+    #  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.DataStoreServiceClient(client_options=client_options)
+
+    request = discoveryengine.DeleteDataStoreRequest(
+        # The full resource name of the data store
+        name=client.data_store_path(project_id, location, data_store_id)
+    )
+
+    # Make the request
+    operation = client.delete_data_store(request=request)
+
+    print(f"Operation: {operation.operation.name}")
+
+    return operation.operation.name
+
+
+# [END genappbuilder_delete_data_store]

discoveryengine/requirements.txt

@@ -1,2 +1,2 @@
-google-cloud-discoveryengine==0.11.11
-google-api-core==2.17.1
+google-cloud-discoveryengine==0.11.13
+google-api-core==2.19.0