chore: Lazy load Vertex AI dependencies in ADK modules

This is about 35% decrease. This change refactors several ADK modules to import `vertexai` and its submodules only when they are first used, rather than at the top of the file. This improves module load times by avoiding unnecessary imports of large dependencies. Imports are also placed within `if TYPE_CHECKING:` blocks where appropriate.

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 829017293

src/google/adk/code_executors/agent_engine_sandbox_code_executor.py

@@ -21,8 +21,6 @@ import re
 from typing import Optional
 
 from typing_extensions import override
-import vertexai
-from vertexai import types
 
 from ..agents.invocation_context import InvocationContext
 from ..utils.feature_decorator import experimental
@@ -78,6 +76,8 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor):
           )
       )
     elif agent_engine_resource_name is not None:
+      from vertexai import types
+
       self._project_id, self._location = (
           self._get_project_id_and_location_from_resource_name(
               agent_engine_resource_name, agent_engine_resource_name_pattern
@@ -174,6 +174,8 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor):
     Returns:
       An API client for the given project and location.
     """
+    import vertexai
+
     return vertexai.Client(project=self._project_id, location=self._location)
 
   def _get_project_id_and_location_from_resource_name(

src/google/adk/code_executors/vertex_ai_code_executor.py

@@ -21,7 +21,6 @@ from typing import Any
 from typing import Optional
 
 from typing_extensions import override
-from vertexai.preview.extensions import Extension
 
 from ..agents.invocation_context import InvocationContext
 from .base_code_executor import BaseCodeExecutor
@@ -88,6 +87,8 @@ Total columns: {df.shape[1]}
 
 def _get_code_interpreter_extension(resource_name: str = None):
   """Returns: Load or create the code interpreter extension."""
+  from vertexai.preview.extensions import Extension
+
   if not resource_name:
     resource_name = os.environ.get('CODE_INTERPRETER_EXTENSION_NAME')
   if resource_name:

src/google/adk/evaluation/response_evaluator.py

@@ -18,7 +18,6 @@ from typing import Optional
 
 from typing_extensions import override
 
-from ..dependencies.vertexai import vertexai
 from .eval_case import Invocation
 from .eval_metrics import EvalMetric
 from .eval_metrics import Interval
@@ -30,8 +29,6 @@ from .evaluator import Evaluator
 from .final_response_match_v1 import RougeEvaluator
 from .vertex_ai_eval_facade import _VertexAiEvalFacade
 
-vertexai_types = vertexai.types
-
 
 class ResponseEvaluator(Evaluator):
   """Evaluates Agent's responses.
@@ -68,7 +65,9 @@ class ResponseEvaluator(Evaluator):
       metric_name = eval_metric.metric_name
 
     if PrebuiltMetrics.RESPONSE_EVALUATION_SCORE.value == metric_name:
-      self._metric_name = vertexai_types.PrebuiltMetric.COHERENCE
+      from ..dependencies.vertexai import vertexai
+
+      self._metric_name = vertexai.types.PrebuiltMetric.COHERENCE
     elif PrebuiltMetrics.RESPONSE_MATCH_SCORE.value == metric_name:
       self._metric_name = metric_name
     else:

src/google/adk/evaluation/safety_evaluator.py

@@ -18,7 +18,6 @@ from typing import Optional
 
 from typing_extensions import override
 
-from ..dependencies.vertexai import vertexai
 from .eval_case import Invocation
 from .eval_metrics import EvalMetric
 from .eval_metrics import Interval
@@ -29,8 +28,6 @@ from .evaluator import EvaluationResult
 from .evaluator import Evaluator
 from .vertex_ai_eval_facade import _VertexAiEvalFacade
 
-vertexai_types = vertexai.types
-
 
 class SafetyEvaluatorV1(Evaluator):
   """Evaluates safety (harmlessness) of an Agent's Response.
@@ -70,7 +67,9 @@ class SafetyEvaluatorV1(Evaluator):
       actual_invocations: list[Invocation],
       expected_invocations: Optional[list[Invocation]],
   ) -> EvaluationResult:
+    from ..dependencies.vertexai import vertexai
+
     return _VertexAiEvalFacade(
         threshold=self._eval_metric.threshold,
-        metric_name=vertexai_types.PrebuiltMetric.SAFETY,
+        metric_name=vertexai.types.PrebuiltMetric.SAFETY,
     ).evaluate_invocations(actual_invocations, expected_invocations)

src/google/adk/evaluation/vertex_ai_eval_facade.py

@@ -17,20 +17,20 @@ from __future__ import annotations
 import math
 import os
 from typing import Optional
+from typing import TYPE_CHECKING
 
 from google.genai import types as genai_types
 import pandas as pd
 from typing_extensions import override
 
-from ..dependencies.vertexai import vertexai
 from .eval_case import Invocation
 from .evaluator import EvalStatus
 from .evaluator import EvaluationResult
 from .evaluator import Evaluator
 from .evaluator import PerInvocationResult
 
-vertexai_types = vertexai.types
-VertexAiClient = vertexai.Client
+if TYPE_CHECKING:
+  from vertexai import types as vertexai_types
 
 _ERROR_MESSAGE_SUFFIX = """
 You should specify both project id and location. This metric uses Vertex Gen AI
@@ -162,7 +162,10 @@ class _VertexAiEvalFacade(Evaluator):
     if not location:
       raise ValueError("Missing location." + _ERROR_MESSAGE_SUFFIX)
 
-    client = VertexAiClient(project=project_id, location=location)
+    from vertexai import Client
+    from vertexai import types as vertexai_types
+
+    client = Client(project=project_id, location=location)
 
     return client.evals.evaluate(
         dataset=vertexai_types.EvaluationDataset(eval_dataset_df=dataset),

src/google/adk/examples/vertex_ai_example_store.py

@@ -17,7 +17,6 @@ from __future__ import annotations
 from google.genai import types
 from typing_extensions import override
 
-from ..dependencies.vertexai import example_stores
 from .base_example_provider import BaseExampleProvider
 from .example import Example
 
@@ -37,6 +36,8 @@ class VertexAiExampleStore(BaseExampleProvider):
 
   @override
   def get_examples(self, query: str) -> list[Example]:
+    from ..dependencies.vertexai import example_stores
+
     example_store = example_stores.ExampleStore(self.examples_store_name)
     # Retrieve relevant examples.
     request = {

src/google/adk/memory/vertex_ai_memory_bank_service.py

@@ -20,7 +20,6 @@ from typing import TYPE_CHECKING
 
 from google.genai import types
 from typing_extensions import override
-import vertexai
 
 from ..utils.vertex_ai_utils import get_express_mode_api_key
 from .base_memory_service import BaseMemoryService
@@ -138,6 +137,8 @@ class VertexAiMemoryBankService(BaseMemoryService):
     Returns:
       An API client for the given project and location or express mode api key.
     """
+    import vertexai
+
     return vertexai.Client(
         project=self._project,
         location=self._location,

src/google/adk/memory/vertex_ai_rag_memory_service.py

@@ -26,7 +26,6 @@ from google.genai import types
 from typing_extensions import override
 
 from . import _utils
-from ..dependencies.vertexai import rag
 from .base_memory_service import BaseMemoryService
 from .base_memory_service import SearchMemoryResponse
 from .memory_entry import MemoryEntry
@@ -93,6 +92,8 @@ class VertexAiRagMemoryService(BaseMemoryService):
     if not self._vertex_rag_store.rag_resources:
       raise ValueError("Rag resources must be set.")
 
+    from ..dependencies.vertexai import rag
+
     for rag_resource in self._vertex_rag_store.rag_resources:
       rag.upload_file(
           corpus_name=rag_resource.rag_corpus,
@@ -109,6 +110,7 @@ class VertexAiRagMemoryService(BaseMemoryService):
       self, *, app_name: str, user_id: str, query: str
   ) -> SearchMemoryResponse:
     """Searches for sessions that match the query using rag.retrieval_query."""
+    from ..dependencies.vertexai import rag
     from ..events.event import Event
 
     response = rag.retrieval_query(

src/google/adk/sessions/vertex_ai_session_service.py

@@ -20,16 +20,14 @@ import logging
 import re
 from typing import Any
 from typing import Optional
+from typing import TYPE_CHECKING
 from typing import Union
 
 from google.genai import types
-from google.genai.errors import ClientError
-from tenacity import retry
-from tenacity import retry_if_result
-from tenacity import stop_after_attempt
-from tenacity import wait_exponential
 from typing_extensions import override
-import vertexai
+
+if TYPE_CHECKING:
+  import vertexai
 
 from . import _session_util
 from ..events.event import Event
@@ -326,6 +324,8 @@ class VertexAiSessionService(BaseSessionService):
     Returns:
       An API client for the given project and location or express mode api key.
     """
+    import vertexai
+
     return vertexai.Client(
         project=self._project,
         location=self._location,

src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py

@@ -23,12 +23,12 @@ from typing import TYPE_CHECKING
 from google.genai import types
 from typing_extensions import override
 
-from ...dependencies.vertexai import rag
 from ...utils.model_name_utils import is_gemini_2_or_above
 from ..tool_context import ToolContext
 from .base_retrieval_tool import BaseRetrievalTool
 
 if TYPE_CHECKING:
+  from ...dependencies.vertexai import rag
   from ...models import LlmRequest
 
 logger = logging.getLogger('google_adk.' + __name__)
@@ -90,6 +90,7 @@ class VertexAiRagRetrieval(BaseRetrievalTool):
       args: dict[str, Any],
       tool_context: ToolContext,
   ) -> Any:
+    from ...dependencies.vertexai import rag
 
     response = rag.retrieval_query(
         text=args['query'],