From e631a16d6451384f5b35552faf21427bdd072f6f Mon Sep 17 00:00:00 2001 From: Yifan Yang Date: Wed, 7 Oct 2026 16:57:00 -0700 Subject: [PATCH] fix: raise google.genai.errors.APIError when Agent Engine create or update LROs fail PiperOrigin-RevId: 995445165 --- agentplatform/_genai/_runtimes_utils.py | 18 +++++++++++++++ agentplatform/_genai/runtimes.py | 16 +++++++++---- .../agentplatform/genai/test_agent_engines.py | 23 +++++++++++++++---- vertexai/_genai/_agent_engines_utils.py | 18 +++++++++++++++ vertexai/_genai/agent_engines.py | 16 +++++++++---- 5 files changed, 79 insertions(+), 12 deletions(-) diff --git a/agentplatform/_genai/_runtimes_utils.py b/agentplatform/_genai/_runtimes_utils.py index b5f0b3d541..b037caebb8 100644 --- a/agentplatform/_genai/_runtimes_utils.py +++ b/agentplatform/_genai/_runtimes_utils.py @@ -52,6 +52,7 @@ import proto from google.api_core import exceptions +from google.genai import errors as google_genai_errors from google.genai import types as google_genai_types from google.api import httpbody_pb2 from google.protobuf import struct_pb2 @@ -564,6 +565,23 @@ def _await_operation( return operation +def _raise_for_operation_error( + *, + operation: Any, + message: str, +) -> None: + """Raises an APIError for a failed operation with a context note.""" + try: + google_genai_errors.APIError.raise_error( + operation.error.get("code"), operation.error, None + ) + except google_genai_errors.APIError as e: + # TODO(yifany): Drop guard once Python 3.10 support is removed. + if sys.version_info >= (3, 11): + e.add_note(message) + raise + + def _compare_requirements( *, requirements: Mapping[str, str], diff --git a/agentplatform/_genai/runtimes.py b/agentplatform/_genai/runtimes.py index ab235e3765..2cc5bbd676 100644 --- a/agentplatform/_genai/runtimes.py +++ b/agentplatform/_genai/runtimes.py @@ -1944,8 +1944,12 @@ def create( f"runtime=client.runtimes.get(name='{runtime.api_resource.name}')" ) elif operation.error: - raise RuntimeError( - f"Operation {operation.name} failed to create Agent Runtime: {operation.error}" + _runtimes_utils._raise_for_operation_error( + operation=operation, + message=( + f"Operation {operation.name} failed to create Agent Runtime:" + f" {operation.error}" + ), ) else: logger.warning("The operation returned an empty response.") @@ -2747,8 +2751,12 @@ def update( f"runtime=client.runtimes.get(name='{runtime.api_resource.name}')" ) elif operation.error: - raise RuntimeError( - f"Operation {operation.name} failed to update Agent Runtime: {operation.error}" + _runtimes_utils._raise_for_operation_error( + operation=operation, + message=( + f"Operation {operation.name} failed to update Agent Runtime:" + f" {operation.error}" + ), ) if runtime.api_resource.spec: self._register_api_methods(runtime=runtime) diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index 08f8c37e60..d8d0ab90e1 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -36,6 +36,7 @@ from agentplatform._genai import runtimes from agentplatform._genai import types as _genai_types from google.genai import client as genai_client +from google.genai import errors as genai_errors from google.genai import types as genai_types from google.api import httpbody_pb2 import pytest @@ -4154,7 +4155,7 @@ def test_create_runtime_error( self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - with pytest.raises(RuntimeError) as excinfo: + with pytest.raises(genai_errors.APIError) as excinfo: self.client.runtimes.create( agent=self.test_agent, config=_genai_types.AgentRuntimeConfig( @@ -4167,7 +4168,14 @@ def test_create_runtime_error( env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, ), ) - assert "failed to create Agent Runtime" in str(excinfo.value) + assert excinfo.value.code == _TEST_AGENT_ENGINE_ERROR["code"] + assert excinfo.value.message == _TEST_AGENT_ENGINE_ERROR["message"] + assert excinfo.value.details == _TEST_AGENT_ENGINE_ERROR + if sys.version_info >= (3, 11): + assert any( + "failed to create Agent Runtime" in note + for note in excinfo.value.__notes__ + ) @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_runtime_description(self, mock_await_operation): @@ -4178,14 +4186,21 @@ def test_update_runtime_description(self, mock_await_operation): self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - with pytest.raises(RuntimeError) as excinfo: + with pytest.raises(genai_errors.APIError) as excinfo: self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config=_genai_types.AgentRuntimeConfig( description=_TEST_AGENT_ENGINE_DESCRIPTION, ), ) - assert "failed to update Agent Runtime" in str(excinfo.value) + assert excinfo.value.code == _TEST_AGENT_ENGINE_ERROR["code"] + assert excinfo.value.message == _TEST_AGENT_ENGINE_ERROR["message"] + assert excinfo.value.details == _TEST_AGENT_ENGINE_ERROR + if sys.version_info >= (3, 11): + assert any( + "failed to update Agent Runtime" in note + for note in excinfo.value.__notes__ + ) @pytest.mark.parametrize( "test_case_name, test_operation_schemas, want_log_output", diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 52c9c58e3a..2da175b745 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -51,6 +51,7 @@ import proto from google.api_core import exceptions +from google.genai import errors as google_genai_errors from google.genai import types as google_genai_types from google.protobuf import struct_pb2 from google.protobuf import json_format @@ -550,6 +551,23 @@ def _await_operation( return operation +def _raise_for_operation_error( + *, + operation: Any, + message: str, +) -> None: + """Raises an APIError for a failed operation with a context note.""" + try: + google_genai_errors.APIError.raise_error( + operation.error.get("code"), operation.error, None + ) + except google_genai_errors.APIError as e: + # TODO(yifany): Drop guard once Python 3.10 support is removed. + if sys.version_info >= (3, 11): + e.add_note(message) + raise + + def _compare_requirements( *, requirements: Mapping[str, str], diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index e8aa8be81e..896e3ca95e 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -2039,8 +2039,12 @@ def create( f"agent_engine=client.agent_engines.get(name='{agent_engine.api_resource.name}')" ) elif operation.error: - raise RuntimeError( - f"Operation {operation.name} failed to create Agent Engine: {operation.error}" + _agent_engines_utils._raise_for_operation_error( + operation=operation, + message=( + f"Operation {operation.name} failed to create Agent Engine:" + f" {operation.error}" + ), ) else: logger.warning("The operation returned an empty response.") @@ -2848,8 +2852,12 @@ def update( f"agent_engine=client.agent_engines.get(name='{agent_engine.api_resource.name}')" ) elif operation.error: - raise RuntimeError( - f"Operation {operation.name} failed to update Agent Engine: {operation.error}" + _agent_engines_utils._raise_for_operation_error( + operation=operation, + message=( + f"Operation {operation.name} failed to update Agent Engine:" + f" {operation.error}" + ), ) if agent_engine.api_resource.spec: self._register_api_methods(agent_engine=agent_engine)