feat: Add VertexAiSearchTool and AgentTools for search

PiperOrigin-RevId: 848220298

Author: google-genai-bot

core/src/main/java/com/google/adk/tools/AgentTool.java

@@ -28,6 +28,7 @@
 import com.google.adk.events.Event;
 import com.google.adk.runner.InMemoryRunner;
 import com.google.adk.runner.Runner;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.genai.types.Content;
@@ -76,6 +77,11 @@ protected AgentTool(BaseAgent agent, boolean skipSummarization) {
     this.skipSummarization = skipSummarization;
   }
 
+  @VisibleForTesting
+  BaseAgent getAgent() {
+    return agent;
+  }
+
   @Override
   public Optional<FunctionDeclaration> declaration() {
     FunctionDeclaration.Builder builder =

core/src/main/java/com/google/adk/tools/GoogleSearchAgentTool.java

@@ -0,0 +1,50 @@
+/*
+ * 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 com.google.adk.tools;
+
+import com.google.adk.agents.LlmAgent;
+import com.google.adk.models.BaseLlm;
+import com.google.common.collect.ImmutableList;
+
+/**
+ * A tool that wraps a sub-agent that only uses google_search tool.
+ *
+ * <p>This is a workaround to support using google_search tool with other tools. TODO(b/448114567):
+ * Remove once the workaround is no longer needed.
+ */
+public class GoogleSearchAgentTool extends AgentTool {
+
+  public static GoogleSearchAgentTool create(BaseLlm model) {
+    LlmAgent googleSearchAgent =
+        LlmAgent.builder()
+            .name("google_search_agent")
+            .model(model)
+            .description("An agent for performing Google search using the `google_search` tool")
+            .instruction(
+                "        You are a specialized Google search agent.\n"
+                    + "\n"
+                    + "        When given a search query, use the `google_search` tool to find the"
+                    + " related information.")
+            .tools(ImmutableList.of(GoogleSearchTool.INSTANCE))
+            .build();
+    return new GoogleSearchAgentTool(googleSearchAgent);
+  }
+
+  protected GoogleSearchAgentTool(LlmAgent agent) {
+    super(agent, false);
+  }
+}

core/src/main/java/com/google/adk/tools/VertexAiSearchAgentTool.java

@@ -0,0 +1,51 @@
+/*
+ * 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 com.google.adk.tools;
+
+import com.google.adk.agents.LlmAgent;
+import com.google.adk.models.BaseLlm;
+import com.google.common.collect.ImmutableList;
+
+/**
+ * A tool that wraps a sub-agent that only uses vertex_ai_search tool.
+ *
+ * <p>This is a workaround to support using {@link VertexAiSearchTool} tool with other tools.
+ */
+public class VertexAiSearchAgentTool extends AgentTool {
+
+  public static VertexAiSearchAgentTool create(
+      BaseLlm model, VertexAiSearchTool vertexAiSearchTool) {
+    LlmAgent vertexAiSearchAgent =
+        LlmAgent.builder()
+            .name("vertex_ai_search_agent")
+            .model(model)
+            .description(
+                "An agent for performing Vertex AI search using the `vertex_ai_search` tool")
+            .instruction(
+                "        You are a specialized Vertex AI search agent.\n"
+                    + "\n"
+                    + "        When given a search query, use the `vertex_ai_search` tool to find"
+                    + " the related information.")
+            .tools(ImmutableList.of(vertexAiSearchTool))
+            .build();
+    return new VertexAiSearchAgentTool(vertexAiSearchAgent);
+  }
+
+  protected VertexAiSearchAgentTool(LlmAgent agent) {
+    super(agent, false);
+  }
+}

core/src/main/java/com/google/adk/tools/VertexAiSearchTool.java

@@ -0,0 +1,133 @@
+/*
+ * 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 com.google.adk.tools;
+
+import com.google.adk.models.LlmRequest;
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
+import com.google.genai.types.GenerateContentConfig;
+import com.google.genai.types.Retrieval;
+import com.google.genai.types.Tool;
+import com.google.genai.types.VertexAISearch;
+import com.google.genai.types.VertexAISearchDataStoreSpec;
+import io.reactivex.rxjava3.core.Completable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * A built-in tool using Vertex AI Search.
+ *
+ * <p>This tool can be configured with either a {@code dataStoreId} (the Vertex AI search data store
+ * resource ID) or a {@code searchEngineId} (the Vertex AI search engine resource ID).
+ */
+@AutoValue
+public abstract class VertexAiSearchTool extends BaseTool {
+  public abstract Optional<String> dataStoreId();
+
+  public abstract Optional<List<VertexAISearchDataStoreSpec>> dataStoreSpecs();
+
+  public abstract Optional<String> searchEngineId();
+
+  public abstract Optional<String> filter();
+
+  public abstract Optional<Integer> maxResults();
+
+  public abstract Optional<String> project();
+
+  public abstract Optional<String> location();
+
+  public abstract Optional<String> dataStore();
+
+  public static Builder builder() {
+    return new AutoValue_VertexAiSearchTool.Builder();
+  }
+
+  VertexAiSearchTool() {
+    super("vertex_ai_search", "vertex_ai_search");
+  }
+
+  @Override
+  public Completable processLlmRequest(
+      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {
+    LlmRequest llmRequest = llmRequestBuilder.build();
+
+    VertexAISearch.Builder vertexAiSearchBuilder = VertexAISearch.builder();
+    dataStoreId().ifPresent(vertexAiSearchBuilder::datastore);
+    searchEngineId().ifPresent(vertexAiSearchBuilder::engine);
+    filter().ifPresent(vertexAiSearchBuilder::filter);
+    maxResults().ifPresent(vertexAiSearchBuilder::maxResults);
+    dataStoreSpecs().ifPresent(vertexAiSearchBuilder::dataStoreSpecs);
+
+    Tool retrievalTool =
+        Tool.builder()
+            .retrieval(Retrieval.builder().vertexAiSearch(vertexAiSearchBuilder.build()).build())
+            .build();
+
+    List<Tool> currentTools = new ArrayList<>();
+    if (llmRequest.config().isPresent() && llmRequest.config().get().tools().isPresent()) {
+      currentTools.addAll(llmRequest.config().get().tools().get());
+    }
+    currentTools.add(retrievalTool);
+
+    GenerateContentConfig newConfig =
+        llmRequest
+            .config()
+            .map(GenerateContentConfig::toBuilder)
+            .orElse(GenerateContentConfig.builder())
+            .tools(ImmutableList.copyOf(currentTools))
+            .build();
+    llmRequestBuilder.config(newConfig);
+    return Completable.complete();
+  }
+
+  /** Builder for {@link VertexAiSearchTool}. */
+  @AutoValue.Builder
+  public abstract static class Builder {
+    public abstract Builder dataStoreId(String dataStoreId);
+
+    public abstract Builder dataStoreSpecs(List<VertexAISearchDataStoreSpec> dataStoreSpecs);
+
+    public abstract Builder searchEngineId(String searchEngineId);
+
+    public abstract Builder filter(String filter);
+
+    public abstract Builder maxResults(Integer maxResults);
+
+    public abstract Builder project(String project);
+
+    public abstract Builder location(String location);
+
+    public abstract Builder dataStore(String dataStore);
+
+    abstract VertexAiSearchTool autoBuild();
+
+    public final VertexAiSearchTool build() {
+      VertexAiSearchTool tool = autoBuild();
+      if ((tool.dataStoreId().isEmpty() && tool.searchEngineId().isEmpty())
+          || (tool.dataStoreId().isPresent() && tool.searchEngineId().isPresent())) {
+        throw new IllegalArgumentException(
+            "Either dataStoreId or searchEngineId must be specified.");
+      }
+      if (tool.dataStoreSpecs().isPresent() && tool.searchEngineId().isEmpty()) {
+        throw new IllegalArgumentException(
+            "searchEngineId must be specified if dataStoreSpecs is specified.");
+      }
+      return tool;
+    }
+  }
+}

core/src/test/java/com/google/adk/tools/GoogleSearchAgentToolTest.java

@@ -0,0 +1,22 @@
+package com.google.adk.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.adk.agents.LlmAgent;
+import com.google.adk.testing.TestLlm;
+import com.google.common.collect.ImmutableList;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public final class GoogleSearchAgentToolTest {
+
+  @Test
+  public void create_createsAgent() {
+    GoogleSearchAgentTool tool = GoogleSearchAgentTool.create(new TestLlm(ImmutableList.of()));
+    assertThat(((LlmAgent) tool.getAgent()).name()).isEqualTo("google_search_agent");
+    assertThat(((LlmAgent) tool.getAgent()).tools().blockingGet())
+        .containsExactly(GoogleSearchTool.INSTANCE);
+  }
+}

core/src/test/java/com/google/adk/tools/VertexAiSearchAgentToolTest.java

@@ -0,0 +1,25 @@
+package com.google.adk.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.adk.agents.LlmAgent;
+import com.google.adk.testing.TestLlm;
+import com.google.common.collect.ImmutableList;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public final class VertexAiSearchAgentToolTest {
+
+  @Test
+  public void create_createsAgent() {
+    VertexAiSearchTool vertexAiSearchTool =
+        VertexAiSearchTool.builder().searchEngineId("test-engine").build();
+    VertexAiSearchAgentTool tool =
+        VertexAiSearchAgentTool.create(new TestLlm(ImmutableList.of()), vertexAiSearchTool);
+    assertThat(((LlmAgent) tool.getAgent()).name()).isEqualTo("vertex_ai_search_agent");
+    assertThat(((LlmAgent) tool.getAgent()).tools().blockingGet())
+        .containsExactly(vertexAiSearchTool);
+  }
+}