Skip to content

Commit dff032f

Browse files
committed
feat(sessions): add flat_layout option for configurable collection hierarchy
Add flat_layout=True constructor parameter so developers can store sessions directly in root_collection/{session_id} instead of the default nested root/{app}/users/{user}/sessions/{session_id} path. Useful when the session id already encodes the user (e.g. {phone}-{date}), matching an existing flat Firestore collection layout. list_sessions adds a userId field filter automatically in flat mode.
1 parent 761842e commit dff032f

1 file changed

Lines changed: 18 additions & 1 deletion

File tree

src/google/adk_community/sessions/firestore_session_service.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ def __init__(
144144
events_collection: str = DEFAULT_EVENTS_COLLECTION,
145145
app_state_collection: str = DEFAULT_APP_STATE_COLLECTION,
146146
user_state_collection: str = DEFAULT_USER_STATE_COLLECTION,
147+
flat_layout: bool = False,
147148
durable_mode: bool = False,
148149
buffer_max_events: int = 10,
149150
flush_interval_seconds: float = 120.0,
@@ -166,6 +167,10 @@ def __init__(
166167
``'app_states'``.
167168
user_state_collection: Root collection for user-scoped state. Defaults
168169
to ``'user_states'``.
170+
flat_layout: When True, session documents live directly in
171+
``root_collection/{session_id}`` (no ``{app}/users/{user}/sessions/``
172+
nesting). Useful when the session id already encodes the user (e.g.
173+
``{phone}-{date}``). Defaults to False.
169174
durable_mode: When True, every event is persisted immediately and no
170175
buffering happens.
171176
buffer_max_events: Flush a session once this many events are buffered.
@@ -190,6 +195,9 @@ def __init__(
190195
self.events_collection = events_collection
191196
self.app_state_collection = app_state_collection
192197
self.user_state_collection = user_state_collection
198+
# flat_layout=True: sessions/{session_id} (no {app}/users/{user} nesting)
199+
# flat_layout=False (default): {root}/{app}/users/{user}/{sessions}/{session_id}
200+
self._flat_layout = flat_layout
193201

194202
self._durable_mode = durable_mode
195203
self._buffer_max_events = buffer_max_events
@@ -211,6 +219,8 @@ def __init__(
211219
# -- Firestore refs / helpers ---------------------------------------------
212220

213221
def _get_sessions_ref(self, app_name: str, user_id: str) -> Any:
222+
if self._flat_layout:
223+
return self.client.collection(self.root_collection)
214224
return (
215225
self.client.collection(self.root_collection)
216226
.document(app_name)
@@ -371,7 +381,14 @@ async def list_sessions(
371381
self, *, app_name: str, user_id: Optional[str] = None
372382
) -> ListSessionsResponse:
373383
"""Lists sessions for an app (optionally a single user)."""
374-
if user_id:
384+
if self._flat_layout:
385+
query = self.client.collection(self.root_collection).where(
386+
"appName", "==", app_name
387+
)
388+
if user_id:
389+
query = query.where("userId", "==", user_id)
390+
docs = await query.get()
391+
elif user_id:
375392
docs = await (
376393
self._get_sessions_ref(app_name, user_id)
377394
.where("appName", "==", app_name)

0 commit comments

Comments
 (0)