diff --git a/src/google/adk/__init__.py b/src/google/adk/__init__.py index be9d2af08b..21d6d0f9bc 100644 --- a/src/google/adk/__init__.py +++ b/src/google/adk/__init__.py @@ -13,13 +13,32 @@ # limitations under the License. from __future__ import annotations +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .agents.context import Context + from .agents.llm_agent import Agent + from .runners import Runner + from .events.event import Event + from .workflow import Workflow +else: + import importlib + + _LAZY_IMPORTS = { + "Agent": ".agents.llm_agent", + "Context": ".agents.context", + "Runner": ".runners", + "Event": ".events.event", + "Workflow": ".workflow", + } + + def __getattr__(name: str): + if name in _LAZY_IMPORTS: + module = importlib.import_module(_LAZY_IMPORTS[name], __name__) + return getattr(module, name) + raise AttributeError(f"module {__name__} has no attribute {name}") from . import version -from .agents.context import Context -from .agents.llm_agent import Agent -from .events.event import Event -from .runners import Runner -from .workflow import Workflow __version__ = version.__version__ __all__ = ["Agent", "Context", "Event", "Runner", "Workflow"] diff --git a/src/google/adk/agents/__init__.py b/src/google/adk/agents/__init__.py index 9d5749f50f..07c22062be 100644 --- a/src/google/adk/agents/__init__.py +++ b/src/google/adk/agents/__init__.py @@ -8,33 +8,76 @@ # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, -# 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. - +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and +# limitations under the License.\n +from __future__ import annotations import importlib from typing import Any from typing import TYPE_CHECKING -from .base_agent import BaseAgent -from .base_agent_config import BaseAgentConfig -from .context import Context -from .invocation_context import InvocationContext -from .live_request_queue import LiveRequest -from .live_request_queue import LiveRequestQueue -from .llm_agent import Agent -from .llm_agent import LlmAgent -from .llm_agent_config import LlmAgentConfig -from .loop_agent import LoopAgent -from .loop_agent_config import LoopAgentConfig -from .parallel_agent import ParallelAgent -from .parallel_agent_config import ParallelAgentConfig -from .run_config import RunConfig -from .sequential_agent import SequentialAgent -from .sequential_agent_config import SequentialAgentConfig - if TYPE_CHECKING: - from .mcp_instruction_provider import McpInstructionProvider + from .base_agent import BaseAgent + from .base_agent_config import BaseAgentConfig + from .context import Context + from .invocation_context import InvocationContext + from .live_request_queue import LiveRequest + from .live_request_queue import LiveRequestQueue + from .llm_agent import Agent + from .llm_agent import LlmAgent + from .llm_agent_config import LlmAgentConfig + from .loop_agent import LoopAgent + from .loop_agent_config import LoopAgentConfig + from .parallel_agent import ParallelAgent + from .parallel_agent_config import ParallelAgentConfig + from .run_config import RunConfig + from .sequential_agent import SequentialAgent + from .sequential_agent_config import SequentialAgentConfig + from .mcp_instruction_provider import McpInstructionProvider +else: + import importlib + + _LAZY_IMPORTS = { + "Agent": ".llm_agent", + "BaseAgent": ".base_agent", + "BaseAgentConfig": ".base_agent_config", + "Context": ".context", + "InvocationContext": ".invocation_context", + "LiveRequest": ".live_request_queue", + "LiveRequestQueue": ".live_request_queue", + "LlmAgent": ".llm_agent", + "LlmAgentConfig": ".llm_agent_config", + "LoopAgent": ".loop_agent", + "LoopAgentConfig": ".loop_agent_config", + "McpInstructionProvider": ".mcp_instruction_provider", + "ParallelAgent": ".parallel_agent", + "ParallelAgentConfig": ".parallel_agent_config", + "RunConfig": ".run_config", + "SequentialAgent": ".sequential_agent", + "SequentialAgentConfig": ".sequential_agent_config", + } + + def __getattr__(name: str) -> Any: + if name in _LAZY_IMPORTS: + if name == "McpInstructionProvider": + try: + module = importlib.import_module( + _LAZY_IMPORTS[name], __name__ + ) + except ImportError as e: + raise ImportError( + "`McpInstructionProvider` requires the `mcp` package." + " Install with: pip install google-adk[extensions]" + ) from e + return getattr(module, name) + else: + module = importlib.import_module(_LAZY_IMPORTS[name], __name__) + return getattr(module, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + + def __dir__() -> list[str]: + return list(globals().keys()) + list(_LAZY_IMPORTS.keys()) + __all__ = [ 'Agent', @@ -55,16 +98,3 @@ 'ParallelAgentConfig', 'SequentialAgentConfig', ] - - -def __getattr__(name: str) -> Any: - if name == 'McpInstructionProvider': - module = importlib.import_module('.mcp_instruction_provider', __name__) - attr = getattr(module, 'McpInstructionProvider') - globals()[name] = attr - return attr - raise AttributeError(f'module {__name__!r} has no attribute {name!r}') - - -def __dir__() -> list[str]: - return list(globals().keys()) + __all__ diff --git a/tests/unittests/test_lazy_loading.py b/tests/unittests/test_lazy_loading.py new file mode 100644 index 0000000000..2134de890b --- /dev/null +++ b/tests/unittests/test_lazy_loading.py @@ -0,0 +1,28 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 sys +import unittest +import google.adk + +class TestLazyLoading(unittest.TestCase): + def test_agent_not_loaded(self): + # Before accessing Agent, it shouldn't be in sys.modules + self.assertNotIn("google.adk.agents.llm_agent", sys.modules) + + # Accessing Agent + _ = google.adk.Agent + + # Now it should be loaded + self.assertIn("google.adk.agents.llm_agent", sys.modules)