diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index ee3a0677f5..e7e843e05c 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -26,6 +26,7 @@ from google.auth import credentials as auth_credentials from google.cloud import storage import vertexai +from google.cloud import aiplatform from google.cloud.aiplatform_v1 import types as aip_types from google.cloud.aiplatform_v1.services import reasoning_engine_service from google.cloud.aiplatform import base @@ -637,6 +638,9 @@ async def test_async_search_memory(self): ("1.16.0", None, False, False, False), ("1.16.0", None, True, False, True), ("1.16.0", None, None, False, False), + ("1.16.0", None, "unspecified", False, False), + ("1.16.0", False, "unspecified", False, False), + ("1.16.0", True, "unspecified", True, False), ("1.17.0", False, False, False, False), ("1.17.0", False, True, False, True), ("1.17.0", False, None, False, False), @@ -646,6 +650,9 @@ async def test_async_search_memory(self): ("1.17.0", None, False, False, False), ("1.17.0", None, True, True, True), ("1.17.0", None, None, False, False), + ("1.17.0", None, "unspecified", False, False), + ("1.17.0", False, "unspecified", False, False), + ("1.17.0", True, "unspecified", True, False), ], ) @mock.patch.dict(os.environ) @@ -690,6 +697,9 @@ def test_default_instrumentor_enablement( ("1.16.0", None, False, False), ("1.16.0", None, True, False), ("1.16.0", None, None, False), + ("1.16.0", None, "unspecified", False), + ("1.16.0", False, "unspecified", False), + ("1.16.0", True, "unspecified", True), ("1.17.0", False, False, False), ("1.17.0", False, True, False), ("1.17.0", False, None, False), @@ -699,6 +709,9 @@ def test_default_instrumentor_enablement( ("1.17.0", None, False, False), ("1.17.0", None, True, True), ("1.17.0", None, None, False), + ("1.17.0", None, "unspecified", False), + ("1.17.0", False, "unspecified", False), + ("1.17.0", True, "unspecified", True), ], ) @mock.patch.dict(os.environ) @@ -1011,100 +1024,107 @@ def update_agent_engine_mock(): yield update_agent_engine_mock -# TODO(jawoszek): Uncomment once we're ready for default-on. -# @pytest.mark.usefixtures("google_auth_mock") -# class TestAgentEngines: -# def setup_method(self): -# importlib.reload(initializer) -# importlib.reload(aiplatform) -# aiplatform.init( -# project=_TEST_PROJECT, -# location=_TEST_LOCATION, -# credentials=_TEST_CREDENTIALS, -# staging_bucket=_TEST_STAGING_BUCKET, -# ) - -# def teardown_method(self): -# initializer.global_pool.shutdown(wait=True) - -# @pytest.mark.parametrize( -# "env_vars,expected_env_vars", -# [ -# ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), -# (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), -# ( -# {"some_env": "some_val"}, -# { -# "some_env": "some_val", -# GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false", -# }, -# ), -# ( -# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, -# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, -# ), -# ], -# ) -# def test_create_default_telemetry_enablement( -# self, -# create_agent_engine_mock: mock.Mock, -# cloud_storage_create_bucket_mock: mock.Mock, -# cloudpickle_dump_mock: mock.Mock, -# cloudpickle_load_mock: mock.Mock, -# get_gca_resource_mock: mock.Mock, -# env_vars: dict[str, str], -# expected_env_vars: dict[str, str], -# ): -# agent_engines.create( -# agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT), -# env_vars=env_vars, -# ) -# create_agent_engine_mock.assert_called_once() -# deployment_spec = create_agent_engine_mock.call_args.kwargs[ -# "reasoning_engine" -# ].spec.deployment_spec -# assert _utils.to_dict(deployment_spec)["env"] == [ -# {"name": key, "value": value} for key, value in expected_env_vars.items() -# ] - -# @pytest.mark.parametrize( -# "env_vars,expected_env_vars", -# [ -# ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), -# (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), -# ( -# {"some_env": "some_val"}, -# { -# "some_env": "some_val", -# GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false", -# }, -# ), -# ( -# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, -# {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, -# ), -# ], -# ) -# def test_update_default_telemetry_enablement( -# self, -# update_agent_engine_mock: mock.Mock, -# cloud_storage_create_bucket_mock: mock.Mock, -# cloudpickle_dump_mock: mock.Mock, -# cloudpickle_load_mock: mock.Mock, -# get_gca_resource_mock: mock.Mock, -# get_agent_engine_mock: mock.Mock, -# env_vars: dict[str, str], -# expected_env_vars: dict[str, str], -# ): -# agent_engines.update( -# resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME, -# description="foobar", # avoid "At least one of ... must be specified" errors. -# env_vars=env_vars, -# ) -# update_agent_engine_mock.assert_called_once() -# deployment_spec = update_agent_engine_mock.call_args.kwargs[ -# "request" -# ].reasoning_engine.spec.deployment_spec -# assert _utils.to_dict(deployment_spec)["env"] == [ -# {"name": key, "value": value} for key, value in expected_env_vars.items() -# ] +@pytest.mark.usefixtures("google_auth_mock") +class TestAgentEngines: + def setup_method(self): + importlib.reload(initializer) + importlib.reload(aiplatform) + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + staging_bucket=_TEST_STAGING_BUCKET, + ) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.parametrize( + "env_vars,expected_env_vars", + [ + ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}), + (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}), + ( + {"some_env": "some_val"}, + { + "some_env": "some_val", + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified", + }, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + ), + ], + ) + def test_create_default_telemetry_enablement( + self, + create_agent_engine_mock: mock.Mock, + cloud_storage_create_bucket_mock: mock.Mock, + cloudpickle_dump_mock: mock.Mock, + cloudpickle_load_mock: mock.Mock, + get_gca_resource_mock: mock.Mock, + env_vars: dict[str, str], + expected_env_vars: dict[str, str], + ): + agent_engines.create( + agent_engine=agent_engines.AdkApp(agent=_TEST_AGENT), + env_vars=env_vars, + ) + create_agent_engine_mock.assert_called_once() + deployment_spec = create_agent_engine_mock.call_args.kwargs[ + "reasoning_engine" + ].spec.deployment_spec + assert _utils.to_dict(deployment_spec)["env"] == [ + {"name": key, "value": value} for key, value in expected_env_vars.items() + ] + + @pytest.mark.parametrize( + "env_vars,expected_env_vars", + [ + ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}), + (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}), + ( + {"some_env": "some_val"}, + { + "some_env": "some_val", + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified", + }, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + ), + ], + ) + def test_update_default_telemetry_enablement( + self, + update_agent_engine_mock: mock.Mock, + cloud_storage_create_bucket_mock: mock.Mock, + cloudpickle_dump_mock: mock.Mock, + cloudpickle_load_mock: mock.Mock, + get_gca_resource_mock: mock.Mock, + get_agent_engine_mock: mock.Mock, + env_vars: dict[str, str], + expected_env_vars: dict[str, str], + ): + agent_engines.update( + resource_name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + description="foobar", # avoid "At least one of ... must be specified" errors. + env_vars=env_vars, + ) + update_agent_engine_mock.assert_called_once() + deployment_spec = update_agent_engine_mock.call_args.kwargs[ + "request" + ].reasoning_engine.spec.deployment_spec + assert _utils.to_dict(deployment_spec)["env"] == [ + {"name": key, "value": value} for key, value in expected_env_vars.items() + ] diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index b8971b8f92..ca97917edf 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -31,6 +31,7 @@ from google.cloud import aiplatform import vertexai from google.cloud.aiplatform import initializer +from vertexai.agent_engines.templates import adk from vertexai._genai import _agent_engines_utils from vertexai._genai import agent_engines from vertexai._genai import types as _genai_types @@ -855,49 +856,52 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare): "description": _TEST_AGENT_ENGINE_DESCRIPTION, } - # TODO(jawoszek): Uncomment once we're ready for default-on. - # @mock.patch.object(_agent_engines_utils, "_prepare") - # @pytest.mark.parametrize( - # "env_vars,expected_env_vars", - # [ - # ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), - # (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}), - # ( - # {"some_env": "some_val"}, - # { - # "some_env": "some_val", - # GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false", - # }, - # ), - # ( - # {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, - # {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, - # ), - # ], - # ) - # def test_agent_engine_adk_telemetry_enablement( - # self, - # mock_prepare: mock.Mock, - # env_vars: dict[str, str], - # expected_env_vars: dict[str, str], - # ): - # agent = mock.Mock(spec=adk.AdkApp) - # agent.clone = lambda: agent - # agent.register_operations = lambda: {} - - # config = self.client.agent_engines._create_config( - # mode="create", - # agent=agent, - # staging_bucket=_TEST_STAGING_BUCKET, - # display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - # description=_TEST_AGENT_ENGINE_DESCRIPTION, - # env_vars=env_vars, - # ) - # assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME - # assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION - # assert config["spec"]["deployment_spec"]["env"] == [ - # {"name": key, "value": value} for key, value in expected_env_vars.items() - # ] + @mock.patch.object(_agent_engines_utils, "_prepare") + @pytest.mark.parametrize( + "env_vars,expected_env_vars", + [ + ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}), + (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}), + ( + {"some_env": "some_val"}, + { + "some_env": "some_val", + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified", + }, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}, + ), + ( + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "false"}, + ), + ], + ) + def test_agent_engine_adk_telemetry_enablement( + self, + mock_prepare: mock.Mock, + env_vars: dict[str, str], + expected_env_vars: dict[str, str], + ): + agent = mock.Mock(spec=adk.AdkApp) + agent.clone = lambda: agent + agent.register_operations = lambda: {} + + config = self.client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + description=_TEST_AGENT_ENGINE_DESCRIPTION, + env_vars=env_vars, + ) + assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME + assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION + assert config["spec"]["deployment_spec"]["env"] == [ + {"name": key, "value": value} for key, value in expected_env_vars.items() + ] @mock.patch.object(_agent_engines_utils, "_prepare") def test_create_agent_engine_config_full(self, mock_prepare): diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 0e063a2c1b..69945a8cd3 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -1920,7 +1920,7 @@ def _add_telemetry_enablement_env( GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" ) - env_to_add = {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"} + env_to_add = {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} if env_vars is None: return env_to_add diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index b0a5fa737f..9b0a882de2 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -1042,9 +1042,8 @@ def _create_config( raise ValueError("location must be set using `vertexai.Client`.") gcs_dir_name = gcs_dir_name or _agent_engines_utils._DEFAULT_GCS_DIR_NAME agent = _agent_engines_utils._validate_agent_or_raise(agent=agent) - # TODO(jawoszek): Uncomment once we're ready for default-on. - # if _agent_engines_utils._is_adk_agent(agent): - # env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) + if _agent_engines_utils._is_adk_agent(agent): + env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise( staging_bucket=staging_bucket, ) diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index 3696fd4a33..b6cad5cdd5 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -518,9 +518,8 @@ def create( if agent_engine is not None: agent_engine = _validate_agent_engine_or_raise(agent_engine) staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - # TODO(jawoszek): Uncomment once we're ready for default-on. - # if _is_adk_agent(None, agent_engine): - # env_vars = _add_telemetry_enablement_env(env_vars=env_vars) + if _is_adk_agent(None, agent_engine): + env_vars = _add_telemetry_enablement_env(env_vars=env_vars) if agent_engine is None: if requirements is not None: @@ -1115,16 +1114,16 @@ def _add_telemetry_enablement_env(*, env_vars: EnvVars) -> EnvVars: ) if env_vars is None: - return {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"} + return {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} if isinstance(env_vars, dict): return ( env_vars if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars - else env_vars | {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"} + else env_vars | {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} ) if isinstance(env_vars, list) or isinstance(env_vars, tuple): if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY not in os.environ: - os.environ[GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] = "true" + os.environ[GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] = "unspecified" if isinstance(env_vars, list): return env_vars + [GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index feea1a6715..817e8475fc 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -1539,12 +1539,15 @@ def _telemetry_enabled(self) -> Optional[bool]: "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" ) - return ( - os.getenv(GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY, "0").lower() - in ("true", "1") - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in os.environ - else None - ) + env_value = os.getenv( + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY, "unspecified" + ).lower() + + if env_value in ("true", "1"): + return True + if env_value in ("false", "0"): + return False + return None # Tracing enablement follows truth table: def _tracing_enabled(self) -> bool: diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 07c3ee44c3..a14d314b68 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -1475,12 +1475,15 @@ def _telemetry_enabled(self) -> Optional[bool]: "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" ) - return ( - os.getenv(GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY, "0").lower() - in ("true", "1") - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in os.environ - else None - ) + env_value = os.getenv( + GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY, "unspecified" + ).lower() + + if env_value in ("true", "1"): + return True + if env_value in ("false", "0"): + return False + return None # Tracing enablement follows truth table: def _tracing_enabled(self) -> bool: