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}")…tr__ scope for better maintainability.
|
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. |
2 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.