Skip to content

Commit 6f9f531

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 bf1851e commit 6f9f531

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
@@ -856,48 +855,49 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare):
856855
"description": _TEST_AGENT_ENGINE_DESCRIPTION,
857856
}
858857

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

902902
@mock.patch.object(_agent_engines_utils, "_prepare")
903903
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
@@ -1042,8 +1042,9 @@ def _create_config(
10421042
raise ValueError("location must be set using `vertexai.Client`.")
10431043
gcs_dir_name = gcs_dir_name or _agent_engines_utils._DEFAULT_GCS_DIR_NAME
10441044
agent = _agent_engines_utils._validate_agent_or_raise(agent=agent)
1045-
if _agent_engines_utils._is_adk_agent(agent):
1046-
env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)
1045+
# TODO(jawoszek): Uncomment once we're ready for default-on.
1046+
# if _agent_engines_utils._is_adk_agent(agent):
1047+
# env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)
10471048
staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise(
10481049
staging_bucket=staging_bucket,
10491050
)

‎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)