Skip to content

Commit 75ecb3b

Browse files
yinghsienwucopybara-github
authored andcommitted
fix: Resolve vertexai.types against agentplatform.types
PiperOrigin-RevId: 983982589
1 parent 6ad6fe8 commit 75ecb3b

5 files changed

Lines changed: 229 additions & 10 deletions

File tree

‎agentplatform/__init__.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,17 @@ def __getattr__(name): # type: ignore[no-untyped-def]
5454
global _genai_types
5555
if _genai_types is None:
5656
_genai_types = importlib.import_module("._genai.types", __name__)
57-
if "vertexai.types" not in sys.modules:
58-
sys.modules["vertexai.types"] = _genai_types
57+
# `types` is an alias for `._genai.types` rather than a real submodule,
58+
# so register it to keep
59+
# `from agentplatform.types import TypeName`
60+
# working without a prior attribute access. Spell it the google3 way:
61+
# Copybara rewrites that prefix in both directions, and the external
62+
# spelling fails its reversibility check. This key was `vertexai.types`,
63+
# which both misnamed this package's own alias and shadowed the real
64+
# `vertexai.types` module.
65+
types_module_name = f"{__name__}.types"
66+
if types_module_name not in sys.modules:
67+
sys.modules[types_module_name] = _genai_types
5968
return _genai_types
6069

6170
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

‎tests/unit/vertexai/test_autorater_yaml.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,14 @@ def setup_method(self):
160160
vertexai.init(
161161
project=_TEST_PROJECT,
162162
location=_TEST_LOCATION,
163+
# Set explicitly rather than inherited. `init` leaves an already
164+
# configured credential in place, and test_extensions installs a
165+
# `Mock(spec=AnonymousCredentials)` globally that it never resets;
166+
# `_upload_string_to_gcs` then hands `global_config.credentials`
167+
# to `storage.Client`, which rejects the Mock's `universe_domain`.
168+
# Whether that leak reaches this module depends on how
169+
# `--dist=loadscope` happens to assign modules to xdist workers.
170+
credentials=auth_credentials.AnonymousCredentials(),
163171
)
164172

165173
def teardown_method(self):
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# -*- coding: utf-8 -*-
2+
3+
# Copyright 2026 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
"""Unit tests for `vertexai.types` resolving against `agentplatform.types`."""
18+
19+
import agentplatform
20+
import vertexai
21+
from agentplatform._genai import types as agentplatform_types
22+
from vertexai._genai import types as legacy_types
23+
24+
import sys
25+
26+
import pytest
27+
28+
# Both modules resolve these out of `_evals_metric_loaders` on first access, so
29+
# they are checked on their own rather than in the bulk comparisons below.
30+
_LAZY_NAMES = frozenset({"PrebuiltMetric", "RubricMetric"})
31+
32+
_SHARED_NAMES = sorted(
33+
(set(agentplatform_types.__all__) & set(legacy_types.__all__)) - _LAZY_NAMES
34+
)
35+
_LEGACY_ONLY_NAMES = sorted(
36+
set(legacy_types.__all__) - set(agentplatform_types.__all__)
37+
)
38+
_AGENT_PLATFORM_ONLY_NAMES = sorted(
39+
set(agentplatform_types.__all__) - set(legacy_types.__all__)
40+
)
41+
42+
43+
def test_vertexai_types_is_the_alias_module():
44+
assert sys.modules[vertexai.types.__name__] is vertexai.types
45+
assert vertexai.types is not legacy_types
46+
assert vertexai.types is not agentplatform_types
47+
48+
49+
def test_shared_names_are_the_agentplatform_objects():
50+
"""The point of the alias: one class per message rather than two."""
51+
assert _SHARED_NAMES, "expected the two modules to share generated types"
52+
mismatched = [
53+
name
54+
for name in _SHARED_NAMES
55+
if getattr(vertexai.types, name) is not getattr(agentplatform_types, name)
56+
]
57+
assert not mismatched
58+
59+
60+
def test_memory_profile_survives_an_isinstance_check():
61+
"""Regression: an agentplatform object checked against the vertexai name."""
62+
profile = agentplatform_types.MemoryProfile(schema_id="user-profile", profile={})
63+
assert isinstance(profile, vertexai.types.MemoryProfile)
64+
65+
66+
def test_shared_types_still_nest_inside_a_legacy_only_model():
67+
"""The two halves have to interoperate, not just coexist.
68+
69+
`GenerateAgentEngineMemoriesConfig` stays a `vertexai` class because
70+
agentplatform renamed it, but its `metadata` values are now agentplatform
71+
objects. Pydantic accepts them because both models derive from the genai
72+
`BaseModel`, which sets `from_attributes=True`.
73+
"""
74+
config = vertexai.types.GenerateAgentEngineMemoriesConfig(
75+
metadata={"record": vertexai.types.MemoryMetadataValue(string_value="123")}
76+
)
77+
assert type(config) is legacy_types.GenerateAgentEngineMemoriesConfig
78+
assert config.metadata["record"].string_value == "123"
79+
80+
81+
def test_legacy_only_names_still_resolve():
82+
"""The `AgentEngine*` surface agentplatform renamed is not taken away."""
83+
assert "AgentEngine" in _LEGACY_ONLY_NAMES
84+
mismatched = [
85+
name
86+
for name in _LEGACY_ONLY_NAMES
87+
if getattr(vertexai.types, name) is not getattr(legacy_types, name)
88+
]
89+
assert not mismatched
90+
91+
92+
def test_agentplatform_only_names_are_reachable():
93+
assert "Runtime" in _AGENT_PLATFORM_ONLY_NAMES
94+
mismatched = [
95+
name
96+
for name in _AGENT_PLATFORM_ONLY_NAMES
97+
if getattr(vertexai.types, name) is not getattr(agentplatform_types, name)
98+
]
99+
assert not mismatched
100+
101+
102+
@pytest.mark.parametrize("name", sorted(_LAZY_NAMES))
103+
def test_lazily_loaded_metric_names_resolve_to_agentplatform(name):
104+
assert getattr(vertexai.types, name) is getattr(agentplatform_types, name)
105+
106+
107+
def test_all_is_the_union_of_both_modules():
108+
expected = set(agentplatform_types.__all__) | set(legacy_types.__all__)
109+
assert set(vertexai.types.__all__) == expected
110+
assert set(dir(vertexai.types)) == expected
111+
112+
113+
def test_unknown_name_still_raises_attribute_error():
114+
with pytest.raises(AttributeError):
115+
_ = vertexai.types.NoSuchTypeName
116+
117+
118+
def test_agentplatform_registers_its_own_types_alias():
119+
"""`agentplatform.types` must no longer claim the `vertexai.types` key."""
120+
assert agentplatform.types is agentplatform_types
121+
assert sys.modules[f"{agentplatform.__name__}.types"] is agentplatform_types
122+
assert sys.modules[f"{vertexai.__name__}.types"] is not agentplatform_types

‎vertexai/__init__.py‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
"""The vertexai module."""
1616

1717
import importlib
18-
import sys
1918

2019
from google.cloud.aiplatform import version as aiplatform_version
2120

@@ -24,7 +23,6 @@
2423
from google.cloud.aiplatform import init
2524

2625
_genai_client = None
27-
_genai_types = None
2826

2927

3028
def __getattr__(name): # type: ignore[no-untyped-def]
@@ -45,12 +43,10 @@ def __getattr__(name): # type: ignore[no-untyped-def]
4543
return getattr(_genai_client, name)
4644

4745
if name == "types":
48-
global _genai_types
49-
if _genai_types is None:
50-
_genai_types = importlib.import_module("._genai.types", __name__)
51-
if "vertexai.types" not in sys.modules:
52-
sys.modules["vertexai.types"] = _genai_types
53-
return _genai_types
46+
# `types` is a real submodule that resolves against
47+
# `agentplatform.types`, so importing it also binds it as an attribute
48+
# here and this runs only once.
49+
return importlib.import_module(".types", __name__)
5450

5551
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
5652

‎vertexai/types.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
"""`vertexai.types` resolved against `agentplatform.types`.
16+
17+
Agent Platform re-generated the Gen AI types under `agentplatform` rather than
18+
moving them, so `vertexai.types.Memory` and `agentplatform.types.Memory` were
19+
two distinct classes carrying identical definitions. Code that hands an
20+
`agentplatform` object to a caller written against the `vertexai` name fails
21+
`isinstance` and does not type check, even though both describe the same
22+
message.
23+
24+
This module removes the split. Every name `agentplatform.types` defines
25+
resolves here to the `agentplatform` class, so the two spellings are one
26+
object. The names that only ever existed under `vertexai` -- the `AgentEngine*`
27+
surface Agent Platform renamed to `Runtime*` and `MemoryBank*` -- keep
28+
resolving to their `vertexai._genai.types` classes, so nothing is taken away
29+
from callers that have not migrated.
30+
31+
The dependency runs one way: `vertexai` reaches into `agentplatform`, never the
32+
reverse. `google-cloud-agentplatform` ships `agentplatform` without `vertexai`
33+
and without the generated clients, and nothing here changes that.
34+
35+
Resolution stays lazy at each step. `vertexai.__getattr__` defers importing
36+
this module until `vertexai.types` is touched; `vertexai._genai.types` is
37+
imported only if a name is not found in `agentplatform`; and neither module's
38+
`PrebuiltMetric`/`RubricMetric` is resolved until asked for, so the evaluation
39+
dependencies are still not pulled in by a bare import.
40+
"""
41+
42+
from __future__ import annotations
43+
44+
import importlib as _importlib
45+
import types as _module_types
46+
import typing as _typing
47+
48+
from agentplatform._genai import types as _agentplatform_types
49+
50+
if _typing.TYPE_CHECKING:
51+
# Read statically and never executed, so every name below is a real class
52+
# to a type checker instead of `__getattr__`'s return type. `agentplatform`
53+
# comes second so it wins the names both modules define, which is the order
54+
# `__getattr__` applies at runtime.
55+
from vertexai._genai.types import * # noqa: F401,F403
56+
from agentplatform._genai.types import * # noqa: F401,F403
57+
58+
_legacy_types: _module_types.ModuleType | None = None
59+
60+
61+
def _get_legacy_types() -> _module_types.ModuleType:
62+
"""Imports `vertexai._genai.types` on first use."""
63+
global _legacy_types
64+
if _legacy_types is None:
65+
_legacy_types = _importlib.import_module("._genai.types", __package__)
66+
return _legacy_types
67+
68+
69+
def __getattr__(name: str) -> _typing.Any:
70+
# See https://peps.python.org/pep-0562/
71+
if name == "__all__":
72+
return __dir__()
73+
try:
74+
return getattr(_agentplatform_types, name)
75+
except AttributeError:
76+
pass
77+
try:
78+
return getattr(_get_legacy_types(), name)
79+
except AttributeError:
80+
raise AttributeError(f"module '{__name__}' has no attribute '{name}'") from None
81+
82+
83+
def __dir__() -> list[str]:
84+
return sorted(set(_agentplatform_types.__all__) | set(_get_legacy_types().__all__))

0 commit comments

Comments
 (0)