Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions agentplatform/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}'")
Expand Down
8 changes: 8 additions & 0 deletions tests/unit/vertexai/test_autorater_yaml.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
121 changes: 121 additions & 0 deletions tests/unit/vertexai/test_types_alias.py
Original file line number Diff line number Diff line change
@@ -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
12 changes: 4 additions & 8 deletions vertexai/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
"""The vertexai module."""

import importlib
import sys

from google.cloud.aiplatform import version as aiplatform_version

Expand All @@ -24,7 +23,6 @@
from google.cloud.aiplatform import init

_genai_client = None
_genai_types = None


def __getattr__(name): # type: ignore[no-untyped-def]
Expand All @@ -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}'")

Expand Down
84 changes: 84 additions & 0 deletions vertexai/types.py
Original file line number Diff line number Diff line change
@@ -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__))
Loading