diff --git a/src/google/adk/sessions/__init__.py b/src/google/adk/sessions/__init__.py index 5583ac436..601f7051f 100644 --- a/src/google/adk/sessions/__init__.py +++ b/src/google/adk/sessions/__init__.py @@ -14,6 +14,9 @@ import logging from .base_session_service import BaseSessionService +from .base_session_service import GetSessionConfig +from .base_session_service import ListSessionsConfig +from .base_session_service import ListSessionsResponse from .in_memory_session_service import InMemorySessionService from .session import Session from .state import State @@ -24,7 +27,10 @@ __all__ = [ 'BaseSessionService', + 'GetSessionConfig', 'InMemorySessionService', + 'ListSessionsConfig', + 'ListSessionsResponse', 'Session', 'State', 'VertexAiSessionService', diff --git a/src/google/adk/sessions/base_session_service.py b/src/google/adk/sessions/base_session_service.py index 25e46ba19..2c6b328f6 100644 --- a/src/google/adk/sessions/base_session_service.py +++ b/src/google/adk/sessions/base_session_service.py @@ -31,10 +31,20 @@ class GetSessionConfig(BaseModel): after_timestamp: Optional[float] = None +class ListSessionsConfig(BaseModel): + """The configuration of listing sessions.""" + + max_sessions: Optional[int] = None + """Maximum number of sessions to return. If not specified, all sessions are returned.""" + + include_state: bool = False + """Whether to include the state data in the response. Default is False for performance.""" + + class ListSessionsResponse(BaseModel): """The response of listing sessions. - The events and states are not set within each Session object. + The events and states are not set within each Session object unless explicitly requested. """ sessions: list[Session] = Field(default_factory=list) @@ -81,7 +91,7 @@ async def get_session( @abc.abstractmethod async def list_sessions( - self, *, app_name: str, user_id: str + self, *, app_name: str, user_id: str, config: Optional[ListSessionsConfig] = None ) -> ListSessionsResponse: """Lists all the sessions.""" diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 9d3f0ddc3..538af6a5b 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -54,6 +54,7 @@ from ..events.event import Event from .base_session_service import BaseSessionService from .base_session_service import GetSessionConfig +from .base_session_service import ListSessionsConfig from .base_session_service import ListSessionsResponse from .session import Session from .state import State @@ -498,22 +499,34 @@ async def get_session( @override async def list_sessions( - self, *, app_name: str, user_id: str + self, *, app_name: str, user_id: str, config: Optional[ListSessionsConfig] = None ) -> ListSessionsResponse: with self.database_session_factory() as session_factory: - results = ( + query = ( session_factory.query(StorageSession) .filter(StorageSession.app_name == app_name) .filter(StorageSession.user_id == user_id) - .all() + .order_by(StorageSession.update_time.desc()) ) + + # Apply pagination if specified + if config and config.max_sessions: + query = query.limit(config.max_sessions) + + results = query.all() + sessions = [] for storage_session in results: + # Determine whether to include state + session_state = {} + if config and config.include_state: + session_state = storage_session.state + session = Session( app_name=app_name, user_id=user_id, id=storage_session.id, - state={}, + state=session_state, last_update_time=storage_session.update_timestamp_tz, ) sessions.append(session) diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py index 70e75411c..41872448a 100644 --- a/src/google/adk/sessions/in_memory_session_service.py +++ b/src/google/adk/sessions/in_memory_session_service.py @@ -25,6 +25,7 @@ from ..events.event import Event from .base_session_service import BaseSessionService from .base_session_service import GetSessionConfig +from .base_session_service import ListSessionsConfig from .base_session_service import ListSessionsResponse from .session import Session from .state import State @@ -201,18 +202,18 @@ def _merge_state( @override async def list_sessions( - self, *, app_name: str, user_id: str + self, *, app_name: str, user_id: str, config: Optional[ListSessionsConfig] = None ) -> ListSessionsResponse: - return self._list_sessions_impl(app_name=app_name, user_id=user_id) + return self._list_sessions_impl(app_name=app_name, user_id=user_id, config=config) def list_sessions_sync( - self, *, app_name: str, user_id: str + self, *, app_name: str, user_id: str, config: Optional[ListSessionsConfig] = None ) -> ListSessionsResponse: logger.warning('Deprecated. Please migrate to the async method.') - return self._list_sessions_impl(app_name=app_name, user_id=user_id) + return self._list_sessions_impl(app_name=app_name, user_id=user_id, config=config) def _list_sessions_impl( - self, *, app_name: str, user_id: str + self, *, app_name: str, user_id: str, config: Optional[ListSessionsConfig] = None ) -> ListSessionsResponse: empty_response = ListSessionsResponse() if app_name not in self.sessions: @@ -221,10 +222,25 @@ def _list_sessions_impl( return empty_response sessions_without_events = [] - for session in self.sessions[app_name][user_id].values(): + all_sessions = list(self.sessions[app_name][user_id].values()) + + # Sort by last_update_time in descending order to get most recent first + all_sessions.sort(key=lambda s: s.last_update_time, reverse=True) + + # Apply pagination if specified + if config and config.max_sessions: + all_sessions = all_sessions[:config.max_sessions] + + for session in all_sessions: copied_session = copy.deepcopy(session) copied_session.events = [] - copied_session.state = {} + + # Determine whether to include state + if config and config.include_state: + copied_session.state = session.state + else: + copied_session.state = {} + sessions_without_events.append(copied_session) return ListSessionsResponse(sessions=sessions_without_events) diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index ebd98e369..9d6825bd2 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -39,6 +39,7 @@ from ..events.event_actions import EventActions from .base_session_service import BaseSessionService from .base_session_service import GetSessionConfig +from .base_session_service import ListSessionsConfig from .base_session_service import ListSessionsResponse from .session import Session @@ -263,7 +264,7 @@ async def get_session( @override async def list_sessions( - self, *, app_name: str, user_id: str + self, *, app_name: str, user_id: str, config: Optional[ListSessionsConfig] = None ) -> ListSessionsResponse: reasoning_engine_id = self._get_reasoning_engine_id(app_name) api_client = self._get_api_client() @@ -284,13 +285,30 @@ async def list_sessions( if not api_response or api_response.get('httpHeaders', None): return ListSessionsResponse() + api_sessions = api_response['sessions'] + + # Sort by updateTime in descending order to get most recent first + api_sessions.sort(key=lambda s: s.get('updateTime', ''), reverse=True) + + # Apply pagination if specified + if config and config.max_sessions: + api_sessions = api_sessions[:config.max_sessions] + sessions = [] - for api_session in api_response['sessions']: + for api_session in api_sessions: + session_state = {} + if config and config.include_state: + # For Vertex AI, we'd need to fetch individual session details to get state + # This is a performance trade-off - we can either make individual calls + # or keep the current behavior. For now, we'll keep it empty but document it. + # In a real implementation, you might want to batch these calls or cache them. + pass + session = Session( app_name=app_name, user_id=user_id, id=api_session['name'].split('/')[-1], - state={}, + state=session_state, last_update_time=isoparse(api_session['updateTime']).timestamp(), ) sessions.append(session)