fix(discoveryengine): Update the samples after GA Launch (GoogleCloudPlatform#10635)

holtskinner · web-flow · commit 7368358fe036 · 2023-09-29T11:38:48.000-05:00
* fix: Add API Endpoint to Samples

- Change `search_engine_id` to `data_store_id`

* fix: Fixes to document samples

- Use correct client (copy/paste mishap)
- Remove erroneous import of annotation

* Add explicit optionals
diff --git a/discoveryengine/documents_sample_test.py b/discoveryengine/documents_sample_test.py
@@ -20,7 +20,7 @@
 
 project_id = os.environ["GOOGLE_CLOUD_PROJECT"]
 location = "global"
-search_engine_id = "test-structured-data-engine"
+data_store_id = "test-structured-data-engine"
 gcs_uri = "gs://cloud-samples-data/gen-app-builder/search/empty.json"
 
 # Empty Dataset
@@ -32,7 +32,7 @@ def test_import_documents_gcs():
     operation_name = import_documents_sample.import_documents_sample(
         project_id=project_id,
         location=location,
-        search_engine_id=search_engine_id,
+        data_store_id=data_store_id,
         gcs_uri=gcs_uri,
     )
 
@@ -43,7 +43,7 @@ def test_import_documents_bigquery():
     operation_name = import_documents_sample.import_documents_sample(
         project_id=project_id,
         location=location,
-        search_engine_id=search_engine_id,
+        data_store_id=data_store_id,
         bigquery_dataset=bigquery_dataset,
         bigquery_table=bigquery_table,
     )
@@ -55,7 +55,7 @@ def test_list_documents():
     response = list_documents_sample.list_documents_sample(
         project_id=project_id,
         location=location,
-        search_engine_id=search_engine_id,
+        data_store_id=data_store_id,
     )
 
     assert response
diff --git a/discoveryengine/get_operation_sample.py b/discoveryengine/get_operation_sample.py
@@ -15,7 +15,6 @@
 
 # [START genappbuilder_get_operation]
 from google.cloud import discoveryengine
-
 from google.longrunning import operations_pb2
 
 # TODO(developer): Uncomment these variables before running the sample.
diff --git a/discoveryengine/import_documents_sample.py b/discoveryengine/import_documents_sample.py
@@ -14,15 +14,15 @@
 #
 
 # [START genappbuilder_import_documents]
-from __future__ import annotations
-
+from typing import Optional
 
+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"
-# search_engine_id = "YOUR_SEARCH_ENGINE_ID"
+# data_store_id = "YOUR_DATA_STORE_ID"
 
 # Must specify either `gcs_uri` or (`bigquery_dataset` and `bigquery_table`)
 # Format: `gs://bucket/directory/object.json` or `gs://bucket/directory/*.json`
@@ -34,20 +34,28 @@
 def import_documents_sample(
     project_id: str,
     location: str,
-    search_engine_id: str,
-    gcs_uri: str | None = None,
-    bigquery_dataset: str | None = None,
-    bigquery_table: str | None = None,
+    data_store_id: str,
+    gcs_uri: Optional[str] = None,
+    bigquery_dataset: Optional[str] = None,
+    bigquery_table: Optional[str] = None,
 ) -> 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.DocumentServiceClient()
+    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}/branches/{branch}
+    # e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}
     parent = client.branch_path(
         project=project_id,
         location=location,
-        data_store=search_engine_id,
+        data_store=data_store_id,
         branch="default_branch",
     )
 
diff --git a/discoveryengine/list_documents_sample.py b/discoveryengine/list_documents_sample.py
@@ -15,30 +15,39 @@
 # [START genappbuilder_list_documents]
 from typing import Any
 
+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"
-# search_engine_id = "YOUR_SEARCH_ENGINE_ID"
+# location = "YOUR_LOCATION"            # Values: "global", "us", "eu"
+# data_store_id = "YOUR_DATA_STORE_ID"
 
 
-def list_documents_sample(project_id: str, location: str, search_engine_id: str) -> Any:
+def list_documents_sample(project_id: str, location: str, data_store_id: str) -> Any:
+    #  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 = discoveryengine.DocumentServiceClient(client_options=client_options)
 
     # The full resource name of the search engine branch.
-    # e.g. projects/{project}/locations/{location}/dataStores/{data_store}/branches/{branch}
+    # e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}
     parent = client.branch_path(
         project=project_id,
         location=location,
-        data_store=search_engine_id,
+        data_store=data_store_id,
         branch="default_branch",
     )
 
     response = client.list_documents(parent=parent)
 
-    print(f"Documents in {search_engine_id}:")
+    print(f"Documents in {data_store_id}:")
     for result in response:
         print(result)
 
diff --git a/discoveryengine/list_operations_sample.py b/discoveryengine/list_operations_sample.py
@@ -13,11 +13,9 @@
 # limitations under the License.
 
 # [START genappbuilder_list_operations]
-from __future__ import annotations
-
+from typing import Optional
 
 from google.cloud import discoveryengine
-
 from google.longrunning import operations_pb2
 
 # TODO(developer): Uncomment these variables before running the sample.
@@ -33,7 +31,7 @@ def list_operations_sample(
     project_id: str,
     location: str,
     search_engine_id: str,
-    operations_filter: str | None = None,
+    operations_filter: Optional[str] = None,
 ) -> operations_pb2.ListOperationsResponse:
     # Create a client
     client = discoveryengine.DocumentServiceClient()
diff --git a/discoveryengine/poll_operation_sample.py b/discoveryengine/poll_operation_sample.py
@@ -17,7 +17,6 @@
 from time import sleep
 
 from google.cloud import discoveryengine
-
 from google.longrunning import operations_pb2
 
 # TODO(developer): Uncomment these variables before running the sample.
diff --git a/discoveryengine/search_sample.py b/discoveryengine/search_sample.py
@@ -16,35 +16,44 @@
 # [START genappbuilder_search]
 from typing import List
 
+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"
-# search_engine_id = "YOUR_SEARCH_ENGINE_ID"
-# serving_config_id = "default_config"          # Values: "default_config"
+# location = "YOUR_LOCATION"                    # Values: "global", "us", "eu"
+# data_store_id = "YOUR_DATA_STORE_ID"
 # search_query = "YOUR_SEARCH_QUERY"
 
 
 def search_sample(
     project_id: str,
     location: str,
-    search_engine_id: str,
-    serving_config_id: str,
+    data_store_id: str,
     search_query: str,
 ) -> List[discoveryengine.SearchResponse.SearchResult]:
+    #  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.SearchServiceClient()
+    client = discoveryengine.SearchServiceClient(client_options=client_options)
 
     # The full resource name of the search engine serving config
-    # e.g. projects/{project_id}/locations/{location}
+    # e.g. projects/{project_id}/locations/{location}/dataStores/{data_store_id}/servingConfigs/{serving_config_id}
     serving_config = client.serving_config_path(
         project=project_id,
         location=location,
-        data_store=search_engine_id,
-        serving_config=serving_config_id,
+        data_store=data_store_id,
+        serving_config="default_config",
     )
 
+    # Refer to the SearchRequest reference for all supported fields:
+    # https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1.types.SearchRequest
     request = discoveryengine.SearchRequest(
         serving_config=serving_config, query=search_query, page_size=10
     )
diff --git a/discoveryengine/search_sample_test.py b/discoveryengine/search_sample_test.py
@@ -18,18 +18,29 @@
 from discoveryengine import search_sample
 
 project_id = os.environ["GOOGLE_CLOUD_PROJECT"]
-location = "global"
-search_engine_id = "test-search-engine_1689960780551"
-serving_config_id = "default_config"
 search_query = "Google"
 
 
 def test_search():
+    location = "global"
+    data_store_id = "test-search-engine_1689960780551"
     response = search_sample.search_sample(
         project_id=project_id,
         location=location,
-        search_engine_id=search_engine_id,
-        serving_config_id=serving_config_id,
+        data_store_id=data_store_id,
+        search_query=search_query,
+    )
+
+    assert response
+
+
+def test_search_eu_endpoint():
+    location = "eu"
+    data_store_id = "alphabet-earnings-reports-eu"
+    response = search_sample.search_sample(
+        project_id=project_id,
+        location=location,
+        data_store_id=data_store_id,
         search_query=search_query,
     )
 
