feat: Add warning on vertexai.Client instantiation and recommend agentplatform.Client

PiperOrigin-RevId: 949154018

Author: sararob

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)

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):