Fix #2433: Lazy load modules to improve import time

Author: 9chait9

src/google/adk/__init__.py

@@ -13,11 +13,25 @@
 # limitations under the License.
 
 from __future__ import annotations
+from typing import TYPE_CHECKING
 
-from . import version
-from .agents.context import Context
-from .agents.llm_agent import Agent
-from .runners import Runner
+if TYPE_CHECKING:
+    from .agents.context import Context
+    from .agents.llm_agent import Agent
+    from .runners import Runner
+else:
+    def __getattr__(name: str):
+        if name == "Agent":
+            from .agents.llm_agent import Agent
+            return Agent
+        if name == "Context":
+            from .agents.context import Context
+            return Context
+        if name == "Runner":
+            from .runners import Runner
+            return Runner
+        raise AttributeError(f"module {__name__} has no attribute {name}")
 
+from . import version
 __version__ = version.__version__
 __all__ = ["Agent", "Context", "Runner"]

src/google/adk/agents/__init__.py

@@ -12,18 +12,61 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base_agent import BaseAgent
-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 .loop_agent import LoopAgent
-from .mcp_instruction_provider import McpInstructionProvider
-from .parallel_agent import ParallelAgent
-from .run_config import RunConfig
-from .sequential_agent import SequentialAgent
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .base_agent import BaseAgent
+    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 .loop_agent import LoopAgent
+    from .mcp_instruction_provider import McpInstructionProvider
+    from .parallel_agent import ParallelAgent
+    from .run_config import RunConfig
+    from .sequential_agent import SequentialAgent
+
+def __getattr__(name: str):
+    if name == 'BaseAgent':
+        from .base_agent import BaseAgent
+        return BaseAgent
+    if name == 'Context':
+        from .context import Context
+        return Context
+    if name == 'InvocationContext':
+        from .invocation_context import InvocationContext
+        return InvocationContext
+    if name == 'LiveRequest':
+        from .live_request_queue import LiveRequest
+        return LiveRequest
+    if name == 'LiveRequestQueue':
+        from .live_request_queue import LiveRequestQueue
+        return LiveRequestQueue
+    if name == 'Agent':
+        from .llm_agent import Agent
+        return Agent
+    if name == 'LlmAgent':
+        from .llm_agent import LlmAgent
+        return LlmAgent
+    if name == 'LoopAgent':
+        from .loop_agent import LoopAgent
+        return LoopAgent
+    if name == 'McpInstructionProvider':
+        from .mcp_instruction_provider import McpInstructionProvider
+        return McpInstructionProvider
+    if name == 'ParallelAgent':
+        from .parallel_agent import ParallelAgent
+        return ParallelAgent
+    if name == 'RunConfig':
+        from .run_config import RunConfig
+        return RunConfig
+    if name == 'SequentialAgent':
+        from .sequential_agent import SequentialAgent
+        return SequentialAgent
+    raise AttributeError(f"module {__name__} has no attribute {name}")
 
 __all__ = [
     'Agent',