diff --git a/tests/unit/agentplatform/genai/test_genai_client.py b/tests/unit/agentplatform/genai/test_genai_client.py index e3be385879..b479e8d716 100644 --- a/tests/unit/agentplatform/genai/test_genai_client.py +++ b/tests/unit/agentplatform/genai/test_genai_client.py @@ -17,12 +17,15 @@ import importlib import pytest +import warnings from unittest import mock from google.cloud import aiplatform import agentplatform from agentplatform._genai import client as agentplatform_client from google.cloud.aiplatform import initializer as aiplatform_initializer +import vertexai +from vertexai._genai import client as vertexai_client _TEST_PROJECT = "test-project" @@ -106,3 +109,16 @@ async def test_call_aclose_async_client(self): ).aio await async_client.aclose() mock_aclose.assert_called() + + @pytest.mark.usefixtures("google_auth_mock") + def test_vertexai_client_deprecation_warning(self): + + with mock.patch.object(vertexai_client, "_CLIENT_WARNING_SHOWN", False): + # Assert that the warning is triggered on the first instantiation + with pytest.warns(FutureWarning, match="The vertexai.Client class is deprecated"): + _ = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) + # Assert that the warning is NOT triggered on subsequent instantiations + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + _ = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) + _ = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) diff --git a/vertexai/_genai/client.py b/vertexai/_genai/client.py index 18b91a4413..7eef0d0712 100644 --- a/vertexai/_genai/client.py +++ b/vertexai/_genai/client.py @@ -16,6 +16,7 @@ import asyncio import importlib import sys +import warnings from typing import Optional, Union, TYPE_CHECKING from types import TracebackType, ModuleType @@ -43,6 +44,7 @@ _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules" +_CLIENT_WARNING_SHOWN = False def _custom_append_library_version_headers(headers: dict[str, str]) -> None: @@ -239,6 +241,14 @@ def __init__( http_options (Union[HttpOptions, HttpOptionsDict]): Http options to use for the client. """ + global _CLIENT_WARNING_SHOWN + if not _CLIENT_WARNING_SHOWN: + warnings.warn( + "The vertexai.Client class is deprecated. Please use agentplatform.Client instead.", + FutureWarning, + stacklevel=2, + ) + _CLIENT_WARNING_SHOWN = True self._debug_config = debug_config or genai_client.DebugConfig() if isinstance(http_options, dict):