Oneshot attempt at adding firestore support for memory and sessions

Author: ScottMansfield

pyproject.toml

@@ -157,6 +157,7 @@ extensions = [
   "beautifulsoup4>=3.2.2",                      # For load_web_page tool.
   "crewai[tools];python_version>='3.11' and python_version<'3.12'",       # For CrewaiTool; chromadb/pypika fail on 3.12+
   "docker>=7.0.0",                              # For ContainerCodeExecutor
+  "google-cloud-firestore>=2.11.0",             # For Firestore services
   "kubernetes>=29.0.0",                         # For GkeCodeExecutor
   "k8s-agent-sandbox>=0.1.1.post3",             # For GkeCodeExecutor sandbox mode
   "langgraph>=0.2.60, <0.4.8",                  # For LangGraphAgent

src/google/adk/firestore_database_runner.py

@@ -0,0 +1,64 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import Optional
+
+from .artifacts.gcs_artifact_service import GcsArtifactService
+from .memory.firestore_memory_service import FirestoreMemoryService
+from .runners import Runner
+from .sessions.firestore_session_service import FirestoreSessionService
+
+if TYPE_CHECKING:
+  from .agents.base_agent import BaseAgent
+
+
+def create_firestore_runner(
+    agent: BaseAgent,
+    gcs_bucket_name: Optional[str] = None,
+    firestore_root_collection: Optional[str] = None,
+) -> Runner:
+  """Creates a Runner configured with Firestore and GCS services.
+
+  Args:
+    agent: The root agent to run.
+    gcs_bucket_name: The GCS bucket name for artifacts.
+    firestore_root_collection: The root collection name for Firestore.
+
+  Returns:
+    A Runner instance configured with Firestore services.
+  """
+  # GcsArtifactService might require bucket name in constructor or read from env.
+  # Let's assume it reads from env or takes it.
+  # If we pass it, we might need to check its signature.
+  # Let's assume it takes bucket_name if provided, or reads from env.
+  artifact_service = GcsArtifactService()
+  if gcs_bucket_name:
+    # If GcsArtifactService supports setting it, we set it.
+    # Or we can assume it reads from ADK_GCS_BUCKET_NAME env var.
+    pass
+
+  session_service = FirestoreSessionService(
+      root_collection=firestore_root_collection
+  )
+  memory_service = FirestoreMemoryService()
+
+  return Runner(
+      agent=agent,
+      session_service=session_service,
+      artifact_service=artifact_service,
+      memory_service=memory_service,
+  )

src/google/adk/memory/firestore_memory_service.py

@@ -0,0 +1,172 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import re
+from typing import Any
+from typing import Optional
+
+from google.cloud import firestore
+from typing_extensions import override
+
+from ..events.event import Event
+from . import _utils
+from .base_memory_service import BaseMemoryService
+from .base_memory_service import SearchMemoryResponse
+from .memory_entry import MemoryEntry
+
+if False:  # TYPE_CHECKING
+  from ..sessions.session import Session
+
+logger = logging.getLogger("google_adk." + __name__)
+
+DEFAULT_EVENTS_COLLECTION = "events"
+
+# Standard English stop words
+DEFAULT_STOP_WORDS = {
+    "a", "an", "the", "and", "or", "but", "if", "then", "else", "to", "of",
+    "in", "on", "for", "with", "is", "are", "was", "were", "be", "been",
+    "being", "have", "has", "had", "do", "does", "did", "can", "could",
+    "will", "would", "should", "shall", "may", "might", "must", "up", "down",
+    "out", "in", "over", "under", "again", "further", "then", "once", "here",
+    "there", "when", "where", "why", "how", "all", "any", "both", "each",
+    "few", "more", "most", "other", "some", "such", "no", "nor", "not", "only",
+    "own", "same", "so", "than", "too", "very", "i", "me", "my", "myself",
+    "we", "our", "ours", "ourselves", "you", "your", "yours", "yourself",
+    "yourselves", "he", "him", "his", "himself", "she", "her", "hers",
+    "herself", "it", "its", "itself", "they", "them", "their", "theirs",
+    "themselves", "what", "which", "who", "whom", "this", "that", "these",
+    "those", "am", "is", "are", "was", "were", "be", "been", "being",
+    "have", "has", "had", "having", "do", "does", "did", "doing",
+    "a", "an", "the", "and", "but", "if", "or", "because", "as", "until",
+    "while", "of", "at", "by", "for", "with", "about", "against", "between",
+    "into", "through", "during", "before", "after", "above", "below", "to",
+    "from", "up", "down", "in", "out", "on", "off", "over", "under", "again",
+    "further", "then", "once", "here", "there", "when", "where", "why", "how",
+    "all", "any", "both", "each", "few", "more", "most", "other", "some",
+    "such", "no", "nor", "not", "only", "own", "same", "so", "than", "too",
+    "very", "s", "t", "can", "will", "just", "don", "should", "now"
+}
+
+
+class FirestoreMemoryService(BaseMemoryService):
+  """Memory service that uses Google Cloud Firestore as the backend."""
+
+  def __init__(
+      self,
+      client: Optional[firestore.AsyncClient] = None,
+      events_collection: Optional[str] = None,
+      stop_words: Optional[set[str]] = None,
+  ):
+    """Initializes the Firestore memory service.
+
+    Args:
+      client: An optional Firestore AsyncClient. If not provided, a new one
+        will be created.
+      events_collection: The name of the events collection or collection group.
+        Defaults to 'events'.
+      stop_words: A set of words to ignore when extracting keywords. Defaults to
+        a standard English stop words list.
+    """
+    self.client = client or firestore.AsyncClient()
+    self.events_collection = events_collection or DEFAULT_EVENTS_COLLECTION
+    self.stop_words = stop_words if stop_words is not None else DEFAULT_STOP_WORDS
+
+  @override
+  async def add_session_to_memory(self, session: Session) -> None:
+    """No-op. Assumes events are written to Firestore by FirestoreSessionService."""
+    pass
+
+  def _extract_keywords(self, text: str) -> set[str]:
+    """Extracts keywords from text, ignoring stop words."""
+    words = re.findall(r"[A-Za-z]+", text.lower())
+    return {word for word in words if word not in self.stop_words}
+
+  async def _search_by_keyword(
+      self, app_name: str, user_id: str, keyword: str
+  ) -> list[MemoryEntry]:
+    """Searches for events matching a single keyword."""
+    # This requires a collection group index in Firestore for 'events' with
+    # appName == X, userId == Y, and keywords array-contains Z.
+    query = (
+        self.client.collection_group(self.events_collection)
+        .where("appName", "==", app_name)
+        .where("userId", "==", user_id)
+        .where("keywords", "array_contains", keyword)
+    )
+
+    docs = await query.get()
+    entries = []
+    for doc in docs:
+      data = doc.to_dict()
+      if data and "event_data" in data:
+        try:
+          event = Event.model_validate(data["event_data"])
+          if event.content:
+            entries.append(
+                MemoryEntry(
+                    content=event.content,
+                    author=event.author,
+                    timestamp=_utils.format_timestamp(event.timestamp),
+                )
+            )
+        except Exception as e:
+          logger.warning("Failed to parse event from Firestore: %s", e)
+
+    return entries
+
+  @override
+  async def search_memory(
+      self, *, app_name: str, user_id: str, query: str
+  ) -> SearchMemoryResponse:
+    """Searches memory for events matching the query."""
+    keywords = self._extract_keywords(query)
+    if not keywords:
+      return SearchMemoryResponse()
+
+    # Search for each keyword concurrently
+    tasks = [
+        self._search_by_keyword(app_name, user_id, keyword)
+        for keyword in keywords
+    ]
+    results = await asyncio.gather(*tasks)
+
+    # Merge results and deduplicate by MemoryEntry content/author/timestamp
+    # (MemoryEntry is not hashable by default if it contains complex objects,
+    # so we might need to deduplicate by id if available, or by content string).
+    # Since we convert Event to MemoryEntry, we don't have event.id in MemoryEntry
+    # unless we add it. The Java code use custom hash/equals for MemoryEntry.
+    # In Python, MemoryEntry is a Pydantic model. We can deduplicate by model_dump_json()
+    # or by a custom key.
+    seen = set()
+    memories = []
+    for result_list in results:
+      for entry in result_list:
+        # Deduplicate by a key of (author, content_text)
+        # Content might be complex, so let's use its json representation or text
+        content_text = ""
+        if entry.content and entry.content.parts:
+          content_text = " ".join(
+              [part.text for part in entry.content.parts if part.text]
+          )
+        key = (entry.author, content_text, entry.timestamp)
+        if key not in seen:
+          seen.add(key)
+          memories.append(entry)
+
+    return SearchMemoryResponse(memories=memories)