fix: Resolve vertexai.types against agentplatform.types

PiperOrigin-RevId: 984095463

Author: yinghsienwu

agentplatform/__init__.py

@@ -54,8 +54,17 @@ def __getattr__(name):  # type: ignore[no-untyped-def]
         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
+        # `types` is an alias for `._genai.types` rather than a real submodule,
+        # so register it to keep
+        # `from agentplatform.types import TypeName`
+        # working without a prior attribute access. Spell it the google3 way:
+        # Copybara rewrites that prefix in both directions, and the external
+        # spelling fails its reversibility check. This key was `vertexai.types`,
+        # which both misnamed this package's own alias and shadowed the real
+        # `vertexai.types` module.
+        types_module_name = f"{__name__}.types"
+        if types_module_name not in sys.modules:
+            sys.modules[types_module_name] = _genai_types
         return _genai_types
 
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

tests/unit/vertexai/test_autorater_yaml.py

@@ -160,6 +160,14 @@ def setup_method(self):
         vertexai.init(
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
+            # Set explicitly rather than inherited. `init` leaves an already
+            # configured credential in place, and test_extensions installs a
+            # `Mock(spec=AnonymousCredentials)` globally that it never resets;
+            # `_upload_string_to_gcs` then hands `global_config.credentials`
+            # to `storage.Client`, which rejects the Mock's `universe_domain`.
+            # Whether that leak reaches this module depends on how
+            # `--dist=loadscope` happens to assign modules to xdist workers.
+            credentials=auth_credentials.AnonymousCredentials(),
         )
 
     def teardown_method(self):

tests/unit/vertexai/test_types_alias.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Unit tests for `vertexai.types` resolving against `agentplatform.types`."""
+
+import sys
+
+import agentplatform
+import vertexai
+from agentplatform._genai import types as agentplatform_types
+from vertexai._genai import types as legacy_types
+import pytest
+
+# Both modules resolve these out of `_evals_metric_loaders` on first access, so
+# they are checked on their own rather than in the bulk comparisons below.
+_LAZY_NAMES = frozenset({"PrebuiltMetric", "RubricMetric"})
+
+_SHARED_NAMES = sorted(
+    (set(agentplatform_types.__all__) & set(legacy_types.__all__)) - _LAZY_NAMES
+)
+_LEGACY_ONLY_NAMES = sorted(
+    set(legacy_types.__all__) - set(agentplatform_types.__all__)
+)
+_AGENT_PLATFORM_ONLY_NAMES = sorted(
+    set(agentplatform_types.__all__) - set(legacy_types.__all__)
+)
+
+
+def test_vertexai_types_is_the_alias_module():
+    assert sys.modules[vertexai.types.__name__] is vertexai.types
+    assert vertexai.types is not legacy_types
+    assert vertexai.types is not agentplatform_types
+
+
+def test_shared_names_are_the_agentplatform_objects():
+    """The point of the alias: one class per message rather than two."""
+    assert _SHARED_NAMES, "expected the two modules to share generated types"
+    mismatched = [
+        name
+        for name in _SHARED_NAMES
+        if getattr(vertexai.types, name) is not getattr(agentplatform_types, name)
+    ]
+    assert not mismatched
+
+
+def test_memory_profile_survives_an_isinstance_check():
+    """Regression: an agentplatform object checked against the vertexai name."""
+    profile = agentplatform_types.MemoryProfile(schema_id="user-profile", profile={})
+    assert isinstance(profile, vertexai.types.MemoryProfile)
+
+
+def test_shared_types_still_nest_inside_a_legacy_only_model():
+    """The two halves have to interoperate, not just coexist.
+
+    `GenerateAgentEngineMemoriesConfig` stays a `vertexai` class because
+    agentplatform renamed it, but its `metadata` values are now agentplatform
+    objects. Pydantic accepts them because both models derive from the genai
+    `BaseModel`, which sets `from_attributes=True`.
+    """
+    config = vertexai.types.GenerateAgentEngineMemoriesConfig(
+        metadata={"record": vertexai.types.MemoryMetadataValue(string_value="123")}
+    )
+    assert type(config) is legacy_types.GenerateAgentEngineMemoriesConfig
+    assert config.metadata["record"].string_value == "123"
+
+
+def test_legacy_only_names_still_resolve():
+    """The `AgentEngine*` surface agentplatform renamed is not taken away."""
+    assert "AgentEngine" in _LEGACY_ONLY_NAMES
+    mismatched = [
+        name
+        for name in _LEGACY_ONLY_NAMES
+        if getattr(vertexai.types, name) is not getattr(legacy_types, name)
+    ]
+    assert not mismatched
+
+
+def test_agentplatform_only_names_are_reachable():
+    assert "Runtime" in _AGENT_PLATFORM_ONLY_NAMES
+    mismatched = [
+        name
+        for name in _AGENT_PLATFORM_ONLY_NAMES
+        if getattr(vertexai.types, name) is not getattr(agentplatform_types, name)
+    ]
+    assert not mismatched
+
+
+@pytest.mark.parametrize("name", sorted(_LAZY_NAMES))
+def test_lazily_loaded_metric_names_resolve_to_agentplatform(name):
+    assert getattr(vertexai.types, name) is getattr(agentplatform_types, name)
+
+
+def test_all_is_the_union_of_both_modules():
+    expected = set(agentplatform_types.__all__) | set(legacy_types.__all__)
+    assert set(vertexai.types.__all__) == expected
+    assert set(dir(vertexai.types)) == expected
+
+
+def test_unknown_name_still_raises_attribute_error():
+    with pytest.raises(AttributeError):
+        _ = vertexai.types.NoSuchTypeName
+
+
+def test_agentplatform_registers_its_own_types_alias():
+    """`agentplatform.types` must no longer claim the `vertexai.types` key."""
+    assert agentplatform.types is agentplatform_types
+    assert sys.modules[f"{agentplatform.__name__}.types"] is agentplatform_types
+    assert sys.modules[f"{vertexai.__name__}.types"] is not agentplatform_types

vertexai/__init__.py

@@ -15,7 +15,6 @@
 """The vertexai module."""
 
 import importlib
-import sys
 
 from google.cloud.aiplatform import version as aiplatform_version
 
@@ -24,7 +23,6 @@
 from google.cloud.aiplatform import init
 
 _genai_client = None
-_genai_types = None
 
 
 def __getattr__(name):  # type: ignore[no-untyped-def]
@@ -45,12 +43,10 @@ def __getattr__(name):  # type: ignore[no-untyped-def]
         return getattr(_genai_client, name)
 
     if name == "types":
-        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
+        # `types` is a real submodule that resolves against
+        # `agentplatform.types`, so importing it also binds it as an attribute
+        # here and this runs only once.
+        return importlib.import_module(".types", __name__)
 
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 

vertexai/types.py

@@ -0,0 +1,84 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""`vertexai.types` resolved against `agentplatform.types`.
+
+Agent Platform re-generated the Gen AI types under `agentplatform` rather than
+moving them, so `vertexai.types.Memory` and `agentplatform.types.Memory` were
+two distinct classes carrying identical definitions. Code that hands an
+`agentplatform` object to a caller written against the `vertexai` name fails
+`isinstance`, even though both describe the same message.
+
+This module removes the split. Every name `agentplatform.types` defines
+resolves here to the `agentplatform` class, so the two spellings are one
+object. The names that only ever existed under `vertexai` -- the `AgentEngine*`
+surface Agent Platform renamed to `Runtime*` and `MemoryBank*` -- keep
+resolving to their `vertexai._genai.types` classes, so nothing is taken away
+from callers that have not migrated.
+
+The unification is a runtime one. Names resolve through `__getattr__`, so a
+type checker still reads `vertexai.types.X` as `Any`, exactly as it did while
+this name was a lazy alias for `vertexai._genai.types`. Re-exporting both
+modules under `typing.TYPE_CHECKING` would make them resolve statically too,
+but it also makes the checker read the generated `Optional` fields for the
+first time, which turns long-standing unguarded accesses in downstream callers
+into new type errors. Static resolution is being rolled out separately, once
+those callers are fixed.
+
+The dependency runs one way: `vertexai` reaches into `agentplatform`, never the
+reverse. `google-cloud-agentplatform` ships `agentplatform` without `vertexai`
+and without the generated clients, and nothing here changes that.
+
+Resolution stays lazy at each step. `vertexai.__getattr__` defers importing
+this module until `vertexai.types` is touched; `vertexai._genai.types` is
+imported only if a name is not found in `agentplatform`; and neither module's
+`PrebuiltMetric`/`RubricMetric` is resolved until asked for, so the evaluation
+dependencies are still not pulled in by a bare import.
+"""
+
+from __future__ import annotations
+
+import importlib as _importlib
+import types as _module_types
+import typing as _typing
+
+from agentplatform._genai import types as _agentplatform_types
+
+_legacy_types: _module_types.ModuleType | None = None
+
+
+def _get_legacy_types() -> _module_types.ModuleType:
+    """Imports `vertexai._genai.types` on first use."""
+    global _legacy_types
+    if _legacy_types is None:
+        _legacy_types = _importlib.import_module("._genai.types", __package__)
+    return _legacy_types
+
+
+def __getattr__(name: str) -> _typing.Any:
+    # See https://peps.python.org/pep-0562/
+    if name == "__all__":
+        return __dir__()
+    try:
+        return getattr(_agentplatform_types, name)
+    except AttributeError:
+        pass
+    try:
+        return getattr(_get_legacy_types(), name)
+    except AttributeError:
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'") from None
+
+
+def __dir__() -> list[str]:
+    return sorted(set(_agentplatform_types.__all__) | set(_get_legacy_types().__all__))