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 b4f41aeb3d..ee3a0677f5 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -26,7 +26,6 @@ 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 @@ -1012,99 +1011,100 @@ def update_agent_engine_mock(): yield update_agent_engine_mock -@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: "true"}), - (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - ( - {"some_env": "some_val"}, - { - "some_env": "some_val", - 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: "true"}), - (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - ( - {"some_env": "some_val"}, - { - "some_env": "some_val", - 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() - ] +# 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() +# ] diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index ade08e7b2f..b8971b8f92 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -31,7 +31,6 @@ 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 @@ -856,48 +855,49 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare): "description": _TEST_AGENT_ENGINE_DESCRIPTION, } - @mock.patch.object(_agent_engines_utils, "_prepare") - @pytest.mark.parametrize( - "env_vars,expected_env_vars", - [ - ({}, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - (None, {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "true"}), - ( - {"some_env": "some_val"}, - { - "some_env": "some_val", - 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() - ] + # 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") def test_create_agent_engine_config_full(self, mock_prepare): diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 9b0a882de2..b0a5fa737f 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -1042,8 +1042,9 @@ 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) - if _agent_engines_utils._is_adk_agent(agent): - env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) + # 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) 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 596056b15a..3696fd4a33 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -518,8 +518,9 @@ 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) - if _is_adk_agent(None, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) + # 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 agent_engine is None: if requirements is not None: