chore: Make VertexAiSessionService fully asynchronous

This CL refactors VertexAiSessionService to use the asynchronous aio client for all Vertex AI API calls. This ensures that the service methods are non-blocking and can be used effectively in an asyncio environment.

PiperOrigin-RevId: 824573356

Author: DeanChensj

src/google/adk/sessions/vertex_ai_session_service.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import asyncio
 import datetime
 import json
 import logging
@@ -104,7 +105,7 @@ async def create_session(
 
     if _is_vertex_express_mode(self._project, self._location):
       config['wait_for_completion'] = False
-      api_response = api_client.agent_engines.sessions.create(
+      api_response = await api_client.aio.agent_engines.sessions.create(
           name=f'reasoningEngines/{reasoning_engine_id}',
           user_id=user_id,
           config=config,
@@ -123,7 +124,7 @@ async def create_session(
       )
       async def _poll_session_resource():
         try:
-          return api_client.agent_engines.sessions.get(
+          return await api_client.aio.agent_engines.sessions.get(
               name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
           )
         except ClientError:
@@ -135,11 +136,11 @@ async def _poll_session_resource():
       except Exception as exc:
         raise ValueError('Failed to create session.') from exc
 
-      get_session_response = api_client.agent_engines.sessions.get(
+      get_session_response = await api_client.aio.agent_engines.sessions.get(
           name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
       )
     else:
-      api_response = api_client.agent_engines.sessions.create(
+      api_response = await api_client.aio.agent_engines.sessions.create(
           name=f'reasoningEngines/{reasoning_engine_id}',
           user_id=user_id,
           config=config,
@@ -168,10 +169,28 @@ async def get_session(
   ) -> Optional[Session]:
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
     api_client = self._get_api_client()
+    session_resource_name = (
+        f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
+    )
+
+    # Get session resource and events in parallel.
+    list_events_kwargs = {}
+    if config and not config.num_recent_events and config.after_timestamp:
+      # Filter events based on timestamp.
+      list_events_kwargs['config'] = {
+          'filter': 'timestamp>="{}"'.format(
+              datetime.datetime.fromtimestamp(
+                  config.after_timestamp, tz=datetime.timezone.utc
+              ).isoformat()
+          )
+      }
 
-    # Get session resource
-    get_session_response = api_client.agent_engines.sessions.get(
-        name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
+    get_session_response, events_iterator = await asyncio.gather(
+        api_client.aio.agent_engines.sessions.get(name=session_resource_name),
+        api_client.aio.agent_engines.sessions.events.list(
+            name=session_resource_name,
+            **list_events_kwargs,
+        ),
     )
 
     if get_session_response.user_id != user_id:
@@ -187,29 +206,14 @@ async def get_session(
         state=getattr(get_session_response, 'session_state', None) or {},
         last_update_time=update_timestamp,
     )
-
-    list_events_kwargs = {}
-    if config and not config.num_recent_events and config.after_timestamp:
-      list_events_kwargs['config'] = {
-          'filter': 'timestamp>="{}"'.format(
-              datetime.datetime.fromtimestamp(
-                  config.after_timestamp, tz=datetime.timezone.utc
-              ).isoformat()
-          )
-      }
-
-    events_iterator = api_client.agent_engines.sessions.events.list(
-        name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
-        **list_events_kwargs,
-    )
     session.events += [
         _from_api_event(event)
         for event in events_iterator
         if event.timestamp.timestamp() <= update_timestamp
     ]
 
-    # Filter events based on config
     if config:
+      # Filter events based on num_recent_events.
       if config.num_recent_events:
         session.events = session.events[-config.num_recent_events :]
 
@@ -226,7 +230,7 @@ async def list_sessions(
     config = {}
     if user_id is not None:
       config['filter'] = f'user_id="{user_id}"'
-    sessions_iterator = api_client.agent_engines.sessions.list(
+    sessions_iterator = await api_client.aio.agent_engines.sessions.list(
         name=f'reasoningEngines/{reasoning_engine_id}',
         config=config,
     )
@@ -251,7 +255,7 @@ async def delete_session(
     api_client = self._get_api_client()
 
     try:
-      api_client.agent_engines.sessions.delete(
+      await api_client.aio.agent_engines.sessions.delete(
           name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
       )
     except Exception as e:
@@ -308,7 +312,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
       )
     config['event_metadata'] = metadata_dict
 
-    api_client.agent_engines.sessions.events.append(
+    await api_client.aio.agent_engines.sessions.events.append(
         name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}',
         author=event.author,
         invocation_id=event.invocation_id,

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -235,22 +235,24 @@ def __init__(self) -> None:
     """Initializes MockClient."""
     self.session_dict: dict[str, Any] = {}
     self.event_dict: dict[str, Tuple[List[Any], Optional[str]]] = {}
-    self.agent_engines = mock.Mock()
-    self.agent_engines.sessions.get.side_effect = self._get_session
-    self.agent_engines.sessions.list.side_effect = self._list_sessions
-    self.agent_engines.sessions.delete.side_effect = self._delete_session
-    self.agent_engines.sessions.create.side_effect = self._create_session
-    self.agent_engines.sessions.events.list.side_effect = self._list_events
-    self.agent_engines.sessions.events.append.side_effect = self._append_event
+    self.aio = mock.Mock()
+    self.aio.agent_engines.sessions.get.side_effect = self._get_session
+    self.aio.agent_engines.sessions.list.side_effect = self._list_sessions
+    self.aio.agent_engines.sessions.delete.side_effect = self._delete_session
+    self.aio.agent_engines.sessions.create.side_effect = self._create_session
+    self.aio.agent_engines.sessions.events.list.side_effect = self._list_events
+    self.aio.agent_engines.sessions.events.append.side_effect = (
+        self._append_event
+    )
     self.last_create_session_config: dict[str, Any] = {}
 
-  def _get_session(self, name: str):
+  async def _get_session(self, name: str):
     session_id = name.split('/')[-1]
     if session_id in self.session_dict:
       return _convert_to_object(self.session_dict[session_id])
     raise api_core_exceptions.NotFound(f'Session not found: {session_id}')
 
-  def _list_sessions(self, name: str, config: dict[str, Any]):
+  async def _list_sessions(self, name: str, config: dict[str, Any]):
     filter_val = config.get('filter', '')
     user_id_match = re.search(r'user_id="([^"]+)"', filter_val)
     if user_id_match:
@@ -271,11 +273,13 @@ def _list_sessions(self, name: str, config: dict[str, Any]):
         _convert_to_object(session) for session in self.session_dict.values()
     ]
 
-  def _delete_session(self, name: str):
+  async def _delete_session(self, name: str):
     session_id = name.split('/')[-1]
     self.session_dict.pop(session_id)
 
-  def _create_session(self, name: str, user_id: str, config: dict[str, Any]):
+  async def _create_session(
+      self, name: str, user_id: str, config: dict[str, Any]
+  ):
     self.last_create_session_config = config
     new_session_id = '4'
     self.session_dict[new_session_id] = {
@@ -299,7 +303,7 @@ def _create_session(self, name: str, user_id: str, config: dict[str, Any]):
         'response': self.session_dict['4'],
     })
 
-  def _list_events(self, name: str, **kwargs):
+  async def _list_events(self, name: str, **kwargs):
     session_id = name.split('/')[-1]
     events = []
     if session_id in self.event_dict:
@@ -322,7 +326,7 @@ def _list_events(self, name: str, **kwargs):
         ]
     return [_convert_to_object(event) for event in events]
 
-  def _append_event(
+  async def _append_event(
       self,
       name: str,
       author: str,