fix: GenAI SDK client - Provide the operation error message if it is available

PiperOrigin-RevId: 799667737

Author: yeesian

tests/unit/vertexai/genai/replays/test_list_agent_engine_memories.py

@@ -35,6 +35,8 @@ def test_list_memories(client):
     )
     assert len(memory_list) == 1
     assert isinstance(memory_list[0], types.Memory)
+    # Clean up resources.
+    agent_engine.delete(force=True)
 
 
 pytestmark = pytest_helper.setup(

tests/unit/vertexai/genai/replays/test_retrieve_agent_engine_memories.py

@@ -68,6 +68,8 @@ def test_retrieve_memories_with_similarity_search_params(client):
         )
         == 2
     )
+    # Clean up resources.
+    agent_engine.delete(force=True)
 
 
 def test_retrieve_memories_with_simple_retrieval_params(client):
@@ -87,6 +89,8 @@ def test_retrieve_memories_with_simple_retrieval_params(client):
     assert isinstance(memories, pagers.Pager)
     assert isinstance(memories.page[0], types.RetrieveMemoriesResponseRetrievedMemory)
     assert memories.page_size == 1
+    # Clean up resources.
+    agent_engine.delete(force=True)
 
 
 pytestmark = pytest_helper.setup(

tests/unit/vertexai/genai/test_agent_engines.py

@@ -617,6 +617,10 @@ def register_operations(self) -> Dict[str, List[str]]:
     "pydantic": ["pydantic"],
 }
 _TEST_OPERATION_NAME = "test_operation_name"
+_TEST_AGENT_ENGINE_ERROR = {
+    "message": "The following quotas are exceeded",
+    "code": 8,
+}
 
 
 def _create_empty_fake_package(package_name: str) -> str:
@@ -1830,6 +1834,49 @@ def setup_method(self):
         )
         self.test_agent = CapitalizeEngine()
 
+    @mock.patch.object(_agent_engines_utils, "_prepare")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
+    def test_create_agent_engine_error(self, mock_await_operation, mock_prepare):
+        mock_await_operation.return_value = _genai_types.AgentEngineOperation(
+            error=_TEST_AGENT_ENGINE_ERROR,
+        )
+        with mock.patch.object(
+            self.client.agent_engines._api_client, "request"
+        ) as request_mock:
+            request_mock.return_value = genai_types.HttpResponse(body="")
+            with pytest.raises(RuntimeError) as excinfo:
+                self.client.agent_engines.create(
+                    agent=self.test_agent,
+                    config=_genai_types.AgentEngineConfig(
+                        display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                        description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                        requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
+                        extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
+                        staging_bucket=_TEST_STAGING_BUCKET,
+                        gcs_dir_name=_TEST_GCS_DIR_NAME,
+                        env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT,
+                    ),
+                )
+                assert "Failed to create agent engine" in str(excinfo.value)
+
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
+    def test_update_agent_engine_description(self, mock_await_operation):
+        mock_await_operation.return_value = _genai_types.AgentEngineOperation(
+            error=_TEST_AGENT_ENGINE_ERROR,
+        )
+        with mock.patch.object(
+            self.client.agent_engines._api_client, "request"
+        ) as request_mock:
+            request_mock.return_value = genai_types.HttpResponse(body="")
+            with pytest.raises(RuntimeError) as excinfo:
+                self.client.agent_engines.update(
+                    name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                    config=_genai_types.AgentEngineConfig(
+                        description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                    ),
+                )
+                assert "Failed to update agent engine" in str(excinfo.value)
+
     @pytest.mark.parametrize(
         "test_case_name, test_operation_schemas, want_log_output",
         [

vertexai/_genai/agent_engines.py

@@ -1125,6 +1125,8 @@ def create(
             logger.info(
                 f"agent_engine=client.agent_engines.get('{agent_engine.api_resource.name}')"
             )
+        elif operation.error:
+            raise RuntimeError(f"Failed to create Agent Engine: {operation.error}")
         else:
             logger.warning("The operation returned an empty response.")
         if agent is not None:
@@ -1516,6 +1518,8 @@ def update(
             logger.info(
                 f"agent_engine=client.agent_engines.get('{agent_engine.api_resource.name}')"
             )
+        elif operation.error:
+            raise RuntimeError(f"Failed to update Agent Engine: {operation.error}")
         if agent_engine.api_resource.spec:
             self._register_api_methods(agent_engine=agent_engine)
         return agent_engine

vertexai/_genai/memories.py

@@ -1222,9 +1222,12 @@ def create(
                 )
             # We need to make a call to get the memory because the operation
             # response might not contain the relevant fields.
-            if not operation.response:
-                raise ValueError("Error retrieving memory.")
-            operation.response = self.get(name=operation.response.name)
+            if operation.response:
+                operation.response = self.get(name=operation.response.name)
+            elif operation.error:
+                raise RuntimeError(f"Failed to create memory: {operation.error}")
+            else:
+                raise RuntimeError("Error creating memory.")
         return operation
 
     def generate(
@@ -1285,10 +1288,14 @@ def generate(
         elif isinstance(config, dict):
             config = types.GenerateAgentEngineMemoriesConfig.model_validate(config)
         if config.wait_for_completion and not operation.done:
-            return _agent_engines_utils._await_operation(
+            operation = _agent_engines_utils._await_operation(
                 operation_name=operation.name,
                 get_operation_fn=self._get_generate_memories_operation,
             )
+            if not operation.response:
+                if operation.error:
+                    raise RuntimeError(f"Failed to generate memory: {operation.error}")
+                raise RuntimeError(f"Error generating memory: {operation}")
         return operation
 
     def list(

vertexai/_genai/sessions.py

@@ -792,12 +792,15 @@ def create(
                 operation_name=operation.name,
                 get_operation_fn=self._get_session_operation,
             )
-            if not operation.response:
-                raise ValueError(
+            if operation.response:
+                operation.response = self.get(name=operation.response.name)
+            elif operation.error:
+                raise RuntimeError(f"Failed to create session: {operation.error}")
+            else:
+                raise RuntimeError(
                     "Error retrieving session from the operation response. "
                     f"Operation name: {operation.name}"
                 )
-            operation.response = self.get(name=operation.response.name)
         return operation
 
     def list(