Skip to content

Commit 468e341

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 468e341

File tree

2 files changed

+49
-0
lines changed

2 files changed

+49
-0
lines changed

tests/unit/vertexai/genai/test_agent_engines.py

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

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