Skip to content

Commit 2ada2c5

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: default to "unspecified" for telemetry enablement
PiperOrigin-RevId: 828517813
1 parent 531d223 commit 2ada2c5

File tree

7 files changed

+190
-162
lines changed

7 files changed

+190
-162
lines changed

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 117 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from google.auth import credentials as auth_credentials
2727
from google.cloud import storage
2828
import vertexai
29+
from google.cloud import aiplatform
2930
from google.cloud.aiplatform_v1 import types as aip_types
3031
from google.cloud.aiplatform_v1.services import reasoning_engine_service
3132
from google.cloud.aiplatform import base
@@ -637,6 +638,9 @@ async def test_async_search_memory(self):
637638
("1.16.0", None, False, False, False),
638639
("1.16.0", None, True, False, True),
639640
("1.16.0", None, None, False, False),
641+
("1.16.0", None, "unspecified", False, False),
642+
("1.16.0", False, "unspecified", False, False),
643+
("1.16.0", True, "unspecified", True, False),
640644
("1.17.0", False, False, False, False),
641645
("1.17.0", False, True, False, True),
642646
("1.17.0", False, None, False, False),
@@ -646,6 +650,9 @@ async def test_async_search_memory(self):
646650
("1.17.0", None, False, False, False),
647651
("1.17.0", None, True, True, True),
648652
("1.17.0", None, None, False, False),
653+
("1.17.0", None, "unspecified", False, False),
654+
("1.17.0", False, "unspecified", False, False),
655+
("1.17.0", True, "unspecified", True, False),
649656
],
650657
)
651658
@mock.patch.dict(os.environ)
@@ -690,6 +697,9 @@ def test_default_instrumentor_enablement(
690697
("1.16.0", None, False, False),
691698
("1.16.0", None, True, False),
692699
("1.16.0", None, None, False),
700+
("1.16.0", None, "unspecified", False),
701+
("1.16.0", False, "unspecified", False),
702+
("1.16.0", True, "unspecified", True),
693703
("1.17.0", False, False, False),
694704
("1.17.0", False, True, False),
695705
("1.17.0", False, None, False),
@@ -699,6 +709,9 @@ def test_default_instrumentor_enablement(
699709
("1.17.0", None, False, False),
700710
("1.17.0", None, True, True),
701711
("1.17.0", None, None, False),
712+
("1.17.0", None, "unspecified", False),
713+
("1.17.0", False, "unspecified", False),
714+
("1.17.0", True, "unspecified", True),
702715
],
703716
)
704717
@mock.patch.dict(os.environ)
@@ -1011,100 +1024,107 @@ def update_agent_engine_mock():
10111024
yield update_agent_engine_mock
10121025

10131026

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-
# ]
1027+
@pytest.mark.usefixtures("google_auth_mock")
1028+
class TestAgentEngines:
1029+
def setup_method(self):
1030+
importlib.reload(initializer)
1031+
importlib.reload(aiplatform)
1032+
aiplatform.init(
1033+
project=_TEST_PROJECT,
1034+
location=_TEST_LOCATION,
1035+
credentials=_TEST_CREDENTIALS,
1036+
staging_bucket=_TEST_STAGING_BUCKET,
1037+
)
1038+
1039+
def teardown_method(self):
1040+
initializer.global_pool.shutdown(wait=True)
1041+
1042+
@pytest.mark.parametrize(
1043+
"env_vars,expected_env_vars",
1044+
[
1045+
({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}),
1046+
(None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}),
1047+
(
1048+
{"some_env": "some_val"},
1049+
{
1050+
"some_env": "some_val",
1051+
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified",
1052+
},
1053+
),
1054+
(
1055+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1056+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1057+
),
1058+
(
1059+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1060+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1061+
),
1062+
],
1063+
)
1064+
def test_create_default_telemetry_enablement(
1065+
self,
1066+
create_agent_engine_mock: mock.Mock,
1067+
cloud_storage_create_bucket_mock: mock.Mock,
1068+
cloudpickle_dump_mock: mock.Mock,
1069+
cloudpickle_load_mock: mock.Mock,
1070+
get_gca_resource_mock: mock.Mock,
1071+
env_vars: dict[str, str],
1072+
expected_env_vars: dict[str, str],
1073+
):
1074+
agent_engines.create(
1075+
agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT),
1076+
env_vars=env_vars,
1077+
)
1078+
create_agent_engine_mock.assert_called_once()
1079+
deployment_spec = create_agent_engine_mock.call_args.kwargs[
1080+
"reasoning_engine"
1081+
].spec.deployment_spec
1082+
assert _utils.to_dict(deployment_spec)["env"] == [
1083+
{"name": key, "value": value} for key, value in expected_env_vars.items()
1084+
]
1085+
1086+
@pytest.mark.parametrize(
1087+
"env_vars,expected_env_vars",
1088+
[
1089+
({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}),
1090+
(None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}),
1091+
(
1092+
{"some_env": "some_val"},
1093+
{
1094+
"some_env": "some_val",
1095+
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified",
1096+
},
1097+
),
1098+
(
1099+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1100+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
1101+
),
1102+
(
1103+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1104+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
1105+
),
1106+
],
1107+
)
1108+
def test_update_default_telemetry_enablement(
1109+
self,
1110+
update_agent_engine_mock: mock.Mock,
1111+
cloud_storage_create_bucket_mock: mock.Mock,
1112+
cloudpickle_dump_mock: mock.Mock,
1113+
cloudpickle_load_mock: mock.Mock,
1114+
get_gca_resource_mock: mock.Mock,
1115+
get_agent_engine_mock: mock.Mock,
1116+
env_vars: dict[str, str],
1117+
expected_env_vars: dict[str, str],
1118+
):
1119+
agent_engines.update(
1120+
resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1121+
description="foobar", # avoid "At least one of ... must be specified" errors.
1122+
env_vars=env_vars,
1123+
)
1124+
update_agent_engine_mock.assert_called_once()
1125+
deployment_spec = update_agent_engine_mock.call_args.kwargs[
1126+
"request"
1127+
].reasoning_engine.spec.deployment_spec
1128+
assert _utils.to_dict(deployment_spec)["env"] == [
1129+
{"name": key, "value": value} for key, value in expected_env_vars.items()
1130+
]

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

Lines changed: 47 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from google.cloud import aiplatform
3232
import vertexai
3333
from google.cloud.aiplatform import initializer
34+
from vertexai.agent_engines.templates import adk
3435
from vertexai._genai import _agent_engines_utils
3536
from vertexai._genai import agent_engines
3637
from vertexai._genai import types as _genai_types
@@ -855,49 +856,52 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare):
855856
"description": _TEST_AGENT_ENGINE_DESCRIPTION,
856857
}
857858

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-
# ]
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: "unspecified"}),
864+
(None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}),
865+
(
866+
{"some_env": "some_val"},
867+
{
868+
"some_env": "some_val",
869+
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified",
870+
},
871+
),
872+
(
873+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
874+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
875+
),
876+
(
877+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
878+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
879+
),
880+
],
881+
)
882+
def test_agent_engine_adk_telemetry_enablement(
883+
self,
884+
mock_prepare: mock.Mock,
885+
env_vars: dict[str, str],
886+
expected_env_vars: dict[str, str],
887+
):
888+
agent = mock.Mock(spec=adk.AdkApp)
889+
agent.clone = lambda: agent
890+
agent.register_operations = lambda: {}
891+
892+
config = self.client.agent_engines._create_config(
893+
mode="create",
894+
agent=agent,
895+
staging_bucket=_TEST_STAGING_BUCKET,
896+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
897+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
898+
env_vars=env_vars,
899+
)
900+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
901+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
902+
assert config["spec"]["deployment_spec"]["env"] == [
903+
{"name": key, "value": value} for key, value in expected_env_vars.items()
904+
]
901905

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

‎vertexai/_genai/_agent_engines_utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1920,7 +1920,7 @@ def _add_telemetry_enablement_env(
19201920
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = (
19211921
"GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY"
19221922
)
1923-
env_to_add = {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}
1923+
env_to_add = {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}
19241924

19251925
if env_vars is None:
19261926
return env_to_add

‎vertexai/_genai/agent_engines.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1042,9 +1042,8 @@ 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-
# 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)
1045+
if _agent_engines_utils._is_adk_agent(agent):
1046+
env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)
10481047
staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise(
10491048
staging_bucket=staging_bucket,
10501049
)

0 commit comments

Comments
 (0)