fix: Parallelize tool resolution in LlmAgent.canonical_tools()

Previously we resolved tools sequentially by awaiting _convert_tool_union_to_tools() in a loop -- reduce the latency by resolving tools concurrently.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 872979105
This commit is contained in:
Kathy Wu
2026-02-20 10:45:18 -08:00
committed by Copybara-Service
parent bef3f117b4
commit 7478bdaa98
2 changed files with 33 additions and 8 deletions
+12 -8
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
import asyncio
import importlib
import inspect
import logging
@@ -589,24 +590,27 @@ class LlmAgent(BaseAgent):
return global_instruction, True
async def canonical_tools(
self, ctx: ReadonlyContext = None
self, ctx: Optional[ReadonlyContext] = None
) -> list[BaseTool]:
"""The resolved self.tools field as a list of BaseTool based on the context.
This method is only for use by Agent Development Kit.
"""
resolved_tools = []
# We may need to wrap some built-in tools if there are other tools
# because the built-in tools cannot be used together with other tools.
# TODO(b/448114567): Remove once the workaround is no longer needed.
multiple_tools = len(self.tools) > 1
model = self.canonical_model
for tool_union in self.tools:
resolved_tools.extend(
await _convert_tool_union_to_tools(
tool_union, ctx, model, multiple_tools
)
)
results = await asyncio.gather(*(
_convert_tool_union_to_tools(tool_union, ctx, model, multiple_tools)
for tool_union in self.tools
))
resolved_tools = []
for tools in results:
resolved_tools.extend(tools)
return resolved_tools
@property
@@ -451,6 +451,27 @@ class TestCanonicalTools:
assert tools[0].name == 'vertex_ai_search'
assert tools[0].__class__.__name__ == 'VertexAiSearchTool'
async def test_multiple_tools_resolution(self):
"""Test that multiple tools are resolved correctly."""
def _tool_1():
pass
def _tool_2():
pass
agent = LlmAgent(
name='test_agent',
model='gemini-pro',
tools=[_tool_1, _tool_2],
)
ctx = await _create_readonly_context(agent)
tools = await agent.canonical_tools(ctx)
assert len(tools) == 2
assert tools[0].name == '_tool_1'
assert tools[1].name == '_tool_2'
# Tests for multi-provider model support via string model names
@pytest.mark.parametrize(