From 8d4ce38c3b02672e06458554c7fb3475d0d8c052 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Sun, 23 Nov 2025 00:00:17 -0800 Subject: [PATCH] fix: Gen AI SDK client - Fix bug in GCS bucket creation for new agent engines. PiperOrigin-RevId: 835792830 --- .../unit/vertexai/genai/test_agent_engines.py | 174 ++++++++++++++++++ vertexai/_genai/_agent_engines_utils.py | 10 +- vertexai/_genai/agent_engines.py | 1 + 3 files changed, 182 insertions(+), 3 deletions(-) diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index cc2b9ec813..a9560ebaa8 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -1450,6 +1450,77 @@ def test_create_base64_encoded_tarball_outside_project_dir_raises(self): finally: os.chdir(origin_dir) + @mock.patch.object(_agent_engines_utils, "_upload_requirements") + @mock.patch.object(_agent_engines_utils, "_upload_extra_packages") + @mock.patch.object(_agent_engines_utils, "_upload_agent_engine") + @mock.patch.object(_agent_engines_utils, "_scan_requirements") + @mock.patch.object(_agent_engines_utils, "_get_gcs_bucket") + def test_prepare_with_creds( + self, + mock_get_gcs_bucket, + mock_scan_requirements, + mock_upload_agent_engine, + mock_upload_extra_packages, + mock_upload_requirements, + ): + mock_scan_requirements.return_value = {} + mock_creds = mock.Mock(spec=auth_credentials.AnonymousCredentials()) + mock_creds.universe_domain = "googleapis.com" + _agent_engines_utils._prepare( + agent=self.test_agent, + project=_TEST_PROJECT, + location=_TEST_LOCATION, + staging_bucket=_TEST_STAGING_BUCKET, + credentials=mock_creds, + gcs_dir_name=_TEST_GCS_DIR_NAME, + requirements=[], + extra_packages=[], + ) + mock_upload_agent_engine.assert_called_once_with( + agent=self.test_agent, + gcs_bucket=mock.ANY, + gcs_dir_name=_TEST_GCS_DIR_NAME, + ) + + @mock.patch.object(_agent_engines_utils, "_upload_requirements") + @mock.patch.object(_agent_engines_utils, "_upload_extra_packages") + @mock.patch.object(_agent_engines_utils, "_upload_agent_engine") + @mock.patch.object(_agent_engines_utils, "_scan_requirements") + @mock.patch("google.auth.default") + @mock.patch.object(_agent_engines_utils, "_get_gcs_bucket") + def test_prepare_without_creds( + self, + mock_get_gcs_bucket, + mock_auth_default, + mock_scan_requirements, + mock_upload_agent_engine, + mock_upload_extra_packages, + mock_upload_requirements, + ): + mock_scan_requirements.return_value = {} + mock_creds = mock.Mock(spec=auth_credentials.AnonymousCredentials()) + mock_auth_default.return_value = (mock_creds, _TEST_PROJECT) + _agent_engines_utils._prepare( + agent=self.test_agent, + project=_TEST_PROJECT, + location=_TEST_LOCATION, + staging_bucket=_TEST_STAGING_BUCKET, + gcs_dir_name=_TEST_GCS_DIR_NAME, + requirements=[], + extra_packages=[], + ) + mock_get_gcs_bucket.assert_called_once_with( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + staging_bucket=_TEST_STAGING_BUCKET, + credentials=None, + ) + mock_upload_agent_engine.assert_called_once_with( + agent=self.test_agent, + gcs_bucket=mock.ANY, + gcs_dir_name=_TEST_GCS_DIR_NAME, + ) + @pytest.mark.usefixtures("google_auth_mock") class TestAgentEngine: @@ -2622,6 +2693,109 @@ def test_operation_schemas( want_operation_schemas.append(want_operation_schema) assert test_agent_engine.operation_schemas() == want_operation_schemas + @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(agent_engines.AgentEngines, "_create") + @mock.patch.object(_agent_engines_utils, "_await_operation") + def test_create_agent_engine_with_creds( + self, mock_await_operation, mock_create, mock_prepare + ): + mock_operation = mock.Mock() + mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME + mock_create.return_value = mock_operation + mock_await_operation.return_value = _genai_types.AgentEngineOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + self.client.agent_engines.create( + agent=self.test_agent, + config=_genai_types.AgentEngineConfig( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + staging_bucket=_TEST_STAGING_BUCKET, + ), + ) + mock_args, mock_kwargs = mock_prepare.call_args + assert mock_kwargs["agent"] == self.test_agent + assert mock_kwargs["extra_packages"] == [] + assert mock_kwargs["project"] == _TEST_PROJECT + assert mock_kwargs["location"] == _TEST_LOCATION + assert mock_kwargs["staging_bucket"] == _TEST_STAGING_BUCKET + assert mock_kwargs["credentials"] == _TEST_CREDENTIALS + assert mock_kwargs["gcs_dir_name"] == "agent_engine" + + @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(agent_engines.AgentEngines, "_create") + @mock.patch("google.auth.default") + @mock.patch.object(_agent_engines_utils, "_await_operation") + def test_create_agent_engine_without_creds( + self, mock_await_operation, mock_auth_default, mock_create, mock_prepare + ): + mock_operation = mock.Mock() + mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME + mock_create.return_value = mock_operation + mock_await_operation.return_value = _genai_types.AgentEngineOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + mock_creds = mock.Mock(spec=auth_credentials.AnonymousCredentials()) + mock_creds.quota_project_id = _TEST_PROJECT + mock_auth_default.return_value = (mock_creds, _TEST_PROJECT) + client = vertexai.Client( + project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=mock_creds + ) + client.agent_engines.create( + agent=self.test_agent, + config=_genai_types.AgentEngineConfig( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + staging_bucket=_TEST_STAGING_BUCKET, + ), + ) + mock_args, mock_kwargs = mock_prepare.call_args + assert mock_kwargs["agent"] == self.test_agent + assert mock_kwargs["extra_packages"] == [] + assert mock_kwargs["project"] == _TEST_PROJECT + assert mock_kwargs["location"] == _TEST_LOCATION + assert mock_kwargs["staging_bucket"] == _TEST_STAGING_BUCKET + assert mock_kwargs["credentials"] == mock_creds + assert mock_kwargs["gcs_dir_name"] == "agent_engine" + + @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(agent_engines.AgentEngines, "_create") + @mock.patch.object(_agent_engines_utils, "_await_operation") + def test_create_agent_engine_with_no_creds_in_client( + self, mock_await_operation, mock_create, mock_prepare + ): + mock_operation = mock.Mock() + mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME + mock_create.return_value = mock_operation + mock_await_operation.return_value = _genai_types.AgentEngineOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + client = vertexai.Client( + project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=None + ) + client.agent_engines.create( + agent=self.test_agent, + config=_genai_types.AgentEngineConfig( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + staging_bucket=_TEST_STAGING_BUCKET, + ), + ) + mock_args, mock_kwargs = mock_prepare.call_args + assert mock_kwargs["agent"] == self.test_agent + assert mock_kwargs["extra_packages"] == [] + assert mock_kwargs["project"] == _TEST_PROJECT + assert mock_kwargs["location"] == _TEST_LOCATION + assert mock_kwargs["staging_bucket"] == _TEST_STAGING_BUCKET + assert mock_kwargs["credentials"] is None + assert mock_kwargs["gcs_dir_name"] == "agent_engine" + @pytest.mark.usefixtures("google_auth_mock") class TestAgentEngineErrors: diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 0beb0e78af..987277ca47 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -772,10 +772,11 @@ def _get_gcs_bucket( project: str, location: str, staging_bucket: str, + credentials: Optional[Any] = None, ) -> _StorageBucket: """Gets or creates the GCS bucket.""" storage = _import_cloud_storage_or_raise() - storage_client = storage.Client(project=project) + storage_client = storage.Client(project=project, credentials=credentials) staging_bucket = staging_bucket.replace("gs://", "") try: gcs_bucket = storage_client.get_bucket(staging_bucket) @@ -910,6 +911,7 @@ def _prepare( location: str, staging_bucket: str, gcs_dir_name: str, + credentials: Optional[Any] = None, ) -> None: """Prepares the agent engine for creation or updates in Vertex AI. @@ -926,8 +928,9 @@ def _prepare( project (str): The project for the staging bucket. location (str): The location for the staging bucket. staging_bucket (str): The staging bucket name in the form "gs://...". - gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. + gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to use + for staging the artifacts needed. + credentials: The credentials to use for the storage client. """ if agent is None: return @@ -935,6 +938,7 @@ def _prepare( project=project, location=location, staging_bucket=staging_bucket, + credentials=credentials, ) _upload_agent_engine( agent=agent, diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 9638ff6d82..6cfcb8be62 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -1082,6 +1082,7 @@ def _create_config( staging_bucket=staging_bucket, gcs_dir_name=gcs_dir_name, extra_packages=extra_packages, + credentials=self._api_client._credentials, ) # Update the package spec. update_masks.append("spec.package_spec.pickle_object_gcs_uri")