Skip to content

Commit f2367d6

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: add telemetry enablement env for agent engines deployed using module
This is to fix telemetry enablement in `adk deploy` CLI. FUTURE_COPYBARA_INTEGRATE_REVIEW=#6093 from googleapis:release-please--branches--main cb969c2 PiperOrigin-RevId: 828955929
1 parent 25873a6 commit f2367d6

File tree

2 files changed

+51
-0
lines changed

2 files changed

+51
-0
lines changed

tests/unit/vertexai/genai/test_agent_engines.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -903,6 +903,54 @@ def test_agent_engine_adk_telemetry_enablement(
903903
{"name": key, "value": value} for key, value in expected_env_vars.items()
904904
]
905905

906+
907+
@mock.patch.object(_agent_engines_utils, "_prepare")
908+
@pytest.mark.parametrize(
909+
"env_vars,expected_env_vars",
910+
[
911+
({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}),
912+
(None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}),
913+
(
914+
{"some_env": "some_val"},
915+
{
916+
"some_env": "some_val",
917+
GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified",
918+
},
919+
),
920+
(
921+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
922+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"},
923+
),
924+
(
925+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
926+
{GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"},
927+
),
928+
],
929+
)
930+
def test_agent_engine_adk_telemetry_enablement_through_source_packages(
931+
self,
932+
mock_prepare: mock.Mock,
933+
env_vars: dict[str, str],
934+
expected_env_vars: dict[str, str],
935+
):
936+
config = self.client.agent_engines._create_config(
937+
mode="create",
938+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
939+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
940+
source_packages=[],
941+
class_methods=[],
942+
entrypoint_module=".",
943+
entrypoint_object=".",
944+
env_vars=env_vars,
945+
agent_framework="google-adk",
946+
)
947+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
948+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
949+
assert config["spec"]["deployment_spec"]["env"] == [
950+
{"name": key, "value": value} for key, value in expected_env_vars.items()
951+
]
952+
953+
906954
@mock.patch.object(_agent_engines_utils, "_prepare")
907955
def test_create_agent_engine_config_full(self, mock_prepare):
908956
config = self.client.agent_engines._create_config(

vertexai/_genai/agent_engines.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1024,6 +1024,9 @@ def _create_config(
10241024
update_masks.append("labels")
10251025
config["labels"] = labels
10261026

1027+
if agent_framework == "google-adk":
1028+
env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)
1029+
10271030
sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
10281031
agent_engine_spec = None
10291032
if agent is not None:

0 commit comments

Comments
 (0)