perf: lazy-load optional providers and auth chain to cut cold start ~25%

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 905227934

Author: GWeale

src/google/adk/agents/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
+from typing import TYPE_CHECKING
+
 from .base_agent import BaseAgent
 from .context import Context
 from .invocation_context import InvocationContext
@@ -20,11 +23,13 @@
 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
 
+if TYPE_CHECKING:
+  from .mcp_instruction_provider import McpInstructionProvider
+
 __all__ = [
     'Agent',
     'BaseAgent',
@@ -39,3 +44,16 @@
     'LiveRequestQueue',
     'RunConfig',
 ]
+
+
+def __getattr__(name: str):
+  if name == 'McpInstructionProvider':
+    try:
+      module = importlib.import_module(f'{__name__}.mcp_instruction_provider')
+    except ImportError as e:
+      raise ImportError(
+          '`McpInstructionProvider` requires the `mcp` package.'
+          ' Install with: pip install google-adk[extensions]'
+      ) from e
+    return module.McpInstructionProvider
+  raise AttributeError(f'module {__name__!r} has no attribute {name!r}')

src/google/adk/auth/__init__.py

@@ -12,12 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+import importlib
+from typing import TYPE_CHECKING
+
 from .auth_credential import AuthCredential
 from .auth_credential import AuthCredentialTypes
 from .auth_credential import OAuth2Auth
-from .auth_handler import AuthHandler
 from .auth_schemes import AuthScheme
 from .auth_schemes import AuthSchemeType
 from .auth_schemes import OpenIdConnectWithConfig
 from .auth_tool import AuthConfig
 from .base_auth_provider import BaseAuthProvider
+
+if TYPE_CHECKING:
+  from .auth_handler import AuthHandler
+
+
+def __getattr__(name: str):
+  if name == 'AuthHandler':
+    return importlib.import_module(f'{__name__}.auth_handler').AuthHandler
+  raise AttributeError(f'module {__name__!r} has no attribute {name!r}')

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -36,9 +36,7 @@
 from ...agents.live_request_queue import LiveRequestQueue
 from ...agents.readonly_context import ReadonlyContext
 from ...agents.run_config import StreamingMode
-from ...auth.auth_handler import AuthHandler
 from ...auth.auth_tool import AuthConfig
-from ...auth.credential_manager import CredentialManager
 from ...events.event import Event
 from ...models.base_llm_connection import BaseLlmConnection
 from ...models.llm_request import LlmRequest
@@ -144,6 +142,8 @@ async def _resolve_toolset_auth(
       continue
 
     auth_config_copy = auth_config.model_copy(deep=True)
+    from ...auth.credential_manager import CredentialManager
+
     try:
       credential = await CredentialManager(
           auth_config_copy
@@ -173,7 +173,8 @@ async def _resolve_toolset_auth(
   if not pending_auth_requests:
     return
 
-  # Build auth requests dict with generated auth requests
+  from ...auth.auth_handler import AuthHandler
+
   auth_requests = {
       credential_id: AuthHandler(auth_config).generate_auth_request()
       for credential_id, auth_config in pending_auth_requests.items()

src/google/adk/flows/llm_flows/single_flow.py

@@ -29,14 +29,15 @@
 from . import instructions
 from . import interactions_processor
 from . import request_confirmation
-from ...auth import auth_preprocessor
 from .base_llm_flow import BaseLlmFlow
 
 logger = logging.getLogger('google_adk.' + __name__)
 
 
 def _create_request_processors():
   """Create the standard request processor list for a single-agent flow."""
+  from ...auth import auth_preprocessor
+
   return [
       basic.request_processor,
       auth_preprocessor.request_processor,

src/google/adk/models/__init__.py

@@ -14,52 +14,85 @@
 
 """Defines the interface to support a model."""
 
-from .apigee_llm import ApigeeLlm
+from __future__ import annotations
+
+import importlib
+from typing import TYPE_CHECKING
+
 from .base_llm import BaseLlm
-from .gemma_llm import Gemma
-from .google_llm import Gemini
 from .llm_request import LlmRequest
 from .llm_response import LlmResponse
 from .registry import LLMRegistry
 
+if TYPE_CHECKING:
+  from .anthropic_llm import Claude
+  from .apigee_llm import ApigeeLlm
+  from .gemma_llm import Gemma
+  from .gemma_llm import Gemma3Ollama
+  from .google_llm import Gemini
+  from .lite_llm import LiteLlm
+
 __all__ = [
+    'ApigeeLlm',
     'BaseLlm',
+    'Claude',
     'Gemini',
     'Gemma',
+    'Gemma3Ollama',
     'LLMRegistry',
+    'LiteLlm',
 ]
 
+_LAZY_PROVIDERS: dict[str, tuple[list[str], str]] = {
+    'Gemini': (
+        [
+            r'gemini-.*',
+            r'model-optimizer-.*',
+            r'projects\/.+\/locations\/.+\/endpoints\/.+',
+            r'projects\/.+\/locations\/.+\/publishers\/google\/models\/gemini.+',
+        ],
+        'google_llm',
+    ),
+    'Gemma': ([r'gemma-.*'], 'gemma_llm'),
+    'ApigeeLlm': ([r'.*-apigee$'], 'apigee_llm'),
+    'Claude': ([r'claude-3-.*', r'claude-.*-4.*'], 'anthropic_llm'),
+    'Gemma3Ollama': ([r'ollama/gemma3.*'], 'gemma_llm'),
+    'LiteLlm': (
+        [
+            r'openai/.*',
+            r'azure/.*',
+            r'azure_ai/.*',
+            r'groq/.*',
+            r'anthropic/.*',
+            r'bedrock/.*',
+            r'ollama/(?!gemma3).*',
+            r'ollama_chat/.*',
+            r'together_ai/.*',
+            r'vertex_ai/.*',
+            r'mistral/.*',
+            r'deepseek/.*',
+            r'fireworks_ai/.*',
+            r'cohere/.*',
+            r'databricks/.*',
+            r'ai21/.*',
+        ],
+        'lite_llm',
+    ),
+}
 
-LLMRegistry.register(Gemini)
-LLMRegistry.register(Gemma)
-LLMRegistry.register(ApigeeLlm)
+for _name, (_patterns, _module) in _LAZY_PROVIDERS.items():
+  LLMRegistry._register_lazy(_patterns, f'{__name__}.{_module}', _name)
 
-# Optionally register Claude if anthropic package is installed
-try:
-  from .anthropic_llm import Claude
-
-  LLMRegistry.register(Claude)
-  __all__.append('Claude')
-except Exception:
-  # Claude support requires: pip install google-adk[extensions]
-  pass
-
-# Optionally register LiteLlm if litellm package is installed
-try:
-  from .lite_llm import LiteLlm
-
-  LLMRegistry.register(LiteLlm)
-  __all__.append('LiteLlm')
-except Exception:
-  # LiteLLM support requires: pip install google-adk[extensions]
-  pass
-
-# Optionally register Gemma3Ollama if litellm package is installed
-try:
-  from .gemma_llm import Gemma3Ollama
 
-  LLMRegistry.register(Gemma3Ollama)
-  __all__.append('Gemma3Ollama')
-except Exception:
-  # Gemma3Ollama requires LiteLLM: pip install google-adk[extensions]
-  pass
+def __getattr__(name: str):
+  if name in _LAZY_PROVIDERS:
+    module_name = _LAZY_PROVIDERS[name][1]
+    try:
+      module = importlib.import_module(f'{__name__}.{module_name}')
+    except ImportError as e:
+      raise ImportError(
+          f'`{name}` requires an optional dependency that is not installed.'
+          ' Install with: pip install google-adk[extensions]'
+      ) from e
+    return getattr(module, name)
+  raise AttributeError(f'module {__name__!r} has no attribute {name!r}')

src/google/adk/models/registry.py

@@ -17,22 +17,20 @@
 from __future__ import annotations
 
 from functools import lru_cache
+import importlib
 import logging
 import re
 from typing import TYPE_CHECKING
+from typing import Union
 
 if TYPE_CHECKING:
   from .base_llm import BaseLlm
 
 logger = logging.getLogger('google_adk.' + __name__)
 
 
-_llm_registry_dict: dict[str, type[BaseLlm]] = {}
-"""Registry for LLMs.
-
-Key is the regex that matches the model name.
-Value is the class that implements the model.
-"""
+_LazyEntry = tuple[str, str]
+_llm_registry_dict: dict[str, Union[type['BaseLlm'], _LazyEntry]] = {}
 
 
 class LLMRegistry:
@@ -81,6 +79,14 @@ def register(llm_cls: type[BaseLlm]):
     for regex in llm_cls.supported_models():
       LLMRegistry._register(regex, llm_cls)
 
+  @staticmethod
+  def _register_lazy(
+      model_name_regexes: list[str], module_path: str, class_name: str
+  ):
+    """Pre-registers a lazily-imported LLM class."""
+    for regex in model_name_regexes:
+      _llm_registry_dict[regex] = (module_path, class_name)
+
   @staticmethod
   @lru_cache(maxsize=32)
   def resolve(model: str) -> type[BaseLlm]:
@@ -95,9 +101,20 @@ def resolve(model: str) -> type[BaseLlm]:
         ValueError: If the model is not found.
     """
 
-    for regex, llm_class in _llm_registry_dict.items():
-      if re.compile(regex).fullmatch(model):
+    for regex, entry in list(_llm_registry_dict.items()):
+      if not re.compile(regex).fullmatch(model):
+        continue
+      if isinstance(entry, tuple):
+        module_path, class_name = entry
+        try:
+          module = importlib.import_module(module_path)
+        except ImportError:
+          _llm_registry_dict.pop(regex, None)
+          continue
+        llm_class = getattr(module, class_name)
+        _llm_registry_dict[regex] = llm_class
         return llm_class
+      return entry
 
     # Provide helpful error messages for known patterns
     error_msg = f'Model {model} not found.'

src/google/adk/telemetry/_experimental_semconv.py

@@ -22,15 +22,19 @@
 import contextvars
 import json
 import os
+import sys
 from typing import Any
 from typing import Literal
+from typing import TYPE_CHECKING
 from typing import TypedDict
 
 from google.genai import types
 from google.genai.models import t as transformers
-from mcp import ClientSession as McpClientSession
-from mcp import Tool as McpTool
 from opentelemetry._logs import Logger
+
+if TYPE_CHECKING:
+  from mcp import ClientSession as McpClientSession
+  from mcp import Tool as McpTool
 from opentelemetry._logs import LogRecord
 from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_INPUT_MESSAGES
 from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_OUTPUT_MESSAGES
@@ -268,12 +272,16 @@ async def _to_tool_definitions(
   if callable(tool):
     return [_tool_definition_from_callable_tool(tool)]
 
-  if isinstance(tool, McpTool):
-    return [_tool_definition_from_mcp_tool(tool)]
+  if 'mcp' in sys.modules:
+    from mcp import ClientSession as McpClientSession
+    from mcp import Tool as McpTool
+
+    if isinstance(tool, McpTool):
+      return [_tool_definition_from_mcp_tool(tool)]
 
-  if isinstance(tool, McpClientSession):
-    result = await tool.list_tools()
-    return [_model_dump_to_tool_definition(t) for t in result.tools]
+    if isinstance(tool, McpClientSession):
+      result = await tool.list_tools()
+      return [_model_dump_to_tool_definition(t) for t in result.tools]
 
   return [
       GenericToolDefinition(

src/google/adk/tools/tool_context.py

@@ -14,17 +14,30 @@
 
 from __future__ import annotations
 
-# Keep CallbackContext for backward compatibility
+import importlib
+from typing import TYPE_CHECKING
+
 from ..agents.callback_context import CallbackContext
 from ..agents.context import Context
-# Keep AuthCredential for backward compatibility
-from ..auth.auth_credential import AuthCredential
-# Keep AuthHandler for backward compatibility
-from ..auth.auth_handler import AuthHandler
-# Keep AuthConfig for backward compatibility
-from ..auth.auth_tool import AuthConfig
-# Keep ToolConfirmation for backward compatibility
 from .tool_confirmation import ToolConfirmation
 
-# ToolContext is unified into Context
+if TYPE_CHECKING:
+  from ..auth.auth_credential import AuthCredential
+  from ..auth.auth_handler import AuthHandler
+  from ..auth.auth_tool import AuthConfig
+
 ToolContext = Context
+
+_LAZY_REEXPORTS: dict[str, tuple[str, str]] = {
+    'AuthCredential': ('google.adk.auth.auth_credential', 'AuthCredential'),
+    'AuthHandler': ('google.adk.auth.auth_handler', 'AuthHandler'),
+    'AuthConfig': ('google.adk.auth.auth_tool', 'AuthConfig'),
+}
+
+
+def __getattr__(name: str):
+  if name in _LAZY_REEXPORTS:
+    module_path, attr = _LAZY_REEXPORTS[name]
+    module = importlib.import_module(module_path)
+    return getattr(module, attr)
+  raise AttributeError(f'module {__name__!r} has no attribute {name!r}')