Skip to content
Merged
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
122 changes: 122 additions & 0 deletions tests/unit/vertexai/test_types_alias.py
Original file line number Diff line number Diff line change
@@ -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
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` 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__))
Loading