diff --git a/agentplatform/agent_engines/templates/adk.py b/agentplatform/agent_engines/templates/adk.py index ce9c9bb967..9eaa49ce46 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/agent_engines/templates/adk.py @@ -2172,7 +2172,8 @@ def _tracing_enabled(self) -> bool: ) def project_id(self) -> Optional[str]: - if project := self._tmpl_attrs.get("project"): + project = self._tmpl_attrs.get("project") + if project and str(project).isdigit(): try: from google.cloud.aiplatform.utils import ( resource_manager_utils, @@ -2184,4 +2185,4 @@ def project_id(self) -> Optional[str]: except (exceptions.PermissionDenied, exceptions.Unauthenticated): return project - return None + return project or None diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index 99e9881cc4..a3725842d9 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py @@ -60,6 +60,7 @@ def __init__(self, name: str, model: str): _TEST_LOCATION = "us-central1" _TEST_PROJECT = "test-project" _TEST_PROJECT_ID = "test-project-id" +_TEST_PROJECT_NUMBER = "123456789" _TEST_API_KEY = "test-api-key" _TEST_MODEL = "gemini-2.0-flash" _TEST_USER_ID = "test_user_id" @@ -1251,7 +1252,7 @@ def test_default_instrumentor_enablement( # Assert default_instrumentor_builder_mock.assert_called_once_with( - _TEST_PROJECT_ID, + _TEST_PROJECT, enable_tracing=want_tracing_setup, enable_logging=want_logging_setup, ) @@ -1314,7 +1315,7 @@ def test_custom_instrumentor_enablement( # Assert if want_custom_instrumentor_called: - custom_instrumentor.assert_called_once_with(_TEST_PROJECT_ID) + custom_instrumentor.assert_called_once_with(_TEST_PROJECT) else: custom_instrumentor.assert_not_called() @@ -1347,7 +1348,7 @@ def test_tracing_setup( headers=mock.ANY, ) - get_project_id_mock.assert_called_with(_TEST_PROJECT) + get_project_id_mock.assert_not_called() user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"] assert ( @@ -1358,6 +1359,38 @@ def test_tracing_setup( is not None ) + def test_project_id_skips_lookup_for_project_id( + self, + get_project_id_mock: mock.Mock, + ): + """Project IDs are returned as-is without a GetProject call.""" + app = adk_template.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT + + assert app.project_id() == _TEST_PROJECT + get_project_id_mock.assert_not_called() + + def test_project_id_resolves_project_number( + self, + get_project_id_mock: mock.Mock, + ): + """Project numbers are resolved to project IDs via GetProject.""" + app = adk_template.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER + + assert app.project_id() == _TEST_PROJECT_ID + get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER) + + def test_project_id_without_project( + self, + get_project_id_mock: mock.Mock, + ): + app = adk_template.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = None + + assert app.project_id() is None + get_project_id_mock.assert_not_called() + @pytest.mark.usefixtures("caplog") def test_enable_tracing( self, 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 a359289406..98078145dc 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -59,6 +59,7 @@ def __init__(self, name: str, model: str): _TEST_LOCATION = "us-central1" _TEST_PROJECT = "test-project" _TEST_PROJECT_ID = "test-project-id" +_TEST_PROJECT_NUMBER = "123456789" _TEST_API_KEY = "test-api-key" _TEST_MODEL = "gemini-2.0-flash" _TEST_USER_ID = "test_user_id" @@ -1064,7 +1065,7 @@ def test_tracing_setup( headers=mock.ANY, ) - get_project_id_mock.assert_called_with(_TEST_PROJECT) + get_project_id_mock.assert_not_called() user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"] assert ( @@ -1075,6 +1076,44 @@ def test_tracing_setup( is not None ) + def test_project_id_skips_lookup_for_project_id( + self, + get_project_id_mock: mock.Mock, + ): + """Project IDs are returned as-is without a GetProject call.""" + app = agent_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + + assert app.project_id() == _TEST_PROJECT + get_project_id_mock.assert_not_called() + + def test_project_id_resolves_project_number( + self, + get_project_id_mock: mock.Mock, + ): + """Project numbers are resolved to project IDs via GetProject.""" + app = agent_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + + assert app.project_id() == _TEST_PROJECT_ID + get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER) + + def test_project_id_without_project( + self, + get_project_id_mock: mock.Mock, + ): + app = agent_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = None + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + + assert app.project_id() is None + get_project_id_mock.assert_not_called() + @pytest.mark.usefixtures("caplog") def test_enable_tracing( self, diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index e956f06e1a..391dc7298a 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py @@ -47,6 +47,7 @@ def __init__(self, name: str, model: str): _TEST_LOCATION = "us-central1" _TEST_PROJECT = "test-project" _TEST_PROJECT_ID = "test-project-id" +_TEST_PROJECT_NUMBER = "123456789" _TEST_MODEL = "gemini-2.0-flash" _TEST_USER_ID = "test_user_id" _TEST_AGENT_NAME = "test_agent" @@ -1010,11 +1011,9 @@ def test_tracing_setup( headers=mock.ANY, ) - calls = [ - mock.call(project_number=_TEST_PROJECT_ID, credentials=mock.ANY), - mock.call(_TEST_PROJECT_ID), - ] - get_project_id_mock.assert_has_calls(calls) + get_project_id_mock.assert_called_once_with( + project_number=_TEST_PROJECT_ID, credentials=mock.ANY + ) user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"] assert ( @@ -1029,6 +1028,44 @@ def test_tracing_setup( == expected_attributes ) + def test_project_id_skips_lookup_for_project_id( + self, + get_project_id_mock: mock.Mock, + ): + """Project IDs are returned as-is without a GetProject call.""" + app = reasoning_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + + assert app.project_id() == _TEST_PROJECT + get_project_id_mock.assert_not_called() + + def test_project_id_resolves_project_number( + self, + get_project_id_mock: mock.Mock, + ): + """Project numbers are resolved to project IDs via GetProject.""" + app = reasoning_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + + assert app.project_id() == _TEST_PROJECT_ID + get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER) + + def test_project_id_without_project( + self, + get_project_id_mock: mock.Mock, + ): + app = reasoning_engines.AdkApp(agent=_TEST_AGENT) + app._tmpl_attrs["project"] = None + # Discard the lookup that vertexai.init() made in setup_method. + get_project_id_mock.reset_mock() + + assert app.project_id() is None + get_project_id_mock.assert_not_called() + @mock.patch.dict(os.environ) def test_span_content_capture_disabled_by_default(self): app = reasoning_engines.AdkApp(agent=_TEST_AGENT) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 76398ba03c..5f814b5b12 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -1868,7 +1868,8 @@ def _tracing_enabled(self) -> bool: ) def project_id(self) -> Optional[str]: - if project := self._tmpl_attrs.get("project"): + project = self._tmpl_attrs.get("project") + if project and str(project).isdigit(): try: from google.cloud.aiplatform.utils import ( resource_manager_utils, @@ -1880,4 +1881,4 @@ def project_id(self) -> Optional[str]: except (exceptions.PermissionDenied, exceptions.Unauthenticated): return project - return None + return project or None diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index ed8909d804..003812ab3a 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -1765,7 +1765,8 @@ def _warn_if_telemetry_api_disabled(self): _warn(_TELEMETRY_API_DISABLED_WARNING % (project, project)) def project_id(self) -> Optional[str]: - if project := self._tmpl_attrs.get("project"): + project = self._tmpl_attrs.get("project") + if project and str(project).isdigit(): try: from google.cloud.aiplatform.utils import ( resource_manager_utils, @@ -1777,4 +1778,4 @@ def project_id(self) -> Optional[str]: except (exceptions.PermissionDenied, exceptions.Unauthenticated): return project - return None + return project or None