Skip to content

Commit 116a051

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
chore: disable default-on telemetry enablement in ADK on AE
PiperOrigin-RevId: 828054121
1 parent 1513d8f commit 116a051

File tree

4 files changed

+146
-144
lines changed

4 files changed

+146
-144
lines changed

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 97 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
from google.auth import credentials as auth_credentials
2727
from google.cloud import storage
2828
import vertexai
29-
from google.cloud import aiplatform
3029
from google.cloud.aiplatform_v1 import types as aip_types
3130
from google.cloud.aiplatform_v1.services import reasoning_engine_service
3231
from google.cloud.aiplatform import base
@@ -1012,99 +1011,100 @@ def update_agent_engine_mock():
10121011
yield update_agent_engine_mock
10131012

10141013

1015-
@pytest.mark.usefixtures("google_auth_mock")
1016-
class TestAgentEngines:
1017-
def setup_method(self):
1018-
importlib.reload(initializer)
1019-
importlib.reload(aiplatform)
1020-
aiplatform.init(
1021-
project=_TEST_PROJECT,
1022-
location=_TEST_LOCATION,
1023-
credentials=_TEST_CREDENTIALS,
1024-
staging_bucket=_TEST_STAGING_BUCKET,
1025-
)
1026-
1027-
def teardown_method(self):
1028-
initializer.global_pool.shutdown(wait=True)
1029-
1030-
@pytest.mark.parametrize(
1031-
"env_vars,expected_env_vars",
1032-
[
1033-
({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}),
1034-
(None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}),
1035-
(
1036-
{"some_env": "some_val"},
1037-
{
1038-
"some_env": "some_val",
1039-
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true",
1040-
},
1041-
),
1042-
(
1043-
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1044-
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1045-
),
1046-
],
1047-
)
1048-
def test_create_default_telemetry_enablement(
1049-
self,
1050-
create_agent_engine_mock: mock.Mock,
1051-
cloud_storage_create_bucket_mock: mock.Mock,
1052-
cloudpickle_dump_mock: mock.Mock,
1053-
cloudpickle_load_mock: mock.Mock,
1054-
get_gca_resource_mock: mock.Mock,
1055-
env_vars: dict[str, str],
1056-
expected_env_vars: dict[str, str],
1057-
):
1058-
agent_engines.create(
1059-
agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT),
1060-
env_vars=env_vars,
1061-
)
1062-
create_agent_engine_mock.assert_called_once()
1063-
deployment_spec = create_agent_engine_mock.call_args.kwargs[
1064-
"reasoning_engine"
1065-
].spec.deployment_spec
1066-
assert _utils.to_dict(deployment_spec)["env"] == [
1067-
{"name": key, "value": value} for key, value in expected_env_vars.items()
1068-
]
1069-
1070-
@pytest.mark.parametrize(
1071-
"env_vars,expected_env_vars",
1072-
[
1073-
({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}),
1074-
(None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}),
1075-
(
1076-
{"some_env": "some_val"},
1077-
{
1078-
"some_env": "some_val",
1079-
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true",
1080-
},
1081-
),
1082-
(
1083-
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1084-
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1085-
),
1086-
],
1087-
)
1088-
def test_update_default_telemetry_enablement(
1089-
self,
1090-
update_agent_engine_mock: mock.Mock,
1091-
cloud_storage_create_bucket_mock: mock.Mock,
1092-
cloudpickle_dump_mock: mock.Mock,
1093-
cloudpickle_load_mock: mock.Mock,
1094-
get_gca_resource_mock: mock.Mock,
1095-
get_agent_engine_mock: mock.Mock,
1096-
env_vars: dict[str, str],
1097-
expected_env_vars: dict[str, str],
1098-
):
1099-
agent_engines.update(
1100-
resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1101-
description="foobar", # avoid "At least one of ... must be specified" errors.
1102-
env_vars=env_vars,
1103-
)
1104-
update_agent_engine_mock.assert_called_once()
1105-
deployment_spec = update_agent_engine_mock.call_args.kwargs[
1106-
"request"
1107-
].reasoning_engine.spec.deployment_spec
1108-
assert _utils.to_dict(deployment_spec)["env"] == [
1109-
{"name": key, "value": value} for key, value in expected_env_vars.items()
1110-
]
1014+
# TODO(jawoszek): Uncomment once we're ready for default-on.
1015+
# @pytest.mark.usefixtures("google_auth_mock")
1016+
# class TestAgentEngines:
1017+
# def setup_method(self):
1018+
# importlib.reload(initializer)
1019+
# importlib.reload(aiplatform)
1020+
# aiplatform.init(
1021+
# project=_TEST_PROJECT,
1022+
# location=_TEST_LOCATION,
1023+
# credentials=_TEST_CREDENTIALS,
1024+
# staging_bucket=_TEST_STAGING_BUCKET,
1025+
# )
1026+
1027+
# def teardown_method(self):
1028+
# initializer.global_pool.shutdown(wait=True)
1029+
1030+
# @pytest.mark.parametrize(
1031+
# "env_vars,expected_env_vars",
1032+
# [
1033+
# ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}),
1034+
# (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}),
1035+
# (
1036+
# {"some_env": "some_val"},
1037+
# {
1038+
# "some_env": "some_val",
1039+
# GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false",
1040+
# },
1041+
# ),
1042+
# (
1043+
# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1044+
# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1045+
# ),
1046+
# ],
1047+
# )
1048+
# def test_create_default_telemetry_enablement(
1049+
# self,
1050+
# create_agent_engine_mock: mock.Mock,
1051+
# cloud_storage_create_bucket_mock: mock.Mock,
1052+
# cloudpickle_dump_mock: mock.Mock,
1053+
# cloudpickle_load_mock: mock.Mock,
1054+
# get_gca_resource_mock: mock.Mock,
1055+
# env_vars: dict[str, str],
1056+
# expected_env_vars: dict[str, str],
1057+
# ):
1058+
# agent_engines.create(
1059+
# agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT),
1060+
# env_vars=env_vars,
1061+
# )
1062+
# create_agent_engine_mock.assert_called_once()
1063+
# deployment_spec = create_agent_engine_mock.call_args.kwargs[
1064+
# "reasoning_engine"
1065+
# ].spec.deployment_spec
1066+
# assert _utils.to_dict(deployment_spec)["env"] == [
1067+
# {"name": key, "value": value} for key, value in expected_env_vars.items()
1068+
# ]
1069+
1070+
# @pytest.mark.parametrize(
1071+
# "env_vars,expected_env_vars",
1072+
# [
1073+
# ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}),
1074+
# (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}),
1075+
# (
1076+
# {"some_env": "some_val"},
1077+
# {
1078+
# "some_env": "some_val",
1079+
# GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false",
1080+
# },
1081+
# ),
1082+
# (
1083+
# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1084+
# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1085+
# ),
1086+
# ],
1087+
# )
1088+
# def test_update_default_telemetry_enablement(
1089+
# self,
1090+
# update_agent_engine_mock: mock.Mock,
1091+
# cloud_storage_create_bucket_mock: mock.Mock,
1092+
# cloudpickle_dump_mock: mock.Mock,
1093+
# cloudpickle_load_mock: mock.Mock,
1094+
# get_gca_resource_mock: mock.Mock,
1095+
# get_agent_engine_mock: mock.Mock,
1096+
# env_vars: dict[str, str],
1097+
# expected_env_vars: dict[str, str],
1098+
# ):
1099+
# agent_engines.update(
1100+
# resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1101+
# description="foobar", # avoid "At least one of ... must be specified" errors.
1102+
# env_vars=env_vars,
1103+
# )
1104+
# update_agent_engine_mock.assert_called_once()
1105+
# deployment_spec = update_agent_engine_mock.call_args.kwargs[
1106+
# "request"
1107+
# ].reasoning_engine.spec.deployment_spec
1108+
# assert _utils.to_dict(deployment_spec)["env"] == [
1109+
# {"name": key, "value": value} for key, value in expected_env_vars.items()
1110+
# ]

‎tests/unit/vertexai/genai/test_agent_engines.py‎

Lines changed: 43 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
from google.cloud import aiplatform
3232
import vertexai
3333
from google.cloud.aiplatform import initializer
34-
from vertexai.agent_engines.templates import adk
3534
from vertexai._genai import _agent_engines_utils
3635
from vertexai._genai import agent_engines
3736
from vertexai._genai import types as _genai_types
@@ -852,48 +851,49 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare):
852851
"description": _TEST_AGENT_ENGINE_DESCRIPTION,
853852
}
854853

855-
@mock.patch.object(_agent_engines_utils, "_prepare")
856-
@pytest.mark.parametrize(
857-
"env_vars,expected_env_vars",
858-
[
859-
({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}),
860-
(None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}),
861-
(
862-
{"some_env": "some_val"},
863-
{
864-
"some_env": "some_val",
865-
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true",
866-
},
867-
),
868-
(
869-
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
870-
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
871-
),
872-
],
873-
)
874-
def test_agent_engine_adk_telemetry_enablement(
875-
self,
876-
mock_prepare: mock.Mock,
877-
env_vars: dict[str, str],
878-
expected_env_vars: dict[str, str],
879-
):
880-
agent = mock.Mock(spec=adk.AdkApp)
881-
agent.clone = lambda: agent
882-
agent.register_operations = lambda: {}
883-
884-
config = self.client.agent_engines._create_config(
885-
mode="create",
886-
agent=agent,
887-
staging_bucket=_TEST_STAGING_BUCKET,
888-
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
889-
description=_TEST_AGENT_ENGINE_DESCRIPTION,
890-
env_vars=env_vars,
891-
)
892-
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
893-
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
894-
assert config["spec"]["deployment_spec"]["env"] == [
895-
{"name": key, "value": value} for key, value in expected_env_vars.items()
896-
]
854+
# TODO(jawoszek): Uncomment once we're ready for default-on.
855+
# @mock.patch.object(_agent_engines_utils, "_prepare")
856+
# @pytest.mark.parametrize(
857+
# "env_vars,expected_env_vars",
858+
# [
859+
# ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}),
860+
# (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}),
861+
# (
862+
# {"some_env": "some_val"},
863+
# {
864+
# "some_env": "some_val",
865+
# GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false",
866+
# },
867+
# ),
868+
# (
869+
# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
870+
# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
871+
# ),
872+
# ],
873+
# )
874+
# def test_agent_engine_adk_telemetry_enablement(
875+
# self,
876+
# mock_prepare: mock.Mock,
877+
# env_vars: dict[str, str],
878+
# expected_env_vars: dict[str, str],
879+
# ):
880+
# agent = mock.Mock(spec=adk.AdkApp)
881+
# agent.clone = lambda: agent
882+
# agent.register_operations = lambda: {}
883+
884+
# config = self.client.agent_engines._create_config(
885+
# mode="create",
886+
# agent=agent,
887+
# staging_bucket=_TEST_STAGING_BUCKET,
888+
# display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
889+
# description=_TEST_AGENT_ENGINE_DESCRIPTION,
890+
# env_vars=env_vars,
891+
# )
892+
# assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
893+
# assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
894+
# assert config["spec"]["deployment_spec"]["env"] == [
895+
# {"name": key, "value": value} for key, value in expected_env_vars.items()
896+
# ]
897897

898898
@mock.patch.object(_agent_engines_utils, "_prepare")
899899
def test_create_agent_engine_config_full(self, mock_prepare):

‎vertexai/_genai/agent_engines.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1040,8 +1040,9 @@ def _create_config(
10401040
raise ValueError("location must be set using `vertexai.Client`.")
10411041
gcs_dir_name = gcs_dir_name or _agent_engines_utils._DEFAULT_GCS_DIR_NAME
10421042
agent = _agent_engines_utils._validate_agent_or_raise(agent=agent)
1043-
if _agent_engines_utils._is_adk_agent(agent):
1044-
env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)
1043+
# TODO(jawoszek): Uncomment once we're ready for default-on.
1044+
# if _agent_engines_utils._is_adk_agent(agent):
1045+
# env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)
10451046
staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise(
10461047
staging_bucket=staging_bucket,
10471048
)

‎vertexai/agent_engines/_agent_engines.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -518,8 +518,9 @@ def create(
518518
if agent_engine is not None:
519519
agent_engine = _validate_agent_engine_or_raise(agent_engine)
520520
staging_bucket = _validate_staging_bucket_or_raise(staging_bucket)
521-
if _is_adk_agent(None, agent_engine):
522-
env_vars = _add_telemetry_enablement_env(env_vars=env_vars)
521+
# TODO(jawoszek): Uncomment once we're ready for default-on.
522+
# if _is_adk_agent(None, agent_engine):
523+
# env_vars = _add_telemetry_enablement_env(env_vars=env_vars)
523524

524525
if agent_engine is None:
525526
if requirements is not None:

0 commit comments

Comments
 (0)