mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
bef3f117b4
commit
7478bdaa98
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user