From 75ecb3b5c6cb610c8fec5c72c698996c0919e962 Mon Sep 17 00:00:00 2001 From: Amy Wu Date: Fri, 18 Sep 2026 11:50:29 -0700 Subject: [PATCH] fix: Resolve vertexai.types against agentplatform.types PiperOrigin-RevId: 983982589 --- agentplatform/__init__.py | 13 ++- tests/unit/vertexai/test_autorater_yaml.py | 8 ++ tests/unit/vertexai/test_types_alias.py | 122 +++++++++++++++++++++ vertexai/__init__.py | 12 +- vertexai/types.py | 84 ++++++++++++++ 5 files changed, 229 insertions(+), 10 deletions(-) create mode 100644 tests/unit/vertexai/test_types_alias.py create mode 100644 vertexai/types.py diff --git a/agentplatform/__init__.py b/agentplatform/__init__.py index feb8436449..fc5923a4b0 100644 --- a/agentplatform/__init__.py +++ b/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}'") diff --git a/tests/unit/vertexai/test_autorater_yaml.py b/tests/unit/vertexai/test_autorater_yaml.py index 669f8c0851..fc767b8a46 100644 --- a/tests/unit/vertexai/test_autorater_yaml.py +++ b/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): diff --git a/tests/unit/vertexai/test_types_alias.py b/tests/unit/vertexai/test_types_alias.py new file mode 100644 index 0000000000..73c4eb20cd --- /dev/null +++ b/tests/unit/vertexai/test_types_alias.py @@ -0,0 +1,122 @@ +# -*- 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 agentplatform +import vertexai +from agentplatform._genai import types as agentplatform_types +from vertexai._genai import types as legacy_types + +import sys + +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 diff --git a/vertexai/__init__.py b/vertexai/__init__.py index 4a7c313a84..3405d2e4c5 100644 --- a/vertexai/__init__.py +++ b/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}'") diff --git a/vertexai/types.py b/vertexai/types.py new file mode 100644 index 0000000000..b1c00f05b2 --- /dev/null +++ b/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` and does not type check, 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 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 + +if _typing.TYPE_CHECKING: + # Read statically and never executed, so every name below is a real class + # to a type checker instead of `__getattr__`'s return type. `agentplatform` + # comes second so it wins the names both modules define, which is the order + # `__getattr__` applies at runtime. + from vertexai._genai.types import * # noqa: F401,F403 + from agentplatform._genai.types import * # noqa: F401,F403 + +_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__))