diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 6fd0dfc7..571ae53a 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -107,19 +107,18 @@ class VertexAiSessionService(BaseSessionService): ) reasoning_engine_id = self._get_reasoning_engine_id(app_name) - api_client = self._get_api_client() config = {'session_state': state} if state else {} config.update(kwargs) - - api_response = await api_client.aio.agent_engines.sessions.create( - name=f'reasoningEngines/{reasoning_engine_id}', - user_id=user_id, - config=config, - ) - logger.debug('Create session response: %s', api_response) - get_session_response = api_response.response - session_id = get_session_response.name.split('/')[-1] + async with self._get_api_client() as api_client: + api_response = await api_client.agent_engines.sessions.create( + name=f'reasoningEngines/{reasoning_engine_id}', + user_id=user_id, + config=config, + ) + logger.debug('Create session response: %s', api_response) + get_session_response = api_response.response + session_id = get_session_response.name.split('/')[-1] session = Session( app_name=app_name, @@ -140,30 +139,29 @@ class VertexAiSessionService(BaseSessionService): config: Optional[GetSessionConfig] = None, ) -> 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}' ) + async with self._get_api_client() as api_client: + # 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 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_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, - ), - ) + get_session_response, events_iterator = await asyncio.gather( + api_client.agent_engines.sessions.get(name=session_resource_name), + api_client.agent_engines.sessions.events.list( + name=session_resource_name, + **list_events_kwargs, + ), + ) if get_session_response.user_id != user_id: raise ValueError( @@ -196,43 +194,45 @@ class VertexAiSessionService(BaseSessionService): self, *, app_name: str, user_id: Optional[str] = None ) -> ListSessionsResponse: reasoning_engine_id = self._get_reasoning_engine_id(app_name) - api_client = self._get_api_client() - sessions = [] - config = {} - if user_id is not None: - config['filter'] = f'user_id="{user_id}"' - sessions_iterator = await api_client.aio.agent_engines.sessions.list( - name=f'reasoningEngines/{reasoning_engine_id}', - config=config, - ) - - for api_session in sessions_iterator: - sessions.append( - Session( - app_name=app_name, - user_id=api_session.user_id, - id=api_session.name.split('/')[-1], - state=getattr(api_session, 'session_state', None) or {}, - last_update_time=api_session.update_time.timestamp(), - ) + async with self._get_api_client() as api_client: + sessions = [] + config = {} + if user_id is not None: + config['filter'] = f'user_id="{user_id}"' + sessions_iterator = await api_client.agent_engines.sessions.list( + name=f'reasoningEngines/{reasoning_engine_id}', + config=config, ) + for api_session in sessions_iterator: + sessions.append( + Session( + app_name=app_name, + user_id=api_session.user_id, + id=api_session.name.split('/')[-1], + state=getattr(api_session, 'session_state', None) or {}, + last_update_time=api_session.update_time.timestamp(), + ) + ) + return ListSessionsResponse(sessions=sessions) async def delete_session( self, *, app_name: str, user_id: str, session_id: str ) -> None: reasoning_engine_id = self._get_reasoning_engine_id(app_name) - api_client = self._get_api_client() - try: - await api_client.aio.agent_engines.sessions.delete( - name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}', - ) - except Exception as e: - logger.error('Error deleting session %s: %s', session_id, e) - raise e + async with self._get_api_client() as api_client: + try: + await api_client.agent_engines.sessions.delete( + name=( + f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}' + ), + ) + except Exception as e: + logger.error('Error deleting session %s: %s', session_id, e) + raise @override async def append_event(self, session: Session, event: Event) -> Event: @@ -240,7 +240,6 @@ class VertexAiSessionService(BaseSessionService): await super().append_event(session=session, event=event) reasoning_engine_id = self._get_reasoning_engine_id(session.app_name) - api_client = self._get_api_client() config = {} if event.content: @@ -284,15 +283,16 @@ class VertexAiSessionService(BaseSessionService): ) config['event_metadata'] = metadata_dict - 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, - timestamp=datetime.datetime.fromtimestamp( - event.timestamp, tz=datetime.timezone.utc - ), - config=config, - ) + async with self._get_api_client() as api_client: + await api_client.agent_engines.sessions.events.append( + name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}', + author=event.author, + invocation_id=event.invocation_id, + timestamp=datetime.datetime.fromtimestamp( + event.timestamp, tz=datetime.timezone.utc + ), + config=config, + ) return event def _get_reasoning_engine_id(self, app_name: str): @@ -318,7 +318,7 @@ class VertexAiSessionService(BaseSessionService): ) -> Optional[Union[types.HttpOptions, types.HttpOptionsDict]]: return None - def _get_api_client(self) -> vertexai.Client: + def _get_api_client(self) -> vertexai.AsyncClient: """Instantiates an API client for the given project and location. Returns: @@ -331,7 +331,7 @@ class VertexAiSessionService(BaseSessionService): location=self._location, http_options=self._api_client_http_options_override(), api_key=self._express_mode_api_key, - ) + ).aio def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event: diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index fd004199..29c3c74c 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -228,24 +228,30 @@ def _convert_to_object(data): return data -class MockApiClient: +class MockAsyncClient: """Mocks the API Client.""" def __init__(self) -> None: """Initializes MockClient.""" self.session_dict: dict[str, Any] = {} self.event_dict: dict[str, Tuple[List[Any], Optional[str]]] = {} - 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.agent_engines = mock.AsyncMock() + 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.last_create_session_config: dict[str, Any] = {} + async def __aenter__(self): + """Enters the asynchronous context.""" + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Exits the asynchronous context.""" + pass + async def _get_session(self, name: str): session_id = name.split('/')[-1] if session_id in self.session_dict: @@ -374,7 +380,7 @@ def mock_vertex_ai_session_service( @pytest.fixture def mock_api_client_instance(): """Creates a mock API client instance for testing.""" - api_client = MockApiClient() + api_client = MockAsyncClient() api_client.session_dict = { '1': MOCK_SESSION_JSON_1, '2': MOCK_SESSION_JSON_2,