From 69627b699fc99ffd293d55656c21d4e14c09255f Mon Sep 17 00:00:00 2001 From: Liang Wu Date: Wed, 12 Nov 2025 14:16:51 -0800 Subject: [PATCH] chore: Implement lazy loading for modules within `google.adk.tools` This is ~27% of the current cold start latency after previous changes are submitted. This change refactors `google.adk.tools/__init__.py` to use `__getattr__` for lazy loading of all tools and related classes. Previously, all modules were imported directly upon `google.adk.tools` import, leading to potentially long initial import times. With lazy loading, modules are only imported when they are first accessed, improving the initial startup performance. Co-authored-by: Liang Wu PiperOrigin-RevId: 831537187 --- src/google/adk/tools/__init__.py | 139 +++++++++++++++++++------------ 1 file changed, 84 insertions(+), 55 deletions(-) diff --git a/src/google/adk/tools/__init__.py b/src/google/adk/tools/__init__.py index f5250d0a..1777bd93 100644 --- a/src/google/adk/tools/__init__.py +++ b/src/google/adk/tools/__init__.py @@ -11,65 +11,94 @@ # 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. +import importlib import logging import sys +from typing import Any +from typing import TYPE_CHECKING -from ..auth.auth_tool import AuthToolArguments -from .agent_tool import AgentTool -from .apihub_tool.apihub_toolset import APIHubToolset -from .base_tool import BaseTool -from .discovery_engine_search_tool import DiscoveryEngineSearchTool -from .enterprise_search_tool import enterprise_web_search_tool as enterprise_web_search -from .example_tool import ExampleTool -from .exit_loop_tool import exit_loop -from .function_tool import FunctionTool -from .get_user_choice_tool import get_user_choice_tool as get_user_choice -from .google_maps_grounding_tool import google_maps_grounding -from .google_search_tool import google_search -from .load_artifacts_tool import load_artifacts_tool as load_artifacts -from .load_memory_tool import load_memory_tool as load_memory -from .long_running_tool import LongRunningFunctionTool -from .preload_memory_tool import preload_memory_tool as preload_memory -from .tool_context import ToolContext -from .transfer_to_agent_tool import transfer_to_agent -from .url_context_tool import url_context -from .vertex_ai_search_tool import VertexAiSearchTool +# The TYPE_CHECKING block is needed for autocomplete to work. +if TYPE_CHECKING: + from ..auth.auth_tool import AuthToolArguments + from .agent_tool import AgentTool + from .apihub_tool.apihub_toolset import APIHubToolset + from .base_tool import BaseTool + from .discovery_engine_search_tool import DiscoveryEngineSearchTool + from .enterprise_search_tool import enterprise_web_search_tool as enterprise_web_search + from .example_tool import ExampleTool + from .exit_loop_tool import exit_loop + from .function_tool import FunctionTool + from .get_user_choice_tool import get_user_choice_tool as get_user_choice + from .google_maps_grounding_tool import google_maps_grounding + from .google_search_tool import google_search + from .load_artifacts_tool import load_artifacts_tool as load_artifacts + from .load_memory_tool import load_memory_tool as load_memory + from .long_running_tool import LongRunningFunctionTool + from .preload_memory_tool import preload_memory_tool as preload_memory + from .tool_context import ToolContext + from .transfer_to_agent_tool import transfer_to_agent + from .url_context_tool import url_context + from .vertex_ai_search_tool import VertexAiSearchTool -__all__ = [ - 'AgentTool', - 'APIHubToolset', - 'AuthToolArguments', - 'BaseTool', - 'DiscoveryEngineSearchTool', - 'enterprise_web_search', - 'google_maps_grounding', - 'google_search', - 'url_context', - 'VertexAiSearchTool', - 'ExampleTool', - 'exit_loop', - 'FunctionTool', - 'get_user_choice', - 'load_artifacts', - 'load_memory', - 'LongRunningFunctionTool', - 'preload_memory', - 'ToolContext', - 'transfer_to_agent', -] +# If you are adding a new tool to this file, please make sure you add it to the +# lazy mapping to avoid expensive imports. If the tool is not using any third +# party dependencies, please feel free to import it eagerly at the top of this +# file. +_LAZY_MAPPING = { + 'AuthToolArguments': ('..auth.auth_tool', 'AuthToolArguments'), + 'AgentTool': ('.agent_tool', 'AgentTool'), + 'APIHubToolset': ('.apihub_tool.apihub_toolset', 'APIHubToolset'), + 'BaseTool': ('.base_tool', 'BaseTool'), + 'DiscoveryEngineSearchTool': ( + '.discovery_engine_search_tool', + 'DiscoveryEngineSearchTool', + ), + 'enterprise_web_search': ( + '.enterprise_search_tool', + 'enterprise_web_search_tool', + ), + 'ExampleTool': ('.example_tool', 'ExampleTool'), + 'exit_loop': ('.exit_loop_tool', 'exit_loop'), + 'FunctionTool': ('.function_tool', 'FunctionTool'), + 'get_user_choice': ('.get_user_choice_tool', 'get_user_choice_tool'), + 'google_maps_grounding': ( + '.google_maps_grounding_tool', + 'google_maps_grounding', + ), + 'google_search': ('.google_search_tool', 'google_search'), + 'load_artifacts': ('.load_artifacts_tool', 'load_artifacts_tool'), + 'load_memory': ('.load_memory_tool', 'load_memory_tool'), + 'LongRunningFunctionTool': ( + '.long_running_tool', + 'LongRunningFunctionTool', + ), + 'preload_memory': ('.preload_memory_tool', 'preload_memory_tool'), + 'ToolContext': ('.tool_context', 'ToolContext'), + 'transfer_to_agent': ('.transfer_to_agent_tool', 'transfer_to_agent'), + 'url_context': ('.url_context_tool', 'url_context'), + 'VertexAiSearchTool': ('.vertex_ai_search_tool', 'VertexAiSearchTool'), + 'MCPToolset': ('.mcp_tool.mcp_toolset', 'MCPToolset'), + 'McpToolset': ('.mcp_tool.mcp_toolset', 'McpToolset'), +} + +__all__ = list(_LAZY_MAPPING.keys()) -if sys.version_info < (3, 10): - logger = logging.getLogger('google_adk.' + __name__) - logger.warning( - 'MCP requires Python 3.10 or above. Please upgrade your Python' - ' version in order to use it.' - ) -else: - from .mcp_tool.mcp_toolset import MCPToolset - from .mcp_tool.mcp_toolset import McpToolset +def __getattr__(name: str) -> Any: + """Lazy loads tools to avoid expensive imports.""" + if name not in _LAZY_MAPPING: + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') - __all__.extend([ - 'MCPToolset', - 'McpToolset', - ]) + module_path, attr_name = _LAZY_MAPPING[name] + # __name__ is `google.adk.tools` and we are doing a relative import + # from there. + module = importlib.import_module(module_path, __name__) + attr = getattr(module, attr_name) + globals()[name] = attr + return attr + + +# __dir__ is used to expose all public interfaces to keep mocking with autoscope +# working. +def __dir__() -> list[str]: + return list(globals().keys()) + __all__