From c0dcf74a952b193e7d0f61f04dcbe4ff657cab78 Mon Sep 17 00:00:00 2001 From: Juan Dominguez Date: Thu, 21 Aug 2025 11:56:17 -0400 Subject: [PATCH 1/3] feat: add new GenAI Vais tool sample --- .../java/genai/tools/ToolsVaisWithText.java | 77 +++++++++++++++++++ .../src/test/java/genai/tools/ToolsIT.java | 9 +++ 2 files changed, 86 insertions(+) create mode 100644 genai/snippets/src/main/java/genai/tools/ToolsVaisWithText.java diff --git a/genai/snippets/src/main/java/genai/tools/ToolsVaisWithText.java b/genai/snippets/src/main/java/genai/tools/ToolsVaisWithText.java new file mode 100644 index 00000000000..b3eb0c4e285 --- /dev/null +++ b/genai/snippets/src/main/java/genai/tools/ToolsVaisWithText.java @@ -0,0 +1,77 @@ +/* + * Copyright 2025 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. + */ + +package genai.tools; + +// [START googlegenaisdk_tools_vais_with_txt] + +import com.google.genai.Client; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponse; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.Retrieval; +import com.google.genai.types.Tool; +import com.google.genai.types.VertexAISearch; + +public class ToolsVaisWithText { + + public static void main(String[] args) { + // TODO(developer): Replace these variables before running the sample. + String modelId = "gemini-2.5-flash"; + // Load Data Store ID from Vertex AI Search + // E.g datastoreId = + // "projects/project-id/locations/global/collections/default_collection/dataStores/datastore-id" + String datastoreId = "your-datastore"; + generateContent(modelId, datastoreId); + } + + // Generates text with Vertex AI Search tool + public static String generateContent(String modelId, String datastoreId) { + // Initialize client that will be used to send requests. This client only needs to be created + // once, and can be reused for multiple requests. + try (Client client = + Client.builder() + .location("global") + .vertexAI(true) + .httpOptions(HttpOptions.builder().apiVersion("v1").build()) + .build()) { + + // Set the VertexAI Search tool and the datastore that the model can use to retrieve data from + Tool vaisSearchTool = + Tool.builder() + .retrieval( + Retrieval.builder() + .vertexAiSearch(VertexAISearch.builder().datastore(datastoreId).build()) + .build()) + .build(); + + // Create a GenerateContentConfig and set the Vertex AI Search tool + GenerateContentConfig contentConfig = + GenerateContentConfig.builder().tools(vaisSearchTool).build(); + + GenerateContentResponse response = + client.models.generateContent( + modelId, "How do I make an appointment to renew my driver's license?", contentConfig); + + System.out.print(response.text()); + // Example response: + // The process for making an appointment to renew your driver's license varies depending + // on your location. To provide you with the most accurate instructions... + return response.text(); + } + } +} +// [END googlegenaisdk_tools_vais_with_txt] \ No newline at end of file diff --git a/genai/snippets/src/test/java/genai/tools/ToolsIT.java b/genai/snippets/src/test/java/genai/tools/ToolsIT.java index b18fb4a9bc1..f6bd21081f4 100644 --- a/genai/snippets/src/test/java/genai/tools/ToolsIT.java +++ b/genai/snippets/src/test/java/genai/tools/ToolsIT.java @@ -32,6 +32,7 @@ public class ToolsIT { private static final String GEMINI_FLASH = "gemini-2.5-flash"; + private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT"); private ByteArrayOutputStream bout; private PrintStream out; @@ -79,4 +80,12 @@ public void testGenerateContentWithFunctionDescription() { assertThat(response).contains("copies_sold=350000"); assertThat(response).contains("album_name=Echoes of the Night"); } + + @Test + public void testToolsVaisWithText() { + String datastore = "projects/" + PROJECT_ID + "/locations/global/" + + "collections/default_collection/dataStores/grounding-test-datastore"; + String response = ToolsVaisWithText.generateContent(GEMINI_FLASH, datastore); + assertThat(response).isNotEmpty(); + } } From 43143ba93c9b1bdf03dd31e24bbf3f7a9d9b0313 Mon Sep 17 00:00:00 2001 From: Juan Dominguez Date: Thu, 21 Aug 2025 12:11:30 -0400 Subject: [PATCH 2/3] refactor: change the way the datastore string was consctructed for better readibility --- genai/snippets/src/test/java/genai/tools/ToolsIT.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/genai/snippets/src/test/java/genai/tools/ToolsIT.java b/genai/snippets/src/test/java/genai/tools/ToolsIT.java index f6bd21081f4..fb5ddbdcde6 100644 --- a/genai/snippets/src/test/java/genai/tools/ToolsIT.java +++ b/genai/snippets/src/test/java/genai/tools/ToolsIT.java @@ -83,9 +83,14 @@ public void testGenerateContentWithFunctionDescription() { @Test public void testToolsVaisWithText() { - String datastore = "projects/" + PROJECT_ID + "/locations/global/" - + "collections/default_collection/dataStores/grounding-test-datastore"; + + String datastore = + String.format( + "projects/%s/locations/global/collections/default_collection/" + + "dataStores/grounding-test-datastore", + PROJECT_ID); String response = ToolsVaisWithText.generateContent(GEMINI_FLASH, datastore); assertThat(response).isNotEmpty(); + } } From 11b51af2cffb4cd003259608158c76998cc1781d Mon Sep 17 00:00:00 2001 From: Juan Dominguez Date: Wed, 17 Sep 2025 18:16:57 -0300 Subject: [PATCH 3/3] include mockito in pom and change tool vais test to mock it instead of using an actual datastore --- genai/snippets/pom.xml | 6 ++ .../src/test/java/genai/tools/ToolsIT.java | 59 ++++++++++++++++--- 2 files changed, 56 insertions(+), 9 deletions(-) diff --git a/genai/snippets/pom.xml b/genai/snippets/pom.xml index b93e29cebe3..32015ec0bf4 100644 --- a/genai/snippets/pom.xml +++ b/genai/snippets/pom.xml @@ -59,6 +59,12 @@ test 4.13.2 + + org.mockito + mockito-core + 5.19.0 + test + com.google.truth truth diff --git a/genai/snippets/src/test/java/genai/tools/ToolsIT.java b/genai/snippets/src/test/java/genai/tools/ToolsIT.java index 71365e4447f..1bbf109580d 100644 --- a/genai/snippets/src/test/java/genai/tools/ToolsIT.java +++ b/genai/snippets/src/test/java/genai/tools/ToolsIT.java @@ -18,16 +18,31 @@ import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertWithMessage; - +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.RETURNS_SELF; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.genai.Client; +import com.google.genai.Models; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponse; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.PrintStream; +import java.lang.reflect.Field; import org.junit.After; import org.junit.Before; import org.junit.BeforeClass; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.MockedStatic; + @RunWith(JUnit4.class) public class ToolsIT { @@ -106,15 +121,41 @@ public void testToolsGoogleSearchWithText() { } @Test - public void testToolsVaisWithText() { + public void testToolsVaisWithText() throws NoSuchFieldException, IllegalAccessException { + String response = "The process for making an appointment to renew your driver's license" + + " varies depending on your location."; String datastore = - String.format( - "projects/%s/locations/global/collections/default_collection/" - + "dataStores/grounding-test-datastore", - PROJECT_ID); - String response = ToolsVaisWithText.generateContent(GEMINI_FLASH, datastore); - assertThat(response).isNotEmpty(); - + String.format( + "projects/%s/locations/global/collections/default_collection/" + + "dataStores/grounding-test-datastore", + PROJECT_ID); + + Client.Builder mockedBuilder = mock(Client.Builder.class, RETURNS_SELF); + Client mockedClient = mock(Client.class); + Models mockedModels = mock(Models.class); + GenerateContentResponse mockedResponse = mock(GenerateContentResponse.class); + + try (MockedStatic mockedStatic = mockStatic(Client.class)) { + mockedStatic.when(Client::builder).thenReturn(mockedBuilder); + when(mockedBuilder.build()).thenReturn(mockedClient); + + // Using reflection because 'models' is a final field and cannot be mockable directly + Field field = Client.class.getDeclaredField("models"); + field.setAccessible(true); + field.set(mockedClient, mockedModels); + + when(mockedClient.models.generateContent( + anyString(), anyString(), any(GenerateContentConfig.class))) + .thenReturn(mockedResponse); + when(mockedResponse.text()).thenReturn(response); + + String generatedResponse = ToolsVaisWithText.generateContent(GEMINI_FLASH, datastore); + + verify(mockedClient.models, times(1)) + .generateContent(anyString(), anyString(), any(GenerateContentConfig.class)); + assertThat(generatedResponse).isNotEmpty(); + assertThat(response).isEqualTo(generatedResponse); + } } } \ No newline at end of file