diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt index dfbbe413..09e75d91 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt @@ -40,7 +40,10 @@ import com.google.adk.kt.tools.BaseTool import com.google.adk.kt.tools.GoogleSearchAgentTool import com.google.adk.kt.tools.GoogleSearchTool import com.google.adk.kt.tools.ToolContext +import com.google.adk.kt.tools.VertexAiSearchAgentTool +import com.google.adk.kt.tools.VertexAiSearchTool import com.google.adk.kt.tools.createGoogleSearchAgent +import com.google.adk.kt.tools.createVertexAiSearchAgent import com.google.adk.kt.types.UsageMetadata import kotlin.time.Clock import kotlinx.coroutines.CancellationException @@ -184,6 +187,8 @@ internal class LlmAgentTurn( when { tool is GoogleSearchTool && tool.bypassMultiToolsLimit -> GoogleSearchAgentTool(createGoogleSearchAgent(agent.model)) + tool is VertexAiSearchTool && tool.bypassMultiToolsLimit -> + VertexAiSearchAgentTool(createVertexAiSearchAgent(agent.model, tool)) else -> tool } } diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/tools/VertexAiSearchAgentTool.kt b/core/src/commonMain/kotlin/com/google/adk/kt/tools/VertexAiSearchAgentTool.kt new file mode 100644 index 00000000..28ea68fb --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/tools/VertexAiSearchAgentTool.kt @@ -0,0 +1,61 @@ +/* + * Copyright 2026 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.kt.tools + +import com.google.adk.kt.agents.Instruction +import com.google.adk.kt.agents.LlmAgent +import com.google.adk.kt.models.Model + +/** Name of the sub-agent (and resulting function tool) created by [createVertexAiSearchAgent]. */ +internal const val VERTEX_AI_SEARCH_AGENT_NAME = "vertex_ai_search_agent" + +/** + * Creates a sub-agent whose only tool is the given [VertexAiSearchTool]. + * + * Used by the multi-tools-limit workaround so that `vertex_ai_search` can be exposed to the parent + * agent as a function tool (via [VertexAiSearchAgentTool]) and therefore coexist with other tools. + * The original [tool] is reused as the sub-agent's only tool; since the sub-agent has a single + * tool, it is never itself rewritten by the workaround. + */ +internal fun createVertexAiSearchAgent(model: Model, tool: VertexAiSearchTool): LlmAgent = + LlmAgent( + name = VERTEX_AI_SEARCH_AGENT_NAME, + model = model, + description = "An agent for searching using the `vertex_ai_search` tool", + instruction = + Instruction.Text( + "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 = listOf(tool), + ) + +/** + * A tool that wraps a sub-agent which only uses a [VertexAiSearchTool]. + * + * Lets `vertex_ai_search` be used alongside other tools (built-in tools cannot be combined with + * other tools in a single request). The built-in is wrapped in a sub-agent and exposed as a + * function tool, and the sub-agent's grounding metadata is propagated back to the parent. + * + * NOTE: this intentionally diverges from the Python ADK, which instead converts the + * `VertexAiSearchTool` into a `DiscoveryEngineSearchTool` that queries the Discovery Engine API + * directly. That public Discovery Engine client is not available as a dependency here, so this + * reuses the same sub-agent mechanism as [GoogleSearchAgentTool]. + */ +internal class VertexAiSearchAgentTool(agent: LlmAgent) : + AgentTool(agent = agent, propagateGroundingMetadata = true) diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/tools/VertexAiSearchTool.kt b/core/src/commonMain/kotlin/com/google/adk/kt/tools/VertexAiSearchTool.kt index 4f383202..e2278b3a 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/tools/VertexAiSearchTool.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/tools/VertexAiSearchTool.kt @@ -37,6 +37,9 @@ import com.google.adk.kt.types.VertexAISearchDataStoreSpec * `projects/{project}/locations/{location}/collections/{collection}/engines/{engine}`. * @property filter The filter to apply to the search results. * @property maxResults The maximum number of results to return. + * @property bypassMultiToolsLimit When true, allows this built-in tool to be used alongside other + * tools: it is exposed as a function tool (see [VertexAiSearchAgentTool]) since built-in tools + * cannot otherwise be combined with other tools in a single request. * @property model Deprecated and unused. Tool support is verified by the backend. */ class VertexAiSearchTool( @@ -45,6 +48,7 @@ class VertexAiSearchTool( val searchEngineId: String? = null, val filter: String? = null, val maxResults: Int? = null, + val bypassMultiToolsLimit: Boolean = false, @Deprecated( "Model-based tool gating has been removed; tool support is verified by the backend. " + "This parameter is unused and will be removed in a future release." diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentMultiToolsLimitTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentMultiToolsLimitTest.kt index 7dbf470a..75685858 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentMultiToolsLimitTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentMultiToolsLimitTest.kt @@ -24,6 +24,7 @@ import com.google.adk.kt.testing.DummyTool import com.google.adk.kt.testing.modelMessage import com.google.adk.kt.testing.userMessage import com.google.adk.kt.tools.GoogleSearchTool +import com.google.adk.kt.tools.VertexAiSearchTool import kotlin.test.Test import kotlin.test.assertNotNull import kotlin.test.assertNull @@ -110,4 +111,35 @@ class LlmAgentMultiToolsLimitTest { // Did not opt in: the built-in google_search is left as-is (not swapped). assertNotNull(tools.firstOrNull { it.googleSearch != null }) } + + @Test + fun multipleTools_vertexAiSearchWithBypass_swappedToFunctionTool() = runTest { + var captured: LlmRequest? = null + val model = + DummyModel("m") { request -> + captured = request + flowOf(LlmResponse(content = modelMessage("done"))) + } + val agent = + LlmAgent( + name = "main", + model = model, + tools = + listOf( + VertexAiSearchTool(dataStoreId = "ds1", bypassMultiToolsLimit = true), + DummyTool("calc"), + ), + ) + val runner = InMemoryRunner(agent = agent) + + val unused = + runner.runAsync(userId = "u", sessionId = "s", newMessage = userMessage("hi")).toList() + + val tools = captured?.config?.tools + assertNotNull(tools) + // The built-in vertex_ai_search retrieval is replaced and re-exposed as a function tool. + assertNull(tools.firstOrNull { it.retrieval != null }) + val declaredNames = tools.flatMap { it.functionDeclarations ?: emptyList() }.map { it.name } + assertTrue(declaredNames.contains("vertex_ai_search_agent")) + } } diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/tools/VertexAiSearchAgentToolTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/tools/VertexAiSearchAgentToolTest.kt new file mode 100644 index 00000000..2b23e071 --- /dev/null +++ b/core/src/commonTest/kotlin/com/google/adk/kt/tools/VertexAiSearchAgentToolTest.kt @@ -0,0 +1,47 @@ +/* + * Copyright 2026 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.kt.tools + +import com.google.adk.kt.testing.DummyModel +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertIs +import kotlin.test.assertTrue + +class VertexAiSearchAgentToolTest { + + @Test + fun createVertexAiSearchAgent_buildsAgentWithVertexAiSearchTool() { + val tool = VertexAiSearchTool(dataStoreId = "ds1", bypassMultiToolsLimit = true) + + val agent = createVertexAiSearchAgent(DummyModel("test"), tool) + + assertEquals(VERTEX_AI_SEARCH_AGENT_NAME, agent.name) + assertEquals(1, agent.tools.size) + assertIs(agent.tools[0]) + } + + @Test + fun vertexAiSearchAgentTool_declaration_usesAgentName() { + val tool = VertexAiSearchTool(dataStoreId = "ds1", bypassMultiToolsLimit = true) + + val agentTool = VertexAiSearchAgentTool(createVertexAiSearchAgent(DummyModel("test"), tool)) + + assertEquals(VERTEX_AI_SEARCH_AGENT_NAME, agentTool.declaration().name) + assertTrue(agentTool.propagateGroundingMetadata) + } +}