From 1aa9bb13b3e92bf35b5961a0884a5be98cd5f951 Mon Sep 17 00:00:00 2001 From: Ankur Sharma Date: Tue, 28 Oct 2025 15:46:06 -0700 Subject: [PATCH] chore: Update vertex ai in example store and rag retrieval to use proxy PiperOrigin-RevId: 825237803 --- src/google/adk/dependencies/vertexai.py | 1 + src/google/adk/examples/vertex_ai_example_store.py | 4 +++- src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/google/adk/dependencies/vertexai.py b/src/google/adk/dependencies/vertexai.py index 0fb06088..4f254d87 100644 --- a/src/google/adk/dependencies/vertexai.py +++ b/src/google/adk/dependencies/vertexai.py @@ -15,4 +15,5 @@ from __future__ import annotations import vertexai +from vertexai.preview import example_stores from vertexai.preview import rag diff --git a/src/google/adk/examples/vertex_ai_example_store.py b/src/google/adk/examples/vertex_ai_example_store.py index 718003ae..92a6de7e 100644 --- a/src/google/adk/examples/vertex_ai_example_store.py +++ b/src/google/adk/examples/vertex_ai_example_store.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + from google.genai import types from typing_extensions import override -from vertexai.preview import example_stores +from ..dependencies.vertexai import example_stores from .base_example_provider import BaseExampleProvider from .example import Example 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 7edce4f9..15cf3b00 100644 --- a/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py +++ b/src/google/adk/tools/retrieval/vertex_ai_rag_retrieval.py @@ -22,8 +22,8 @@ from typing import TYPE_CHECKING from google.genai import types from typing_extensions import override -from vertexai.preview import rag +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