From 09d9916dc4bb43ee0a4da80b7c6701fb1ff8e69b Mon Sep 17 00:00:00 2001
From: Shangjie Chen <deanchen@google.com>
Date: Thu, 26 Jun 2025 18:52:59 -0700
Subject: [PATCH] chore: Catch exection from GenAI API client in vertex
 services

PiperOrigin-RevId: 776368177
---
 .../memory/vertex_ai_memory_bank_service.py   | 48 +++++++++++--------
 .../adk/sessions/vertex_ai_session_service.py | 33 ++++++++-----
 2 files changed, 50 insertions(+), 31 deletions(-)

diff --git a/src/google/adk/memory/vertex_ai_memory_bank_service.py b/src/google/adk/memory/vertex_ai_memory_bank_service.py
index 083b48e8d..fd5d8c822 100644
--- a/src/google/adk/memory/vertex_ai_memory_bank_service.py
+++ b/src/google/adk/memory/vertex_ai_memory_bank_service.py
@@ -68,6 +68,7 @@ async def add_session_to_memory(self, session: Session):
         events.append({
             'content': event.content.model_dump(exclude_none=True, mode='json')
         })
+
     request_dict = {
         'direct_contents_source': {
             'events': events,
@@ -79,11 +80,16 @@ async def add_session_to_memory(self, session: Session):
     }
 
     if events:
-      api_response = await api_client.async_request(
-          http_method='POST',
-          path=f'reasoningEngines/{self._agent_engine_id}/memories:generate',
-          request_dict=request_dict,
-      )
+      try:
+        api_response = await api_client.async_request(
+            http_method='POST',
+            path=f'reasoningEngines/{self._agent_engine_id}/memories:generate',
+            request_dict=request_dict,
+        )
+      except Exception as e:
+        raise ValueError(f'Error adding session to memory: {e}') from e
+
+      api_response = _convert_api_response(api_response)
       logger.info(f'Generate memory response: {api_response}')
     else:
       logger.info('No events to add to memory.')
@@ -92,21 +98,25 @@ async def add_session_to_memory(self, session: Session):
   async def search_memory(self, *, app_name: str, user_id: str, query: str):
     api_client = self._get_api_client()
 
-    api_response = await api_client.async_request(
-        http_method='POST',
-        path=f'reasoningEngines/{self._agent_engine_id}/memories:retrieve',
-        request_dict={
-            'scope': {
-                'app_name': app_name,
-                'user_id': user_id,
-            },
-            'similarity_search_params': {
-                'search_query': query,
-            },
-        },
-    )
+    try:
+      api_response = await api_client.async_request(
+          http_method='POST',
+          path=f'reasoningEngines/{self._agent_engine_id}/memories:retrieve',
+          request_dict={
+              'scope': {
+                  'app_name': app_name,
+                  'user_id': user_id,
+              },
+              'similarity_search_params': {
+                  'search_query': query,
+              },
+          },
+      )
+    except Exception as e:
+      raise ValueError(f'Error retrieving memory: {e}') from e
+
     api_response = _convert_api_response(api_response)
-    logger.info(f'Search memory response: {api_response}')
+    logger.info(f'Retrieve memory response: {api_response}')
 
     if not api_response or not api_response.get('retrievedMemories', None):
       return SearchMemoryResponse()
diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py
index 06a904c89..662368af0 100644
--- a/src/google/adk/sessions/vertex_ai_session_service.py
+++ b/src/google/adk/sessions/vertex_ai_session_service.py
@@ -78,6 +78,7 @@ async def create_session(
           'User-provided Session id is not supported for'
           ' VertexAISessionService.'
       )
+
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
     api_client = self._get_api_client()
 
@@ -85,11 +86,15 @@ async def create_session(
     if state:
       session_json_dict['session_state'] = state
 
-    api_response = await api_client.async_request(
-        http_method='POST',
-        path=f'reasoningEngines/{reasoning_engine_id}/sessions',
-        request_dict=session_json_dict,
-    )
+    try:
+      api_response = await api_client.async_request(
+          http_method='POST',
+          path=f'reasoningEngines/{reasoning_engine_id}/sessions',
+          request_dict=session_json_dict,
+      )
+    except Exception as e:
+      raise ValueError(f'Error creating session: {e}') from e
+
     api_response = _convert_api_response(api_response)
     logger.info(f'Create Session response {api_response}')
 
@@ -288,8 +293,7 @@ async def delete_session(
           request_dict={},
       )
     except Exception as e:
-      logger.error(f'Error deleting session {session_id}: {e}')
-      raise e
+      raise ValueError(f'Error deleting session: {e}') from e
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
@@ -298,11 +302,16 @@ async def append_event(self, session: Session, event: Event) -> Event:
 
     reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)
     api_client = self._get_api_client()
-    await api_client.async_request(
-        http_method='POST',
-        path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}:appendEvent',
-        request_dict=_convert_event_to_json(event),
-    )
+
+    try:
+      await api_client.async_request(
+          http_method='POST',
+          path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}:appendEvent',
+          request_dict=_convert_event_to_json(event),
+      )
+    except Exception as e:
+      raise ValueError(f'Error appending event: {e}') from e
+
     return event
 
   def _get_reasoning_engine_id(self, app_name: str):