chore: enable 'from vertexai.types import *' for GenAI client

sararob · copybara-github · commit c1b3761a814f · 2025-06-17T08:27:48.000-07:00
PiperOrigin-RevId: 772477280
diff --git a/tests/unit/vertexai/genai/test_genai_client.py b/tests/unit/vertexai/genai/test_genai_client.py
@@ -55,3 +55,8 @@ def test_genai_client(self):
     async def test_async_client(self):
         test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
         assert isinstance(test_client.aio, vertexai._genai.client.AsyncClient)
+
+    @pytest.mark.usefixtures("google_auth_mock")
+    def test_types(self):
+        assert vertexai.types is not None
+        assert vertexai.types.LLMMetric is not None
diff --git a/vertexai/__init__.py b/vertexai/__init__.py
@@ -15,6 +15,7 @@
 """The vertexai module."""
 
 import importlib
+import sys
 
 from google.cloud.aiplatform import version as aiplatform_version
 
@@ -47,6 +48,8 @@ def __getattr__(name):
         global _genai_types
         if _genai_types is None:
             _genai_types = importlib.import_module("._genai.types", __name__)
+        if "vertexai.types" not in sys.modules:
+            sys.modules["vertexai.types"] = _genai_types
         return _genai_types
 
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
