Skip to content

Commit

Permalink
fix(discoveryengine): Update the samples after GA Launch (#10635)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
holtskinner authored Sep 29, 2023
1 parent 8a3a38e commit 7368358
Show file tree
Hide file tree
Showing 8 changed files with 74 additions and 41 deletions.
8 changes: 4 additions & 4 deletions discoveryengine/documents_sample_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
)

Expand All @@ -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,
)
Expand All @@ -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
1 change: 0 additions & 1 deletion discoveryengine/get_operation_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
28 changes: 18 additions & 10 deletions discoveryengine/import_documents_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand All @@ -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",
)

Expand Down
23 changes: 16 additions & 7 deletions discoveryengine/list_documents_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
6 changes: 2 additions & 4 deletions discoveryengine/list_operations_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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()
Expand Down
1 change: 0 additions & 1 deletion discoveryengine/poll_operation_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
27 changes: 18 additions & 9 deletions discoveryengine/search_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down
21 changes: 16 additions & 5 deletions discoveryengine/search_sample_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down

0 comments on commit 7368358

Please sign in to comment.