Ensure memory generation does not go over the firestore batch limit

Author: ScottMansfield

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

@@ -86,7 +86,7 @@ def __init__(
   async def add_session_to_memory(self, session: Session) -> None:
     """Extracts keywords from session events and stores them in the memories collection."""
     batch = self.client.batch()
-    has_updates = False
+    count = 0
 
     for event in session.events:
       if not event.content or not event.content.parts:
@@ -114,9 +114,13 @@ async def add_session_to_memory(self, session: Session) -> None:
               "timestamp": event.timestamp,
           },
       )
-      has_updates = True
+      count += 1
+      if count >= 500:
+        await batch.commit()
+        batch = self.client.batch()
+        count = 0
 
-    if has_updates:
+    if count > 0:
       await batch.commit()
 
   def _extract_keywords(self, text: str) -> set[str]:

tests/unittests/integrations/firestore/test_firestore_memory_service.py

@@ -348,3 +348,37 @@ async def test_add_session_to_memory_commit_error(mock_firestore_client):
 
   with pytest.raises(Exception, match="Firestore commit failed"):
     await service.add_session_to_memory(session)
+
+
+@pytest.mark.asyncio
+async def test_add_session_to_memory_exceeds_batch_limit(mock_firestore_client):
+  service = FirestoreMemoryService(client=mock_firestore_client)
+
+  from google.adk.sessions.session import Session
+
+  session = Session(id="test_session", app_name="test_app", user_id="test_user")
+
+  for i in range(501):
+    content = types.Content(parts=[types.Part.from_text(text=f"event keyword {i}")])
+    event = Event(
+        invocation_id=f"test_inv_{i}",
+        author="user",
+        content=content,
+        timestamp=1234567890.0 + i,
+    )
+    session.events.append(event)
+
+  batch1 = mock.MagicMock()
+  batch2 = mock.MagicMock()
+  batch1.commit = mock.AsyncMock()
+  batch2.commit = mock.AsyncMock()
+  mock_firestore_client.batch.side_effect = [batch1, batch2]
+
+  await service.add_session_to_memory(session)
+
+  assert mock_firestore_client.batch.call_count == 2
+  assert batch1.set.call_count == 500
+  batch1.commit.assert_called_once()
+  assert batch2.set.call_count == 1
+  batch2.commit.assert_called_once()
+