Skip to content

Commit 749b4f6

Browse files
m-badenm-baden
committed
Merge origin/main and resolve conflicts
2 parents 029b87d + 8436e1f commit 749b4f6

7 files changed

Lines changed: 389 additions & 33 deletions

File tree

‎src/google/adk/agents/llm_agent.py‎

Lines changed: 66 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -142,38 +142,39 @@ async def _convert_tool_union_to_tools(
142142
model: Union[str, BaseLlm],
143143
multiple_tools: bool = False,
144144
) -> list[BaseTool]:
145+
from ..tools.enterprise_search_tool import EnterpriseWebSearchTool
145146
from ..tools.google_search_tool import GoogleSearchTool
146147
from ..tools.vertex_ai_search_tool import VertexAiSearchTool
147148

148-
# Wrap google_search tool with AgentTool if there are multiple tools because
149-
# the built-in tools cannot be used together with other tools.
149+
# Handle built-in tool workarounds when multiple tools are present.
150+
# Built-in tools cannot be used together with other tools, so we wrap or
151+
# replace them with compatible alternatives.
150152
# TODO(b/448114567): Remove once the workaround is no longer needed.
151-
if multiple_tools and isinstance(tool_union, GoogleSearchTool):
152-
from ..tools.google_search_agent_tool import create_google_search_agent
153-
from ..tools.google_search_agent_tool import GoogleSearchAgentTool
154-
155-
search_tool = cast(GoogleSearchTool, tool_union)
156-
if search_tool.bypass_multi_tools_limit:
157-
return [GoogleSearchAgentTool(create_google_search_agent(model))]
158-
159-
# Replace VertexAiSearchTool with DiscoveryEngineSearchTool if there are
160-
# multiple tools because the built-in tools cannot be used together with
161-
# other tools.
162-
# TODO(b/448114567): Remove once the workaround is no longer needed.
163-
if multiple_tools and isinstance(tool_union, VertexAiSearchTool):
164-
from ..tools.discovery_engine_search_tool import DiscoveryEngineSearchTool
165-
166-
vais_tool = cast(VertexAiSearchTool, tool_union)
167-
if vais_tool.bypass_multi_tools_limit:
168-
return [
169-
DiscoveryEngineSearchTool(
170-
data_store_id=vais_tool.data_store_id,
171-
data_store_specs=vais_tool.data_store_specs,
172-
search_engine_id=vais_tool.search_engine_id,
173-
filter=vais_tool.filter,
174-
max_results=vais_tool.max_results,
175-
)
176-
]
153+
if multiple_tools:
154+
tool_workarounds = [
155+
# GoogleSearchTool: wrap with AgentTool
156+
{
157+
'tool_class': GoogleSearchTool,
158+
'handler': lambda: _handle_google_search_tool(tool_union, model),
159+
},
160+
# VertexAiSearchTool: replace with DiscoveryEngineSearchTool
161+
{
162+
'tool_class': VertexAiSearchTool,
163+
'handler': lambda: _handle_vertex_ai_search_tool(tool_union),
164+
},
165+
# EnterpriseWebSearchTool: wrap with AgentTool
166+
{
167+
'tool_class': EnterpriseWebSearchTool,
168+
'handler': lambda: _handle_enterprise_search_tool(
169+
tool_union, model
170+
),
171+
},
172+
]
173+
174+
for workaround in tool_workarounds:
175+
if isinstance(tool_union, workaround['tool_class']):
176+
if tool_union.bypass_multi_tools_limit:
177+
return workaround['handler']()
177178

178179
if isinstance(tool_union, BaseTool):
179180
return [tool_union]
@@ -192,6 +193,43 @@ async def _convert_tool_union_to_tools(
192193
return []
193194

194195

196+
def _handle_google_search_tool(
197+
tool_union: ToolUnion, model: Union[str, BaseLlm]
198+
) -> list[BaseTool]:
199+
"""Handle GoogleSearchTool workaround by wrapping with AgentTool."""
200+
from ..tools.google_search_agent_tool import create_google_search_agent
201+
from ..tools.google_search_agent_tool import GoogleSearchAgentTool
202+
203+
return [GoogleSearchAgentTool(create_google_search_agent(model))]
204+
205+
206+
def _handle_vertex_ai_search_tool(tool_union: ToolUnion) -> list[BaseTool]:
207+
"""Handle VertexAiSearchTool workaround by replacing with DiscoveryEngineSearchTool."""
208+
from ..tools.discovery_engine_search_tool import DiscoveryEngineSearchTool
209+
from ..tools.vertex_ai_search_tool import VertexAiSearchTool
210+
211+
vais_tool = cast(VertexAiSearchTool, tool_union)
212+
return [
213+
DiscoveryEngineSearchTool(
214+
data_store_id=vais_tool.data_store_id,
215+
data_store_specs=vais_tool.data_store_specs,
216+
search_engine_id=vais_tool.search_engine_id,
217+
filter=vais_tool.filter,
218+
max_results=vais_tool.max_results,
219+
)
220+
]
221+
222+
223+
def _handle_enterprise_search_tool(
224+
tool_union: ToolUnion, model: Union[str, BaseLlm]
225+
) -> list[BaseTool]:
226+
"""Handle EnterpriseWebSearchTool workaround by wrapping with AgentTool."""
227+
from ..tools.enterprise_search_agent_tool import create_enterprise_search_agent
228+
from ..tools.enterprise_search_agent_tool import EnterpriseSearchAgentTool
229+
230+
return [EnterpriseSearchAgentTool(create_enterprise_search_agent(model))]
231+
232+
195233
class LlmAgent(BaseAgent):
196234
"""LLM-based Agent."""
197235

‎src/google/adk/flows/llm_flows/base_llm_flow.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,10 @@ async def _maybe_add_grounding_metadata(
266266
tools = await agent.canonical_tools(readonly_context)
267267
invocation_context.canonical_tools_cache = tools
268268

269-
if not any(tool.name == 'google_search_agent' for tool in tools):
269+
if not any(
270+
tool.name in {'google_search_agent', 'enterprise_search_agent'}
271+
for tool in tools
272+
):
270273
return response
271274
ground_metadata = invocation_context.session.state.get(
272275
'temp:_adk_grounding_metadata', None

‎src/google/adk/tools/_search_agent_tool.py‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from typing import Any
18+
19+
from google.genai import types
20+
from typing_extensions import override
21+
22+
from ..agents.llm_agent import LlmAgent
23+
from ..memory.in_memory_memory_service import InMemoryMemoryService
24+
from ..runners import Runner
25+
from ..sessions.in_memory_session_service import InMemorySessionService
26+
from ..utils.context_utils import Aclosing
27+
from ._forwarding_artifact_service import ForwardingArtifactService
28+
from .agent_tool import AgentTool
29+
from .tool_context import ToolContext
30+
31+
32+
class _SearchAgentTool(AgentTool):
33+
"""A base class for search agent tools."""
34+
35+
@override
36+
async def run_async(
37+
self,
38+
*,
39+
args: dict[str, Any],
40+
tool_context: ToolContext,
41+
) -> Any:
42+
43+
if isinstance(self.agent, LlmAgent) and self.agent.input_schema:
44+
input_value = self.agent.input_schema.model_validate(args)
45+
content = types.Content(
46+
role='user',
47+
parts=[
48+
types.Part.from_text(
49+
text=input_value.model_dump_json(exclude_none=True)
50+
)
51+
],
52+
)
53+
else:
54+
content = types.Content(
55+
role='user',
56+
parts=[types.Part.from_text(text=args['request'])],
57+
)
58+
runner = Runner(
59+
app_name=self.agent.name,
60+
agent=self.agent,
61+
artifact_service=ForwardingArtifactService(tool_context),
62+
session_service=InMemorySessionService(),
63+
memory_service=InMemoryMemoryService(),
64+
credential_service=tool_context._invocation_context.credential_service,
65+
plugins=list(tool_context._invocation_context.plugin_manager.plugins),
66+
)
67+
try:
68+
state_dict = {
69+
k: v
70+
for k, v in tool_context.state.to_dict().items()
71+
if not k.startswith('_adk') and not k.startswith('temp:')
72+
}
73+
session = await runner.session_service.create_session(
74+
app_name=self.agent.name,
75+
user_id=tool_context._invocation_context.user_id,
76+
state=state_dict,
77+
)
78+
79+
last_content = None
80+
last_grounding_metadata = None
81+
async with Aclosing(
82+
runner.run_async(
83+
user_id=session.user_id,
84+
session_id=session.id,
85+
new_message=content,
86+
)
87+
) as agen:
88+
async for event in agen:
89+
# Forward state delta to parent session.
90+
if event.actions.state_delta:
91+
tool_context.state.update(event.actions.state_delta)
92+
if event.content:
93+
last_content = event.content
94+
last_grounding_metadata = event.grounding_metadata
95+
96+
if not last_content:
97+
return ''
98+
merged_text = '\n'.join(p.text for p in last_content.parts if p.text)
99+
if isinstance(self.agent, LlmAgent) and self.agent.output_schema:
100+
tool_result = self.agent.output_schema.model_validate_json(
101+
merged_text
102+
).model_dump(exclude_none=True)
103+
else:
104+
tool_result = merged_text
105+
106+
if last_grounding_metadata:
107+
tool_context.state['temp:_adk_grounding_metadata'] = (
108+
last_grounding_metadata
109+
)
110+
return tool_result
111+
finally:
112+
await runner.close()

‎src/google/adk/tools/enterprise_search_agent_tool.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from typing import Union
18+
19+
from ..agents.llm_agent import LlmAgent
20+
from ..models.base_llm import BaseLlm
21+
from ._search_agent_tool import _SearchAgentTool
22+
from .enterprise_search_tool import enterprise_web_search_tool
23+
24+
25+
def create_enterprise_search_agent(model: Union[str, BaseLlm]) -> LlmAgent:
26+
"""Create a sub-agent that only uses enterprise_web_search tool."""
27+
return LlmAgent(
28+
name='enterprise_search_agent',
29+
model=model,
30+
description=(
31+
'An agent for performing Enterprise search using the'
32+
' `enterprise_web_search` tool'
33+
),
34+
instruction="""
35+
You are a specialized Enterprise search agent.
36+
37+
When given a search query, use the `enterprise_web_search` tool to find the related information.
38+
""",
39+
tools=[enterprise_web_search_tool],
40+
)
41+
42+
43+
class EnterpriseSearchAgentTool(_SearchAgentTool):
44+
"""A tool that wraps a sub-agent that only uses enterprise_web_search tool.
45+
46+
This is a workaround to support using enterprise_web_search tool with other tools.
47+
TODO(b/448114567): Remove once the workaround is no longer needed.
48+
49+
Attributes:
50+
agent: The sub-agent that this tool wraps.
51+
"""
52+
53+
def __init__(self, agent: LlmAgent):
54+
self.agent = agent
55+
super().__init__(agent=self.agent)

‎src/google/adk/tools/enterprise_search_tool.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,18 @@ class EnterpriseWebSearchTool(BaseTool):
4141
4242
"""
4343

44-
def __init__(self):
45-
"""Initializes the Enterprise Web Search tool."""
44+
def __init__(self, *, bypass_multi_tools_limit: bool = False):
45+
"""Initializes the Enterprise web search tool.
46+
47+
Args:
48+
bypass_multi_tools_limit: Whether to bypass the multi tools limitation,
49+
so that the tool can be used with other tools in the same agent.
50+
"""
4651
# Name and description are not used because this is a model built-in tool.
4752
super().__init__(
4853
name='enterprise_web_search', description='enterprise_web_search'
4954
)
55+
self.bypass_multi_tools_limit = bypass_multi_tools_limit
5056

5157
@override
5258
async def process_llm_request(

‎src/google/adk/tools/google_search_agent_tool.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@
1818

1919
from ..agents.llm_agent import LlmAgent
2020
from ..models.base_llm import BaseLlm
21-
from .agent_tool import AgentTool
21+
from ..utils.context_utils import Aclosing
22+
from ._forwarding_artifact_service import ForwardingArtifactService
23+
from ._search_agent_tool import _SearchAgentTool
2224
from .google_search_tool import google_search
2325

2426

@@ -39,7 +41,7 @@ def create_google_search_agent(model: Union[str, BaseLlm]) -> LlmAgent:
3941
)
4042

4143

42-
class GoogleSearchAgentTool(AgentTool):
44+
class GoogleSearchAgentTool(_SearchAgentTool):
4345
"""A tool that wraps a sub-agent that only uses google_search tool.
4446
4547
This is a workaround to support using google_search tool with other tools.

0 commit comments

Comments
 (0)