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 <wuliang@google.com>
PiperOrigin-RevId: 831537187

Author: wuliang229

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
 
-  __all__.extend([
-      'MCPToolset',
-      '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}')
+
+  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__