fix: Updates to Vertex AI Search samples after GA (GoogleCloudPlatform#8662)

holtskinner · web-flow · commit 0d69c03a9796 · 2023-09-21T09:14:14.000-07:00
* fix: Updates to Vertex AI Search samples after GA

- Change `searchEngineId` to `dataStoreId`
- Add Endpoint for non-global data stores.

* Updates to comments

* fix: Line length

* Fix syntax error
diff --git a/discoveryengine/src/main/java/discoveryengine/v1/Search.java b/discoveryengine/src/main/java/discoveryengine/v1/Search.java
@@ -21,6 +21,7 @@
 import com.google.cloud.discoveryengine.v1.SearchRequest;
 import com.google.cloud.discoveryengine.v1.SearchResponse;
 import com.google.cloud.discoveryengine.v1.SearchServiceClient;
+import com.google.cloud.discoveryengine.v1.SearchServiceSettings;
 import com.google.cloud.discoveryengine.v1.ServingConfigName;
 import java.io.IOException;
 import java.util.concurrent.ExecutionException;
@@ -30,38 +31,45 @@ public static void main() throws IOException, ExecutionException {
     // TODO(developer): Replace these variables before running the sample.
     // Project ID or project number of the Cloud project you want to use.
     String projectId = "PROJECT_ID";
-    // Location of the search engine. Options: "global"
+    // Location of the data store. Options: "global", "us", "eu"
     String location = "global";
-    // Collection containing the datastore.
+    // Collection containing the data store.
     String collectionId = "default_collection";
-    // Datastore/Search Engine ID.
-    String searchEngineId = "DATA_STORE_ID";
+    // Data store ID.
+    String dataStoreId = "DATA_STORE_ID";
     // Serving configuration. Options: "default_search"
     String servingConfigId = "default_search";
-    // Search Query for the search engine.
+    // Search Query for the data store.
     String searchQuery = "Google";
-    search(projectId, location, collectionId, searchEngineId, servingConfigId, searchQuery);
+    search(projectId, location, collectionId, dataStoreId, servingConfigId, searchQuery);
   }
 
-  /** Performs a search on a given datastore/search engine. */
+  /** Performs a search on a given datastore. */
   public static void search(
       String projectId,
       String location,
       String collectionId,
-      String searchEngineId,
+      String dataStoreId,
       String servingConfigId,
       String searchQuery)
       throws IOException, ExecutionException {
+    // For more information, refer to:
+    // https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store
+    String endpoint = (location.equals("global")) 
+        ? String.format("discoveryengine.googleapis.com:443", location) 
+        : String.format("%s-discoveryengine.googleapis.com:443", location);
+    SearchServiceSettings settings =
+        SearchServiceSettings.newBuilder().setEndpoint(endpoint).build();
     // Initialize client that will be used to send requests. This client only needs to be created
     // once, and can be reused for multiple requests. After completing all of your requests, call
     // the `searchServiceClient.close()` method on the client to safely
     // clean up any remaining background resources.
-    try (SearchServiceClient searchServiceClient = SearchServiceClient.create()) {
+    try (SearchServiceClient searchServiceClient = SearchServiceClient.create(settings)) {
       SearchRequest request =
           SearchRequest.newBuilder()
               .setServingConfig(
                   ServingConfigName.formatProjectLocationCollectionDataStoreServingConfigName(
-                      projectId, location, collectionId, searchEngineId, servingConfigId))
+                      projectId, location, collectionId, dataStoreId, servingConfigId))
               .setQuery(searchQuery)
               .setPageSize(10)
               .build();
diff --git a/discoveryengine/src/test/java/discoveryengine/v1/SearchTest.java b/discoveryengine/src/test/java/discoveryengine/v1/SearchTest.java
@@ -33,7 +33,7 @@ public class SearchTest {
   private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
   private static final String LOCATION = "global";
   private static final String COLLECTION_ID = "default_collection";
-  private static final String SEARCH_ENGINE_ID = "test-search-engine";
+  private static final String DATA_STORE_ID = "test-search-engine";
   private static final String SERVING_CONFIG_ID = "default_search";
   private static final String SEARCH_QUERY = "Google";
 
@@ -64,7 +64,7 @@ public void setUp() {
   @Test
   public void testSearch() throws Exception {
     Search.search(
-        PROJECT_ID, LOCATION, COLLECTION_ID, SEARCH_ENGINE_ID, SERVING_CONFIG_ID, SEARCH_QUERY);
+        PROJECT_ID, LOCATION, COLLECTION_ID, DATA_STORE_ID, SERVING_CONFIG_ID, SEARCH_QUERY);
     String got = bout.toString();
 
     assertThat(got).contains("Response content:");
