feat: add support for gemini models in VertexAiRagRetrieval

tilgalas · copybara-github · commit 924fb7174855 · 2026-03-06T05:43:20.000-08:00
PiperOrigin-RevId: 879585211
diff --git a/core/pom.xml b/core/pom.xml
@@ -209,10 +209,29 @@
         <artifactId>maven-surefire-plugin</artifactId>
         <executions>
           <execution>
+            <!-- run all tests with standard env -->
             <id>basic</id>
             <goals>
               <goal>test</goal>
             </goals>
+            <configuration>
+              <environmentVariables>
+                <GOOGLE_GENAI_USE_VERTEXAI>false</GOOGLE_GENAI_USE_VERTEXAI>
+              </environmentVariables>
+            </configuration>
+          </execution>
+          <execution>
+            <id>vertex-ai-rag-retrieval</id>
+            <goals>
+              <goal>test</goal>
+            </goals>
+            <configuration>
+              <environmentVariables>
+                <GOOGLE_GENAI_USE_VERTEXAI>true</GOOGLE_GENAI_USE_VERTEXAI>
+              </environmentVariables>
+              <!-- run a second variant of the following tests -->
+              <test>VertexAiRagRetrievalTest#processLlmRequest_gemini2Model_addVertexRagStoreToConfig, VertexAiRagRetrievalTest#processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig</test>
+            </configuration>
           </execution>
           <execution>
             <id>apigee-llm</id>
diff --git a/core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java b/core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java
@@ -20,6 +20,7 @@
 
 import com.google.adk.models.LlmRequest;
 import com.google.adk.tools.ToolContext;
+import com.google.adk.utils.ModelNameUtils;
 import com.google.cloud.aiplatform.v1.RagContexts;
 import com.google.cloud.aiplatform.v1.RagQuery;
 import com.google.cloud.aiplatform.v1.RetrieveContextsRequest;
@@ -105,10 +106,9 @@ public VertexAiRagRetrieval(
   public Completable processLlmRequest(
       LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {
     LlmRequest llmRequest = llmRequestBuilder.build();
-    // Use Gemini built-in Vertex AI RAG tool for Gemini 2 models or when using Vertex AI API Model
+    // Use Gemini built-in Vertex AI RAG tool for Gemini models when using Vertex AI API Model
     boolean useVertexAi = Boolean.parseBoolean(System.getenv("GOOGLE_GENAI_USE_VERTEXAI"));
-    if (useVertexAi
-        && (llmRequest.model().isPresent() && llmRequest.model().get().startsWith("gemini-2"))) {
+    if (useVertexAi && llmRequest.model().filter(ModelNameUtils::isGeminiModel).isPresent()) {
       GenerateContentConfig config =
           llmRequest.config().orElseGet(() -> GenerateContentConfig.builder().build());
       ImmutableList.Builder<Tool> toolsBuilder = ImmutableList.builder();
diff --git a/core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java b/core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java
@@ -208,7 +208,7 @@ public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
             "projects/test-project/locations/us-central1",
             ragResources,
             vectorDistanceThreshold);
-    LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("gemini-1-pro");
+    LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("other-model");
     ToolContext toolContext = buildToolContext();
     GenerateContentConfig initialConfig = GenerateContentConfig.builder().build();
     llmRequestBuilder.config(initialConfig);
