fix(firestore): make the Firestore session and memory services usable

src/google/adk/integrations/firestore/__init__.py

@@ -12,6 +12,41 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Firestore integrations for ADK.
+
+This module provides session and memory services backed by Google Cloud
+Firestore. They require the optional ``google-cloud-firestore`` package.
+"""
+
 from __future__ import annotations
 
-"""Firestore integrations for ADK."""
+import typing
+
+if typing.TYPE_CHECKING:
+  from .firestore_memory_service import FirestoreMemoryService
+  from .firestore_session_service import FirestoreSessionService
+
+# Map attribute names to relative module paths.
+_lazy_imports = {
+    "FirestoreMemoryService": ".firestore_memory_service",
+    "FirestoreSessionService": ".firestore_session_service",
+}
+
+__all__ = [
+    "FirestoreMemoryService",
+    "FirestoreSessionService",
+]
+
+
+def __getattr__(name: str) -> typing.Any:
+  if name in _lazy_imports:
+    import importlib
+
+    module_path = _lazy_imports[name]
+    module = importlib.import_module(module_path, __name__)
+    return getattr(module, name)
+  raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+  return list(_lazy_imports.keys())

src/google/adk/integrations/firestore/firestore_memory_service.py

@@ -15,6 +15,9 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Mapping
+from collections.abc import Sequence
+import hashlib
 import logging
 import re
 from typing import Optional
@@ -32,6 +35,7 @@
 if TYPE_CHECKING:
   from google.cloud import firestore
 
+  from ...events.event import Event
   from ...sessions.session import Session
 
 logger = logging.getLogger("google_adk." + __name__)
@@ -40,6 +44,18 @@
 DEFAULT_MEMORIES_COLLECTION = "memories"
 
 
+def _memory_doc_id(
+    *, app_name: str, user_id: str, session_id: Optional[str], event_id: str
+) -> str:
+  """Returns a stable memory document ID for an event.
+
+  Hashed because app names and user IDs may contain characters that are not
+  allowed in document IDs, such as "/".
+  """
+  key = "\x00".join((app_name, user_id, session_id or "", event_id))
+  return hashlib.sha256(key.encode("utf-8")).hexdigest()
+
+
 class FirestoreMemoryService(BaseMemoryService):  # type: ignore[misc]
   """Memory service that uses Google Cloud Firestore as the backend.
 
@@ -82,10 +98,50 @@ def __init__(
   @override
   async def add_session_to_memory(self, session: Session) -> None:
     """Extracts keywords from session events and stores them in the memories collection."""
+    await self._write_memories(
+        app_name=session.app_name,
+        user_id=session.user_id,
+        session_id=session.id,
+        events=session.events,
+    )
+
+  @override
+  async def add_events_to_memory(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      events: Sequence[Event],
+      session_id: Optional[str] = None,
+      custom_metadata: Optional[Mapping[str, object]] = None,
+  ) -> None:
+    """Adds events, such as the latest turn, to the memories collection.
+
+    Re-adding an event with the same session ID overwrites its entry. The
+    session ID is part of that entry's ID, so an event added without one and
+    later added with one is stored twice; `search_memory` returns one copy.
+    """
+    _ = custom_metadata
+    await self._write_memories(
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session_id,
+        events=events,
+    )
+
+  async def _write_memories(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: Optional[str],
+      events: Sequence[Event],
+  ) -> None:
+    """Writes one memory document per event that has text keywords."""
     batch = self.client.batch()
     count = 0
 
-    for event in session.events:
+    for event in events:
       if not event.content or not event.content.parts:
         continue
 
@@ -97,12 +153,20 @@ async def add_session_to_memory(self, session: Session) -> None:
       if not keywords:
         continue
 
-      doc_ref = self.client.collection(self.memories_collection).document()
+      doc_ref = self.client.collection(self.memories_collection).document(
+          _memory_doc_id(
+              app_name=app_name,
+              user_id=user_id,
+              session_id=session_id,
+              event_id=event.id,
+          )
+      )
       batch.set(
           doc_ref,
           {
-              "appName": session.app_name,
-              "userId": session.user_id,
+              "appName": app_name,
+              "userId": user_id,
+              "sessionId": session_id,
               "keywords": list(keywords),
               "author": event.author,
               "content": event.content.model_dump(

src/google/adk/integrations/firestore/firestore_session_service.py

@@ -42,6 +42,7 @@
 
 try:
   from google.cloud import firestore
+  from google.cloud.firestore_v1.base_query import FieldFilter
 except ImportError as e:
   raise ImportError(
       "FirestoreSessionService requires google-cloud-firestore. "
@@ -245,14 +246,15 @@ async def _create_txn(
           (user_snap.to_dict() or {}) if user_snap.exists else {}
       )
 
-      # 2. Writes
+      # 2. Writes. Documents are written whole: a merge write deep-merges
+      # nested maps and would keep keys that the new value dropped.
       if app_state_delta:
         current_app.update(app_state_delta)
-        transaction.set(app_ref, current_app, merge=True)
+        transaction.set(app_ref, current_app)
 
       if user_state_delta:
         current_user.update(user_state_delta)
-        transaction.set(user_ref, current_user, merge=True)
+        transaction.set(user_ref, current_user)
 
       transaction.set(session_ref, session_data)
       return current_app, current_user
@@ -324,7 +326,7 @@ async def get_session(
           after_dt = datetime.fromtimestamp(
               config.after_timestamp, tz=timezone.utc
           )
-          query = query.where("timestamp", ">=", after_dt)
+          query = query.where(filter=FieldFilter("timestamp", ">=", after_dt))
         if config.num_recent_events is not None:
           query = query.limit_to_last(config.num_recent_events)
 
@@ -368,12 +370,12 @@ async def list_sessions(
     """Lists sessions from Firestore."""
     if user_id:
       query = self._get_sessions_ref(app_name, user_id).where(
-          "appName", "==", app_name
+          filter=FieldFilter("appName", "==", app_name)
       )
       docs = await query.get()
     else:
       query = self.client.collection_group(self.sessions_collection).where(
-          "appName", "==", app_name
+          filter=FieldFilter("appName", "==", app_name)
       )
       docs = await query.get()
 
@@ -441,6 +443,21 @@ def _iter_sessions_data() -> Iterator[dict[str, Any]]:
     sessions.sort(key=lambda s: (s.last_update_time, s.user_id, s.id))
     return ListSessionsResponse(sessions=sessions)
 
+  async def get_user_state(
+      self, *, app_name: str, user_id: str
+  ) -> dict[str, Any]:
+    """Gets the user-scoped state from Firestore."""
+    user_ref = (
+        self.client.collection(self.user_state_collection)
+        .document(app_name)
+        .collection("users")
+        .document(user_id)
+    )
+    user_doc = await user_ref.get()
+    if not user_doc.exists:
+      return {}
+    return user_doc.to_dict() or {}
+
   async def delete_session(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
@@ -544,16 +561,16 @@ async def _append_txn(transaction: firestore.AsyncTransaction) -> int:
             else None
         )
 
-        # 2. Writes
+        # 2. Writes. Documents are written whole, as in create_session.
         if app_updates and app_snap is not None:
           current_app = (app_snap.to_dict() or {}) if app_snap.exists else {}
           current_app.update(app_updates)
-          transaction.set(app_ref, current_app, merge=True)
+          transaction.set(app_ref, current_app)
 
         if user_updates and user_snap is not None:
-          current_user = user_snap.to_dict() if user_snap.exists else {}
+          current_user = (user_snap.to_dict() or {}) if user_snap.exists else {}
           current_user.update(user_updates)
-          transaction.set(user_ref, current_user, merge=True)
+          transaction.set(user_ref, current_user)
 
         new_revision = current_revision + 1
 
@@ -588,7 +605,10 @@ async def _append_txn(transaction: firestore.AsyncTransaction) -> int:
             event_ref,
             {
                 "event_data": event_data,
-                "timestamp": firestore.SERVER_TIMESTAMP,
+                # Event time, not write time: after_timestamp filters on it.
+                "timestamp": datetime.fromtimestamp(
+                    event.timestamp, tz=timezone.utc
+                ),
                 "appName": session.app_name,
                 "userId": session.user_id,
             },