Skip to content

Commit b48b194

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
feat: Add VertexAiSearchTool and AgentTools for search
PiperOrigin-RevId: 851332413
1 parent f6ff2f2 commit b48b194

File tree

7 files changed

+433
-0
lines changed

7 files changed

+433
-0
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.adk.events.Event;
2929
import com.google.adk.runner.InMemoryRunner;
3030
import com.google.adk.runner.Runner;
31+
import com.google.common.annotations.VisibleForTesting;
3132
import com.google.common.collect.ImmutableList;
3233
import com.google.common.collect.ImmutableMap;
3334
import com.google.genai.types.Content;
@@ -76,6 +77,11 @@ protected AgentTool(BaseAgent agent, boolean skipSummarization) {
7677
this.skipSummarization = skipSummarization;
7778
}
7879

80+
@VisibleForTesting
81+
BaseAgent getAgent() {
82+
return agent;
83+
}
84+
7985
@Override
8086
public Optional<FunctionDeclaration> declaration() {
8187
FunctionDeclaration.Builder builder =

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.tools;
18+
19+
import com.google.adk.agents.LlmAgent;
20+
import com.google.adk.models.BaseLlm;
21+
import com.google.common.collect.ImmutableList;
22+
23+
/**
24+
* A tool that wraps a sub-agent that only uses google_search tool.
25+
*
26+
* <p>This is a workaround to support using google_search tool with other tools. TODO(b/448114567):
27+
* Remove once the workaround is no longer needed.
28+
*/
29+
public class GoogleSearchAgentTool extends AgentTool {
30+
31+
public static GoogleSearchAgentTool create(BaseLlm model) {
32+
LlmAgent googleSearchAgent =
33+
LlmAgent.builder()
34+
.name("google_search_agent")
35+
.model(model)
36+
.description("An agent for performing Google search using the `google_search` tool")
37+
.instruction(
38+
" You are a specialized Google search agent.\n"
39+
+ "\n"
40+
+ " When given a search query, use the `google_search` tool to find the"
41+
+ " related information.")
42+
.tools(ImmutableList.of(GoogleSearchTool.INSTANCE))
43+
.build();
44+
return new GoogleSearchAgentTool(googleSearchAgent);
45+
}
46+
47+
protected GoogleSearchAgentTool(LlmAgent agent) {
48+
super(agent, false);
49+
}
50+
}

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.tools;
18+
19+
import com.google.adk.agents.LlmAgent;
20+
import com.google.adk.models.BaseLlm;
21+
import com.google.common.collect.ImmutableList;
22+
23+
/**
24+
* A tool that wraps a sub-agent that only uses vertex_ai_search tool.
25+
*
26+
* <p>This is a workaround to support using {@link VertexAiSearchTool} tool with other tools.
27+
*/
28+
public class VertexAiSearchAgentTool extends AgentTool {
29+
30+
public static VertexAiSearchAgentTool create(
31+
BaseLlm model, VertexAiSearchTool vertexAiSearchTool) {
32+
LlmAgent vertexAiSearchAgent =
33+
LlmAgent.builder()
34+
.name("vertex_ai_search_agent")
35+
.model(model)
36+
.description(
37+
"An agent for performing Vertex AI search using the `vertex_ai_search` tool")
38+
.instruction(
39+
" You are a specialized Vertex AI search agent.\n"
40+
+ "\n"
41+
+ " When given a search query, use the `vertex_ai_search` tool to find"
42+
+ " the related information.")
43+
.tools(ImmutableList.of(vertexAiSearchTool))
44+
.build();
45+
return new VertexAiSearchAgentTool(vertexAiSearchAgent);
46+
}
47+
48+
protected VertexAiSearchAgentTool(LlmAgent agent) {
49+
super(agent, false);
50+
}
51+
}

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

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.tools;
18+
19+
import com.google.adk.models.LlmRequest;
20+
import com.google.auto.value.AutoValue;
21+
import com.google.common.collect.ImmutableList;
22+
import com.google.genai.types.GenerateContentConfig;
23+
import com.google.genai.types.Retrieval;
24+
import com.google.genai.types.Tool;
25+
import com.google.genai.types.VertexAISearch;
26+
import com.google.genai.types.VertexAISearchDataStoreSpec;
27+
import io.reactivex.rxjava3.core.Completable;
28+
import java.util.ArrayList;
29+
import java.util.List;
30+
import java.util.Optional;
31+
32+
/**
33+
* A built-in tool using Vertex AI Search.
34+
*
35+
* <p>This tool can be configured with either a {@code dataStoreId} (the Vertex AI search data store
36+
* resource ID) or a {@code searchEngineId} (the Vertex AI search engine resource ID).
37+
*/
38+
@AutoValue
39+
public abstract class VertexAiSearchTool extends BaseTool {
40+
public abstract Optional<String> dataStoreId();
41+
42+
public abstract Optional<List<VertexAISearchDataStoreSpec>> dataStoreSpecs();
43+
44+
public abstract Optional<String> searchEngineId();
45+
46+
public abstract Optional<String> filter();
47+
48+
public abstract Optional<Integer> maxResults();
49+
50+
public abstract Optional<String> project();
51+
52+
public abstract Optional<String> location();
53+
54+
public abstract Optional<String> dataStore();
55+
56+
public static Builder builder() {
57+
return new AutoValue_VertexAiSearchTool.Builder();
58+
}
59+
60+
VertexAiSearchTool() {
61+
super("vertex_ai_search", "vertex_ai_search");
62+
}
63+
64+
@Override
65+
public Completable processLlmRequest(
66+
LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {
67+
LlmRequest llmRequest = llmRequestBuilder.build();
68+
69+
if (llmRequest
70+
.model()
71+
.map(model -> !model.startsWith("gemini") && !model.contains("/gemini"))
72+
.orElse(false)) {
73+
return Completable.error(
74+
new IllegalArgumentException(
75+
"Vertex AI Search tool is only supported for Gemini models."));
76+
}
77+
78+
VertexAISearch.Builder vertexAiSearchBuilder = VertexAISearch.builder();
79+
dataStoreId().ifPresent(vertexAiSearchBuilder::datastore);
80+
searchEngineId().ifPresent(vertexAiSearchBuilder::engine);
81+
filter().ifPresent(vertexAiSearchBuilder::filter);
82+
maxResults().ifPresent(vertexAiSearchBuilder::maxResults);
83+
dataStoreSpecs().ifPresent(vertexAiSearchBuilder::dataStoreSpecs);
84+
85+
Tool retrievalTool =
86+
Tool.builder()
87+
.retrieval(Retrieval.builder().vertexAiSearch(vertexAiSearchBuilder.build()).build())
88+
.build();
89+
90+
List<Tool> currentTools =
91+
new ArrayList<>(
92+
llmRequest.config().flatMap(GenerateContentConfig::tools).orElse(ImmutableList.of()));
93+
currentTools.add(retrievalTool);
94+
95+
GenerateContentConfig newConfig =
96+
llmRequest
97+
.config()
98+
.map(GenerateContentConfig::toBuilder)
99+
.orElse(GenerateContentConfig.builder())
100+
.tools(ImmutableList.copyOf(currentTools))
101+
.build();
102+
llmRequestBuilder.config(newConfig);
103+
return Completable.complete();
104+
}
105+
106+
/** Builder for {@link VertexAiSearchTool}. */
107+
@AutoValue.Builder
108+
public abstract static class Builder {
109+
public abstract Builder dataStoreId(String dataStoreId);
110+
111+
public abstract Builder dataStoreSpecs(List<VertexAISearchDataStoreSpec> dataStoreSpecs);
112+
113+
public abstract Builder searchEngineId(String searchEngineId);
114+
115+
public abstract Builder filter(String filter);
116+
117+
public abstract Builder maxResults(Integer maxResults);
118+
119+
public abstract Builder project(String project);
120+
121+
public abstract Builder location(String location);
122+
123+
public abstract Builder dataStore(String dataStore);
124+
125+
abstract VertexAiSearchTool autoBuild();
126+
127+
public final VertexAiSearchTool build() {
128+
VertexAiSearchTool tool = autoBuild();
129+
if ((tool.dataStoreId().isEmpty() && tool.searchEngineId().isEmpty())
130+
|| (tool.dataStoreId().isPresent() && tool.searchEngineId().isPresent())) {
131+
throw new IllegalArgumentException(
132+
"Either dataStoreId or searchEngineId must be specified.");
133+
}
134+
if (tool.dataStoreSpecs().isPresent() && tool.searchEngineId().isEmpty()) {
135+
throw new IllegalArgumentException(
136+
"searchEngineId must be specified if dataStoreSpecs is specified.");
137+
}
138+
return tool;
139+
}
140+
}
141+
}

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package com.google.adk.tools;
2+
3+
import static com.google.common.truth.Truth.assertThat;
4+
5+
import com.google.adk.agents.LlmAgent;
6+
import com.google.adk.testing.TestLlm;
7+
import com.google.common.collect.ImmutableList;
8+
import org.junit.Test;
9+
import org.junit.runner.RunWith;
10+
import org.junit.runners.JUnit4;
11+
12+
@RunWith(JUnit4.class)
13+
public final class GoogleSearchAgentToolTest {
14+
15+
@Test
16+
public void create_createsAgent() {
17+
GoogleSearchAgentTool tool = GoogleSearchAgentTool.create(new TestLlm(ImmutableList.of()));
18+
assertThat(tool.getAgent().name()).isEqualTo("google_search_agent");
19+
assertThat(((LlmAgent) tool.getAgent()).tools().blockingGet())
20+
.containsExactly(GoogleSearchTool.INSTANCE);
21+
}
22+
}

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package com.google.adk.tools;
2+
3+
import static com.google.common.truth.Truth.assertThat;
4+
5+
import com.google.adk.agents.LlmAgent;
6+
import com.google.adk.testing.TestLlm;
7+
import com.google.common.collect.ImmutableList;
8+
import org.junit.Test;
9+
import org.junit.runner.RunWith;
10+
import org.junit.runners.JUnit4;
11+
12+
@RunWith(JUnit4.class)
13+
public final class VertexAiSearchAgentToolTest {
14+
15+
@Test
16+
public void create_createsAgent() {
17+
VertexAiSearchTool vertexAiSearchTool =
18+
VertexAiSearchTool.builder().searchEngineId("test-engine").build();
19+
VertexAiSearchAgentTool tool =
20+
VertexAiSearchAgentTool.create(new TestLlm(ImmutableList.of()), vertexAiSearchTool);
21+
assertThat(tool.getAgent().name()).isEqualTo("vertex_ai_search_agent");
22+
assertThat(((LlmAgent) tool.getAgent()).tools().blockingGet())
23+
.containsExactly(vertexAiSearchTool);
24+
}
25+
}

0 commit comments

Comments
 (0)