fix: raise google.genai.errors.APIError when Agent Engine create or update LROs fail

PiperOrigin-RevId: 995250809

Author: yifanyang

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],

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)

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",

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],

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)