You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
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
This commit is contained in:
Liang Wu
committed by
Copybara-Service
Copybara-Service
parent
0ab79b9502
commit
69627b699f
@@ -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__
|
||||
|
||||
Reference in New Issue
Block a user