From 7368358fe0366980a2a1763ce1ab51845b6a910a Mon Sep 17 00:00:00 2001 From: Holt Skinner <13262395+holtskinner@users.noreply.github.com> Date: Fri, 29 Sep 2023 11:38:48 -0500 Subject: [PATCH] fix(discoveryengine): Update the samples after GA Launch (#10635) * 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 --- discoveryengine/documents_sample_test.py | 8 +++---- discoveryengine/get_operation_sample.py | 1 - discoveryengine/import_documents_sample.py | 28 ++++++++++++++-------- discoveryengine/list_documents_sample.py | 23 ++++++++++++------ discoveryengine/list_operations_sample.py | 6 ++--- discoveryengine/poll_operation_sample.py | 1 - discoveryengine/search_sample.py | 27 ++++++++++++++------- discoveryengine/search_sample_test.py | 21 ++++++++++++---- 8 files changed, 74 insertions(+), 41 deletions(-) diff --git a/discoveryengine/documents_sample_test.py b/discoveryengine/documents_sample_test.py index 153bb8e05449..5caa2833d572 100644 --- 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 index 40d12f591f3a..83fe95a088e6 100644 --- 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 index 21b564f3fcf8..79482455f89c 100644 --- 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 index e9c4b1ee025f..26489b6f8741 100644 --- 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 index f95503caaf72..43e644102b1b 100644 --- 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 index 30bba6e64629..5eaf7edcd6c0 100644 --- 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 index b3505690423e..9745bf6a17f0 100644 --- 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 index 243b64b94110..127b25703a57 100644 --- 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, )