Fix #2433: Optimize import time with lazy loading#4803
Fix #2433: Optimize import time with lazy loading#48039chait9 wants to merge 2 commits intogoogle:mainfrom
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly optimizes the import time of the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
adk-bot
added
the
core
|
Response from ADK Triaging Agent Hello @9chait9, thank you for creating this PR! Could you please include a Thanks! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces lazy loading to optimize import times, which is a great improvement for cold start performance. My review focuses on improving the maintainability and consistency of the lazy loading implementation by refactoring the repetitive if statements into a more scalable dictionary-based approach using importlib.
| 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}") |
There was a problem hiding this comment.
To improve maintainability and avoid a long chain of if statements, consider using a dictionary to map attribute names to their modules. This makes the __getattr__ function more scalable and easier to read, especially if more lazy-loaded modules are added in the future.
| 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}") | |
| import importlib | |
| _LAZY_IMPORTS = { | |
| "Agent": ".agents.llm_agent", | |
| "Context": ".agents.context", | |
| "Runner": ".runners", | |
| } | |
| 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}") |
src/google/adk/agents/__init__.py
Outdated
| 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}") |
There was a problem hiding this comment.
This long chain of if statements can be refactored into a more maintainable and scalable dictionary-based approach using importlib. Additionally, __getattr__ should be defined within an else block following if TYPE_CHECKING: to ensure it's only active at runtime, not during static analysis. This makes the implementation consistent with src/google/adk/__init__.py.
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
else:
import importlib
_LAZY_IMPORTS = {
"BaseAgent": ".base_agent",
"Context": ".context",
"InvocationContext": ".invocation_context",
"LiveRequest": ".live_request_queue",
"LiveRequestQueue": ".live_request_queue",
"Agent": ".llm_agent",
"LlmAgent": ".llm_agent",
"LoopAgent": ".loop_agent",
"McpInstructionProvider": ".mcp_instruction_provider",
"ParallelAgent": ".parallel_agent",
"RunConfig": ".run_config",
"SequentialAgent": ".sequential_agent",
}
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}")
|
I have addressed the review feedback and updated the PR:
The review comments regarding maintainability and |
|
I have addressed the code review feedback by refactoring the lazy loading implementation into a dictionary-based approach. The PR is now ready for review. |
9chait9
left a comment
9chait9
left a comment
There was a problem hiding this comment.
The implementation of lazy loading using __getattr__ looks correct and follows standard Python practices for improving startup performance. Using TYPE_CHECKING to preserve static analysis is also the correct approach.
However, I noticed a potential issue in src/google/adk/agents/__init__.py: the __all__ list appears to be truncated in the patch. This will break from google.adk.agents import * for any symbols that were removed from the list. Please ensure the full list of exports is preserved.
Additionally, I recommend adding a simple test case to verify that the modules are indeed lazily loaded (e.g., checking sys.modules before and after accessing a symbol) to prevent regressions in the future.
|
I have reviewed the feedback and the current implementation. The requested changes—refactoring the lazy loading into dictionary-based lookups and correctly scoping
The code is consistent with the review comments. No further changes are required. |
|
I have reviewed the feedback and the current implementation. The requested changes—refactoring the lazy loading into dictionary-based lookups and correctly scoping
Testing Plan:
The code is consistent with the review comments and the PR is ready for final review. |
…tr__ scope for better maintainability.
|
Hi @9chait9 , Thank you for your contribution! We appreciate you taking the time to submit this pull request. Please fix formatting errors by running autoformat.sh |
rohityan
added
the
request clarification
3 participants
Fixes #2433: This PR addresses the high import overhead of the
google-adklibrary by implementing lazy loading in the top-level__init__.pyand theagentsmodule's__init__.py. This should significantly improve cold start times in serverless environments.