From 44d45fe9cd236664d7db7c77d5208f85fe000a95 Mon Sep 17 00:00:00 2001 From: Liang Wu Date: Thu, 6 Nov 2025 10:40:20 -0800 Subject: [PATCH] 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 PiperOrigin-RevId: 829017293 --- .../agent_engine_sandbox_code_executor.py | 6 ++++-- .../adk/code_executors/vertex_ai_code_executor.py | 3 ++- src/google/adk/evaluation/response_evaluator.py | 7 +++---- src/google/adk/evaluation/safety_evaluator.py | 7 +++---- src/google/adk/evaluation/vertex_ai_eval_facade.py | 11 +++++++---- src/google/adk/examples/vertex_ai_example_store.py | 3 ++- .../adk/memory/vertex_ai_memory_bank_service.py | 3 ++- .../adk/memory/vertex_ai_rag_memory_service.py | 4 +++- src/google/adk/sessions/vertex_ai_session_service.py | 12 ++++++------ .../adk/tools/retrieval/vertex_ai_rag_retrieval.py | 3 ++- .../test_agent_engine_sandbox_code_executor.py | 8 +++----- .../memory/test_vertex_ai_memory_bank_service.py | 4 +--- 12 files changed, 38 insertions(+), 33 deletions(-) diff --git a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py b/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py index 135c56bc..2e3e978b 100644 --- a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py +++ b/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( diff --git a/src/google/adk/code_executors/vertex_ai_code_executor.py b/src/google/adk/code_executors/vertex_ai_code_executor.py index 68c32526..a6a0ec8e 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/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: diff --git a/src/google/adk/evaluation/response_evaluator.py b/src/google/adk/evaluation/response_evaluator.py index d51c3829..685222f2 100644 --- a/src/google/adk/evaluation/response_evaluator.py +++ b/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: diff --git a/src/google/adk/evaluation/safety_evaluator.py b/src/google/adk/evaluation/safety_evaluator.py index 614a6606..e85d62dd 100644 --- a/src/google/adk/evaluation/safety_evaluator.py +++ b/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) diff --git a/src/google/adk/evaluation/vertex_ai_eval_facade.py b/src/google/adk/evaluation/vertex_ai_eval_facade.py index f80b7aba..bddcbe53 100644 --- a/src/google/adk/evaluation/vertex_ai_eval_facade.py +++ b/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), diff --git a/src/google/adk/examples/vertex_ai_example_store.py b/src/google/adk/examples/vertex_ai_example_store.py index 92a6de7e..75a7b789 100644 --- a/src/google/adk/examples/vertex_ai_example_store.py +++ b/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 = { diff --git a/src/google/adk/memory/vertex_ai_memory_bank_service.py b/src/google/adk/memory/vertex_ai_memory_bank_service.py index 6667c2de..b8f434c5 100644 --- a/src/google/adk/memory/vertex_ai_memory_bank_service.py +++ b/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, diff --git a/src/google/adk/memory/vertex_ai_rag_memory_service.py b/src/google/adk/memory/vertex_ai_rag_memory_service.py index dd491a62..236bf4b5 100644 --- a/src/google/adk/memory/vertex_ai_rag_memory_service.py +++ b/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( diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index e0651f0e..6fd0dfc7 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/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, diff --git a/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py b/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py index 15cf3b00..b0acc0fe 100644 --- a/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py +++ b/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'], diff --git a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py index b1f92295..c1bfd29a 100644 --- a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py +++ b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py @@ -58,17 +58,15 @@ class TestAgentEngineSandboxCodeExecutor: ), ) - @patch( - "google.adk.code_executors.agent_engine_sandbox_code_executor.vertexai" - ) + @patch("vertexai.Client") def test_execute_code_success( self, - mock_vertexai, + mock_vertexai_client, mock_invocation_context, ): # Setup Mocks mock_api_client = MagicMock() - mock_vertexai.Client.return_value = mock_api_client + mock_vertexai_client.return_value = mock_api_client mock_response = MagicMock() mock_json_output = MagicMock() mock_json_output.mime_type = "application/json" diff --git a/tests/unittests/memory/test_vertex_ai_memory_bank_service.py b/tests/unittests/memory/test_vertex_ai_memory_bank_service.py index 6a1f0ccb..77e22c94 100644 --- a/tests/unittests/memory/test_vertex_ai_memory_bank_service.py +++ b/tests/unittests/memory/test_vertex_ai_memory_bank_service.py @@ -87,9 +87,7 @@ def mock_vertex_ai_memory_bank_service( @pytest.fixture def mock_vertexai_client(): - with mock.patch( - 'google.adk.memory.vertex_ai_memory_bank_service.vertexai.Client' - ) as mock_client_constructor: + with mock.patch('vertexai.Client') as mock_client_constructor: mock_client = mock.MagicMock() mock_client.agent_engines.memories.generate = mock.MagicMock() mock_client.agent_engines.memories.retrieve = mock.MagicMock()