diff --git a/pyproject.toml b/pyproject.toml index 2acaae84a81..7e750ffe9a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -374,6 +374,7 @@ optional-dependencies.test = [ "llama-index-readers-file>=0.4", "lxml>=5.3", "mcp>=1.24,<3", + "mongomock>=4.3,<5", # In-memory MongoDB used by the mongodb integration tests. "nltk!=3.10.1", # Transitive via rouge-score and llama-index-core; 3.10.1's import hook breaks any venv living inside the working directory (reverted upstream in nltk/nltk#3732). "openai>=2.20,<3", "openpyxl>=3.1.5,<4", diff --git a/src/google/adk/integrations/mongodb/README.md b/src/google/adk/integrations/mongodb/README.md index 23cb6bfbe6f..397ebb49b6b 100644 --- a/src/google/adk/integrations/mongodb/README.md +++ b/src/google/adk/integrations/mongodb/README.md @@ -24,7 +24,9 @@ and memory services work with any MongoDB deployment reachable by PyMongo. (`FeatureName.MONGODB_SESSION_SERVICE`) and emits a one-time warning on first use. - **Long-term Memory (`MongoDbMemoryService`):** Ingests session events into a - memory collection and recalls them by keyword matching. + memory collection and recalls them by keyword matching (default), by + `$vectorSearch` with Google embeddings (`retrieval_mode="vector"`), or with + Atlas Automated Embedding (`retrieval_mode="auto"`). - **Flexible Connection Options:** Pass a connection string and let the integration create/own the PyMongo client, or bring your own pre-configured `pymongo.MongoClient`. @@ -50,7 +52,7 @@ pip install "pymongo>=4.9,<5" | `mongodb_vector_search` | A **vector search index** on the collection, and a deployment that supports `$vectorSearch` (MongoDB Atlas 6.0.11+/7.0.2+, or self-managed MongoDB 8.2+). | | `mongodb_hybrid_search` | A **vector search index** and a **full-text search index** on the collection, and a deployment that supports `$rankFusion` (MongoDB Atlas 8.0+). | | `MongoDbSessionService` | Any MongoDB deployment reachable by PyMongo (no search indexes needed). | -| `MongoDbMemoryService` | Any MongoDB deployment reachable by PyMongo (no search indexes needed). | +| `MongoDbMemoryService` | Keyword mode (default): any MongoDB deployment reachable by PyMongo (no search indexes needed). `"vector"` / `"auto"` modes: a deployment with vector search (Atlas or self-managed 8.0+). | ## Quick Start: Search Tools @@ -279,40 +281,74 @@ runner = Runner( ) ``` -Events ingested via `add_session_to_memory` or `add_events_to_memory` are -written in a single unordered `bulk_write` batch as memory documents (one per -event) keyed by a compound `_id` of `(app_name, user_id, session_id, -event_id)`, each carrying the lowercase keywords extracted from its text -content (a compact English stop-word list is ignored; pass `stop_words` to -customize). `search_memory` returns up to `max_results` (default `50`) entries -whose keyword array intersects the query's keywords, ranked by keyword overlap -with `timestamp` descending as a tie-breaker. Ingestion is idempotent — -re-adding events or a session overwrites existing memories rather than -duplicating them. +Events ingested via `add_session_to_memory` are written in a single +unordered `bulk_write` batch as memory documents (one per event) keyed by a +compound `_id` of `(app_name, user_id, session_id, event_id)`. Ingestion is +idempotent — re-adding a session overwrites its memories rather than +duplicating them — and every search returns at most `max_results` (default +20) ranked entries, so a query carrying a common term returns the best few +memories rather than everything it touched. -Retrieval is intentionally keyword-only: it needs no embedding model, no -Google credentials, and no vector index, so it works on any MongoDB -deployment (not just Atlas) at zero per-query cost. Vector recall is a -possible follow-up that can build on this integration's own search tools. +`retrieval_mode` selects how `search_memory` recalls memories: -### Memory secondary index (`ensure_indexes`) +| Mode | How it works | Requirements | +| :--- | :--- | :--- | +| `"keyword"` (default) | Documents carry the lowercase keywords extracted from their text (a compact English stop-word list is ignored; pass `stop_words` to customize); recall matches the query's keywords, ranked by overlap with `timestamp` descending as a tie-breaker. | Any MongoDB deployment — no embedding model, credentials, or search index. | +| `"vector"` | Texts are embedded at ingest with a Google embedding model (`embedding_model_name`, default `text-embedding-005`) through the genai client; recall embeds the query and runs `$vectorSearch`. | A deployment with vector search (Atlas or self-managed 8.0+), Google credentials, and a vector index on the `embedding` field. | +| `"auto"` | Documents carry plain text only; an `autoEmbed` index over the `text` field makes Atlas embed documents on write and queries on read with a Voyage AI model (Atlas Automated Embedding, Preview). | Atlas, and an `autoEmbed` vector index. Freshly ingested memories become recallable once Atlas embeds them (typically seconds). | + +```python +# Default: keyword recall, no embeddings anywhere. +memory_service = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", +) + +# Vector recall with Google embeddings (needs Google credentials). +vector_memory = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + retrieval_mode="vector", +) + +# Atlas Automated Embedding: Atlas embeds memories and queries for you. +auto_memory = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + retrieval_mode="auto", +) +``` + +Keyword is the default because it is dependency-free and works everywhere; +pick a vector mode when semantic similarity matters more than synonym-blind +exactness, or when no Google credentials are available but Atlas is +(`"auto"`). The retrieval mode must match how memories were ingested: +memories ingested in one mode are invisible to the others (keyword docs +carry no embedding, vector/auto docs carry no keywords). + +### Memory secondary indexes (`ensure_indexes`) Writes are keyed by `_id`, which MongoDB indexes implicitly. However, -`search_memory` reads `memories` by `(app_name, user_id)` with a `$in` match on -the `keywords` array, which becomes a collection scan as data grows. +`search_memory` reads `memories` by `(app_name, user_id)` — with a `$in` +match on the `keywords` array in keyword mode, or a `$vectorSearch` pre-filter +in vector/auto modes — which becomes a collection scan as data grows. -The service therefore creates the recommended secondary index on first use by -default (`ensure_indexes=True`): +With `ensure_indexes=True` (the default) the service creates the index its +mode needs on first use: -| Collection | Index | Keys | +| Collection | Index | Keys / definition | | :--- | :--- | :--- | -| `memories` | `app_user_keywords` | `(app_name, user_id, keywords)` (multikey) | +| `memories` (keyword) | `app_user_keywords` | `(app_name, user_id, keywords)` (multikey) | +| `memories` (vector) | `memory_vector_index` | cosine `vector` field on `embedding` (sized from `embedding_output_dimensionality`, default 768) with `app_name`/`user_id` filter fields | +| `memories` (auto) | `memory_vector_index` | `autoEmbed` field over `text` (model from `auto_embedding_model_name`) with `app_name`/`user_id` filter fields | +| `memories` (all modes) | `memory_ttl` | `expire_at` (TTL) — created when `memory_ttl_days` is set (default 7) | Index creation is idempotent — re-running with the same spec is a no-op. When the runtime user lacks the `createIndex` privilege (see [Required permissions](#required-permissions-least-privilege-setup)), the service logs one warning and falls back to unindexed reads rather than -failing; other index-creation errors still propagate. Pass +failing; other index-creation errors still propagate. Vector search indexes +build asynchronously — recall works once the index reports READY. Pass `ensure_indexes=False` to skip the attempt entirely (and silence the warning) on a deliberately least-privileged user: @@ -324,6 +360,32 @@ memory_service = MongoDbMemoryService( ) ``` +### Memory retention (`memory_ttl_days`) + +Memories are retained for `memory_ttl_days` (default **7**) measured from +ingest: each memory document stores an `expire_at` date that many days in the +future, and the `memory_ttl` TTL index (created when `ensure_indexes=True`) +lets MongoDB's TTL monitor delete a memory once its `expire_at` passes. +Re-ingesting a session refreshes `expire_at`. Pass `memory_ttl_days=None` to +keep memories indefinitely (no `expire_at`, no TTL index). Deletion runs in +the background, so expired memories may linger briefly past their `expire_at`. + +```python +# Keep memories for 30 days. +memory_service = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + memory_ttl_days=30, +) + +# Keep memories indefinitely. +memory_service = MongoDbMemoryService( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="my_app", + memory_ttl_days=None, +) +``` + ## Required permissions (least-privilege setup) Do not hand the integration an admin user. The components need very little, @@ -333,7 +395,7 @@ and separate database users keep the privileges separated: | :--- | :--- | :--- | | **Search user** (read-only) | `MongoDbToolset` | `find` on each collection the agent may search. | | **State user** (read-write) | `MongoDbSessionService` | `find`, `insert`, `update`, `remove` on `sessions` and `events` (session deletion); `find`, `insert`, `update` on `app_states` and `user_states` (upserts only, never deletes). Add `createIndex` on `sessions` and `events` so `ensure_indexes` (on by default) can work; without it the service logs one warning and falls back to unindexed reads, or pass `ensure_indexes=False` to opt out. | -| **Memory user** (read-write) | `MongoDbMemoryService` | `find`, `insert`, `update`, `createIndex` on `memories` (upserts only, never deletes). `createIndex` lets `ensure_indexes` (on by default) build the recall index; without it the service logs one warning and falls back to unindexed reads, or drop it and pass `ensure_indexes=False` to opt out. | +| **Memory user** (read-write) | `MongoDbMemoryService` | `find`, `insert`, `update`, `createIndex` on `memories` (upserts only, never deletes). `createIndex` lets `ensure_indexes` (on by default) build the recall and TTL indexes; without it the service logs one warning and falls back to unindexed reads, or drop it and pass `ensure_indexes=False` to opt out. | ```javascript // mongosh, as an admin user: one read-only role for the search tools, and @@ -374,7 +436,7 @@ Creating and dropping **search indexes** is an administrative action, not something the tools do at run time — run setup with a separate user holding the `createSearchIndexes`, `listSearchIndexes`, and `dropSearchIndexes` actions (on Atlas, the built-in `Project Data Access Admin` role covers -this), and keep those privileges off the runtime user above. +this), and keep those privileges off the runtime users above. ### The model chooses the collection @@ -456,6 +518,10 @@ remote = agent_engines.create( ) ``` +The MongoDB session and memory services cannot be deployed to Agent Engine +yet — wiring them in (including the same pickling and secret handling) is +planned as a follow-up. + ## Configuration: `MongoDbToolSettings` `MongoDbToolSettings` customizes the defaults used by the search tools: diff --git a/src/google/adk/integrations/mongodb/_client.py b/src/google/adk/integrations/mongodb/_client.py index 5a27dfdb2e6..3149f9ee788 100644 --- a/src/google/adk/integrations/mongodb/_client.py +++ b/src/google/adk/integrations/mongodb/_client.py @@ -63,6 +63,25 @@ def get_mongo_client( ) +def connection_string_from_env(env_var: str) -> str: + """Reads the MongoDB connection string from an environment variable. + + Used with `connection_string_env_var` so the value is resolved where the + client is built — on the runtime, not at packaging time — which keeps the + database password out of pickled deployment packages (Agent Engine). + + Raises: + ValueError: If the environment variable is not set. + """ + value = os.environ.get(env_var) + if not value: + raise ValueError( + f"Environment variable {env_var} is not set: it must hold the" + " MongoDB connection string." + ) + return value + + def is_index_privilege_error(exc: Exception) -> bool: """Returns whether `exc` is the deployment denying a `createIndex` call. diff --git a/src/google/adk/integrations/mongodb/_memory_service.py b/src/google/adk/integrations/mongodb/_memory_service.py index a4cf65f8ad2..cad9a863929 100644 --- a/src/google/adk/integrations/mongodb/_memory_service.py +++ b/src/google/adk/integrations/mongodb/_memory_service.py @@ -19,9 +19,13 @@ import asyncio from collections.abc import Mapping from collections.abc import Sequence +from datetime import datetime +from datetime import timedelta +from datetime import timezone import itertools import logging import re +import threading from typing import Any from typing import TYPE_CHECKING import unicodedata @@ -31,6 +35,7 @@ from typing_extensions import override from . import _client +from . import _search_tool from ...features import experimental from ...features import FeatureName from ...memory import _utils @@ -39,6 +44,7 @@ from ...memory.memory_entry import MemoryEntry if TYPE_CHECKING: + from google.genai import Client from pymongo import MongoClient from pymongo import ReplaceOne from pymongo.collection import Collection @@ -70,13 +76,39 @@ def __init__( _DEFAULT_MAX_RESULTS = 50 _UNKNOWN_SESSION_ID = "__unknown_session_id__" +# Retrieval modes for the `retrieval_mode` constructor argument. +KEYWORD_RETRIEVAL = "keyword" +VECTOR_RETRIEVAL = "vector" +AUTO_EMBEDDING_RETRIEVAL = "auto" +_RETRIEVAL_MODES = frozenset( + {KEYWORD_RETRIEVAL, VECTOR_RETRIEVAL, AUTO_EMBEDDING_RETRIEVAL} +) + +_DEFAULT_EMBEDDING_MODEL_NAME = "text-embedding-005" +_DEFAULT_EMBEDDING_DIMENSIONS = 768 # native size of text-embedding-005 +_DEFAULT_AUTO_EMBEDDING_MODEL = "voyage-4" +_DEFAULT_MEMORY_VECTOR_INDEX_NAME = "memory_vector_index" +_DEFAULT_NUM_CANDIDATES = 100 + +# Default retention for stored memories, in days. Each memory document carries +# an `expire_at` date this many days past ingest, and the `memory_ttl` TTL +# index (created when `ensure_indexes=True`) lets MongoDB's TTL monitor delete +# documents once `expire_at` passes. Pass `memory_ttl_days=None` to keep +# memories indefinitely. +_DEFAULT_MEMORY_TTL_DAYS = 7 + # Recommended secondary index, created on first use unless -# `ensure_indexes=False`. It covers the `search_memory` lookup (app/user -# equality plus a multikey `$in` on keywords); writes are keyed by `_id`, which -# MongoDB indexes implicitly. +# `ensure_indexes=False`. It covers the keyword-mode `search_memory` lookup +# (app/user equality plus a multikey `$in` on keywords); writes are keyed by +# `_id`, which MongoDB indexes implicitly. _MEMORIES_RECALL_INDEX_NAME = "app_user_keywords" _MEMORIES_RECALL_INDEX_KEYS = [("app_name", 1), ("user_id", 1), ("keywords", 1)] +# TTL index on the `expire_at` date field: MongoDB's TTL monitor deletes a +# memory document once its `expire_at` passes. Created for every retrieval +# mode when `ensure_indexes=True` and a TTL is configured. +_MEMORIES_TTL_INDEX_NAME = "memory_ttl" + # Compact English stop-word list ignored when extracting keywords. Kept # intentionally short; callers can pass their own set via `stop_words`. DEFAULT_STOP_WORDS = frozenset({ @@ -138,9 +170,23 @@ class MongoDbMemoryService(BaseMemoryService): """Memory service that uses MongoDB as the backend. Events ingested from sessions are stored as memory documents (one per - event) in a single collection, each carrying the lowercase keywords - extracted from its text content. `search_memory` matches documents whose - keyword array intersects the query's keywords, ranked by keyword overlap. + event) in a single collection. `retrieval_mode` selects how + `search_memory` recalls them: + + - `"keyword"` (default): each document carries the lowercase keywords + extracted from its text; recall matches documents whose keyword array + intersects the query's keywords, ranked by overlap (most recent first + within a tie) and capped at `max_results`. Works on any MongoDB + deployment — no embedding model, credentials, or search index needed. + - `"vector"`: texts are embedded at ingest with the Google embedding + model (`embedding_model_name`) through the genai client, and recall + runs `$vectorSearch` with the embedded query. Requires a deployment + with vector search (Atlas or self-managed 8.0+) and Google credentials. + - `"auto"`: Atlas Automated Embedding (Preview). Documents carry plain + text only; an `autoEmbed` vector index over the `text` field makes + Atlas embed documents on write and queries on read with a Voyage AI + model, so no Google credentials are needed. Freshly ingested memories + become recallable once Atlas has embedded them (typically seconds). Example: ```python @@ -167,6 +213,13 @@ def __init__( stop_words: set[str] | None = None, max_results: int = _DEFAULT_MAX_RESULTS, ensure_indexes: bool = True, + retrieval_mode: str = KEYWORD_RETRIEVAL, + genai_client: Client | None = None, + embedding_model_name: str = _DEFAULT_EMBEDDING_MODEL_NAME, + embedding_output_dimensionality: int | None = None, + auto_embedding_model_name: str = _DEFAULT_AUTO_EMBEDDING_MODEL, + vector_index_name: str = _DEFAULT_MEMORY_VECTOR_INDEX_NAME, + memory_ttl_days: int | None = _DEFAULT_MEMORY_TTL_DAYS, ) -> None: """Initializes the MongoDB memory service. @@ -178,17 +231,57 @@ def __init__( mongo_client: An existing PyMongo client to use instead of creating one from `connection_string`. The caller keeps ownership of the client. memories_collection: Collection name for memory documents. - stop_words: Words to ignore when extracting keywords. Defaults to a - standard English stop-word list. - max_results: Maximum number of memories returned per query. - ensure_indexes: Create the recommended secondary index on first use: - `app_user_keywords` on the memories collection (the app/user/keywords - lookup behind `search_memory`). Index creation is idempotent. Without - the `createIndex` privilege, a warning is logged once and reads fall - back to unindexed scans; pass False to skip the attempt entirely. + stop_words: Words to ignore when extracting keywords (keyword mode). + Defaults to a standard English stop-word list. + max_results: Maximum number of memories returned per query (keyword + mode: ranked by overlap and recency; vector modes: nearest + neighbors). + ensure_indexes: Create the recommended indexes for the retrieval mode on + first use (default True): the `app_user_keywords` multikey index in + keyword mode or the `memory_vector_index` vector search index in + vector/auto modes, plus the `memory_ttl` TTL index when + `memory_ttl_days` is set. Index creation is idempotent. Without the + `createIndex`/`createSearchIndexes` privilege a warning is logged once + and reads fall back to unindexed behavior; pass False to skip the + attempt entirely. Vector search indexes build asynchronously — recall + works once the index reports READY. + retrieval_mode: "keyword" (default), "vector", or "auto" — see the + class docstring. The mode must match how the stored memories were + ingested: memories ingested in one mode are invisible to the others. + genai_client: The `google.genai` client used to embed memory texts and + queries in vector mode. Defaults to a client built from the ambient + environment. Unused in keyword and auto modes. + embedding_model_name: Google embedding model for vector mode. Must be + the model the stored embeddings were produced with. + embedding_output_dimensionality: Optional shortened embedding length; + must match the vector index's `numDimensions` when set. + auto_embedding_model_name: Voyage AI model the auto mode's `autoEmbed` + index embeds with. + vector_index_name: Name of the vector search index used in vector/auto + modes. + memory_ttl_days: Number of days a memory is retained before MongoDB's + TTL monitor deletes it, measured from ingest (default 7). Each memory + document stores an `expire_at` date this many days in the future, and + when `ensure_indexes=True` a `memory_ttl` TTL index is created on that + field. Pass None to keep memories indefinitely (no `expire_at`, no TTL + index). Deletion runs in the background, so expired memories may + linger briefly past their `expire_at`. + + Raises: + ValueError: If `max_results` or `memory_ttl_days` is not positive, or + `retrieval_mode` is unknown. """ if max_results <= 0: raise ValueError("max_results must be positive.") + if memory_ttl_days is not None and memory_ttl_days < 1: + raise ValueError( + "memory_ttl_days must be a positive integer or None (no expiry)." + ) + if retrieval_mode not in _RETRIEVAL_MODES: + raise ValueError( + f"retrieval_mode must be one of {sorted(_RETRIEVAL_MODES)}, got" + f" {retrieval_mode!r}." + ) if mongo_client is not None and connection_string is not None: raise ValueError( "Only one of `connection_string` and `mongo_client` may be provided." @@ -213,6 +306,17 @@ def __init__( self._max_results = max_results self._ensure_indexes = ensure_indexes self._indexes_ensured = False + self._retrieval_mode = retrieval_mode + self._genai_client = genai_client + self._embedding_model_name = embedding_model_name + self._embedding_output_dimensionality = embedding_output_dimensionality + self._auto_embedding_model_name = auto_embedding_model_name + self._vector_index_name = vector_index_name + self._memory_ttl_days = memory_ttl_days + # Serializes first-use index creation across the worker threads that + # `asyncio.to_thread` runs `_ensure_indexes_once` on (vector index + # creation is not concurrency-safe on a cold collection). + self._index_lock = threading.Lock() def _memories(self) -> Collection[dict[str, Any]]: return self._client[self._database_name][self._memories_collection] @@ -229,7 +333,7 @@ def _memory_key( } def _ensure_indexes_once(self) -> None: - """Creates the recommended secondary index on first use, if enabled. + """Creates the recommended secondary indexes on first use, if enabled. Index creation is idempotent on the server, and the guard flag only skips work in this process. A missing `createIndex` privilege logs one @@ -238,16 +342,37 @@ def _ensure_indexes_once(self) -> None: """ if not self._ensure_indexes or self._indexes_ensured: return - try: - self._memories().create_index( - _MEMORIES_RECALL_INDEX_KEYS, name=_MEMORIES_RECALL_INDEX_NAME - ) + with self._index_lock: + # Re-check under the lock: a concurrent first call may have done it. + if self._indexes_ensured: + return + if self._retrieval_mode == KEYWORD_RETRIEVAL: + self._create_index( + lambda: self._memories().create_index( + _MEMORIES_RECALL_INDEX_KEYS, name=_MEMORIES_RECALL_INDEX_NAME + ) + ) + else: + self._create_index(self._create_vector_index) + self._create_index(self._create_ttl_index) + self._indexes_ensured = True + + def _create_index(self, create) -> None: + """Runs one index-creation call, tolerating per-index conflicts. + + Code 85 (IndexOptionsConflict) means an equivalent index already exists + under another name/options, so reads are still indexed; a missing + `createIndex` privilege warns once and falls back to unindexed reads. + Either way the remaining indexes are still attempted. Any other failure + propagates so the next operation retries (the guard flag stays unset). + """ + if create is None: + return # pymongo is an optional import, so the error is matched by shape. + try: + create() except Exception as exc: # pylint: disable=broad-exception-caught - # Code 85 (IndexOptionsConflict): an index on these keys already exists - # under another name or with other options, so reads are still indexed. if getattr(exc, "code", None) == 85: - self._indexes_ensured = True return if not _client.is_index_privilege_error(exc): raise @@ -258,7 +383,69 @@ def _ensure_indexes_once(self) -> None: " grant the privilege. Details: %s", exc, ) - self._indexes_ensured = True + + def _create_ttl_index(self) -> None: + """Creates the TTL index that expires memories after `memory_ttl_days`. + + The index targets the `expire_at` date field written at ingest. With + `expireAfterSeconds=0` MongoDB's TTL monitor deletes a document as soon + as its `expire_at` passes, so the per-document date carries the retention + (re-ingesting events refreshes it). No-op when no TTL is configured. + """ + if self._memory_ttl_days is None: + return + self._memories().create_index( + [("expire_at", 1)], + name=_MEMORIES_TTL_INDEX_NAME, + expireAfterSeconds=0, + ) + + def _create_vector_index(self) -> None: + """Creates the vector search index for vector/auto retrieval modes.""" + # pymongo is imported lazily here: it stays an optional dependency until + # a MongoDB-backed service is actually constructed. + from pymongo.operations import SearchIndexModel # pylint: disable=import-outside-toplevel + + # createSearchIndexes rejects a collection that does not exist yet, which + # is exactly the state of a fresh deployment on its first ingest. + database = self._client[self._database_name] + if self._memories_collection not in database.list_collection_names(): + try: + database.create_collection(self._memories_collection) + except Exception as exc: + # CollectionInvalid (code 48): another process created it first. + if getattr(exc, "code", None) != 48: + raise + + fields: list[dict[str, Any]] = [] + if self._retrieval_mode == VECTOR_RETRIEVAL: + fields.append({ + "type": "vector", + "path": "embedding", + "numDimensions": ( + self._embedding_output_dimensionality + or _DEFAULT_EMBEDDING_DIMENSIONS + ), + "similarity": "cosine", + }) + else: # auto mode: Atlas embeds the plain `text` field server-side. + fields.append({ + "type": "autoEmbed", + "modality": "text", + "path": "text", + "model": self._auto_embedding_model_name, + }) + fields += [ + {"type": "filter", "path": "app_name"}, + {"type": "filter", "path": "user_id"}, + ] + self._memories().create_search_index( + model=SearchIndexModel( + definition={"fields": fields}, + name=self._vector_index_name, + type="vectorSearch", + ) + ) def _extract_keywords(self, text: str) -> set[str]: """Extracts lowercase keywords from text, ignoring stop words. @@ -280,6 +467,32 @@ def _extract_keywords(self, text: str) -> set[str]: ) return words - self._stop_words + async def _embed_texts(self, texts: list[str]) -> list[list[float]]: + """Embeds memory texts in one batch call with the configured model.""" + # google.genai is imported lazily here: it stays an optional dependency + # until vector mode is actually used. + from google.genai import Client # pylint: disable=import-outside-toplevel + from google.genai import types as genai_types # pylint: disable=import-outside-toplevel + + config = genai_types.EmbedContentConfig() + if self._embedding_output_dimensionality: + config.output_dimensionality = self._embedding_output_dimensionality + try: + response = await ( + self._genai_client or Client() + ).aio.models.embed_content( + model=self._embedding_model_name, + contents=list(texts), + config=config, + ) + except Exception as ex: + raise RuntimeError(f"Failed to embed memory texts: {ex!r}") from ex + if not response.embeddings: + raise ValueError( + f"Model {self._embedding_model_name} returned no embeddings." + ) + return [[float(v) for v in e.values] for e in response.embeddings] + @override async def add_session_to_memory(self, session: Session) -> None: """Ingests the session's text events into the memory collection.""" @@ -304,11 +517,15 @@ async def add_events_to_memory( Ingestion is idempotent: memory documents are keyed by a compound `_id` of `(app_name, user_id, session_id, event_id)`, so re-adding events - overwrites their memories rather than duplicating them. + overwrites their memories rather than duplicating them. In vector mode + the texts are embedded with one batch call before the write; in auto mode + plain text is stored and Atlas embeds it when the autoEmbed index + processes the document. """ del custom_metadata # Not supported by MongoDB memory service. scoped_session_id = session_id or _UNKNOWN_SESSION_ID - operations: list[ReplaceOne[dict[str, Any]]] = [] + + pairs: list[tuple[Event, str, list[Any]]] = [] for event in events: if not event.content or not event.content.parts: continue @@ -316,9 +533,46 @@ async def add_events_to_memory( text = " ".join([part.text for part in visible_parts if part.text]) if not text: continue - keywords = self._extract_keywords(text) - if not keywords: - continue + pairs.append((event, text, visible_parts)) + if not pairs: + return + + embeddings: list[list[float]] | None = None + if self._retrieval_mode == VECTOR_RETRIEVAL: + embeddings = await self._embed_texts([text for _, text, _ in pairs]) + + # One expiry for the whole batch, so every memory ingested together ages + # out together. Re-ingesting events refreshes expire_at (upsert). + expire_at = None + if self._memory_ttl_days is not None: + expire_at = datetime.now(timezone.utc) + timedelta( + days=self._memory_ttl_days + ) + + operations: list[ReplaceOne[dict[str, Any]]] = [] + for i, (event, text, visible_parts) in enumerate(pairs): + document: dict[str, Any] = { + "app_name": app_name, + "user_id": user_id, + "session_id": scoped_session_id, + "author": event.author, + "text": text, + "content": ( + event.content.model_copy( + update={"parts": visible_parts} + ).model_dump(exclude_none=True, mode="json") + ), + "timestamp": event.timestamp, + } + if expire_at is not None: + document["expire_at"] = expire_at + if self._retrieval_mode == KEYWORD_RETRIEVAL: + keywords = self._extract_keywords(text) + if not keywords: + continue + document["keywords"] = sorted(keywords) + elif embeddings is not None: + document["embedding"] = embeddings[i] operations.append( ReplaceOne( { @@ -326,19 +580,7 @@ async def add_events_to_memory( app_name, user_id, scoped_session_id, event.id ) }, - { - "app_name": app_name, - "user_id": user_id, - "session_id": scoped_session_id, - "author": event.author, - "keywords": sorted(keywords), - "content": ( - event.content.model_copy( - update={"parts": visible_parts} - ).model_dump(exclude_none=True, mode="json") - ), - "timestamp": event.timestamp, - }, + document, upsert=True, ) ) @@ -354,34 +596,37 @@ def _write() -> None: async def search_memory( self, *, app_name: str, user_id: str, query: str ) -> SearchMemoryResponse: - """Searches memory for events matching the query's keywords.""" - keywords = self._extract_keywords(query) - if not keywords: - return SearchMemoryResponse() + """Searches memory for events matching the query. - query_keywords = sorted(keywords) - pipeline: list[dict[str, Any]] = [ - { - "$match": { - "app_name": app_name, - "user_id": user_id, - "keywords": {"$in": query_keywords}, - } - }, - { - "$addFields": { - "_score": { - "$size": {"$setIntersection": ["$keywords", query_keywords]} - } - } - }, - {"$sort": {"_score": -1, "timestamp": -1}}, - {"$limit": self._max_results}, - ] + Keyword mode matches the query's keywords; vector mode embeds the query + and runs `$vectorSearch`; auto mode sends the query text to the autoEmbed + index. All modes cap results at `max_results`. + """ + query_vector: list[float] | None = None + if self._retrieval_mode == VECTOR_RETRIEVAL: + # Embed in async context; the genai client's async transport needs a + # running loop, which a worker thread would not have. + query_vector = await _search_tool._embed_query( + query=query, + model_name=self._embedding_model_name, + output_dimensionality=self._embedding_output_dimensionality, + genai_client=self._genai_client, + ) + + # Keyword mode: a query with no usable keywords matches nothing, so skip + # the aggregate — and index creation — entirely. + query_keywords: list[str] | None = None + if self._retrieval_mode == KEYWORD_RETRIEVAL: + keywords = self._extract_keywords(query) + if not keywords: + return SearchMemoryResponse() + query_keywords = sorted(keywords) def _search() -> list[dict[str, Any]]: self._ensure_indexes_once() - return list(self._memories().aggregate(pipeline)) + if self._retrieval_mode == KEYWORD_RETRIEVAL: + return self._keyword_search(app_name, user_id, query_keywords or []) + return self._vector_search_docs(app_name, user_id, query, query_vector) docs = await asyncio.to_thread(_search) @@ -418,6 +663,59 @@ def _search() -> list[dict[str, Any]]: return SearchMemoryResponse(memories=memories) + def _keyword_search( + self, app_name: str, user_id: str, query_keywords: list[str] + ) -> list[dict[str, Any]]: + """Keyword recall: overlap-ranked, capped — all server-side.""" + pipeline: list[dict[str, Any]] = [ + { + "$match": { + "app_name": app_name, + "user_id": user_id, + "keywords": {"$in": query_keywords}, + } + }, + { + "$addFields": { + "_score": { + "$size": {"$setIntersection": ["$keywords", query_keywords]} + } + } + }, + {"$sort": {"_score": -1, "timestamp": -1}}, + {"$limit": self._max_results}, + ] + return list(self._memories().aggregate(pipeline)) + + def _vector_search_docs( + self, + app_name: str, + user_id: str, + query: str, + query_vector: list[float] | None, + ) -> list[dict[str, Any]]: + """Vector recall: `$vectorSearch` with a pre-filter on app and user.""" + stage: dict[str, Any] = { + "index": self._vector_index_name, + "numCandidates": max(self._max_results * 10, _DEFAULT_NUM_CANDIDATES), + "limit": self._max_results, + "filter": {"app_name": app_name, "user_id": user_id}, + } + if self._retrieval_mode == AUTO_EMBEDDING_RETRIEVAL: + # Atlas embeds the query text with the autoEmbed index's model. + stage["path"] = "text" + stage["query"] = {"text": query} + else: + stage["path"] = "embedding" + stage["queryVector"] = query_vector + + return list( + self._memories().aggregate([ + {"$vectorSearch": stage}, + {"$project": {"embedding": 0}}, + ]) + ) + async def close(self) -> None: """Closes the MongoDB client if it was created by this service.""" if self._owns_client: diff --git a/tests/unittests/integrations/mongodb/test_memory_service.py b/tests/unittests/integrations/mongodb/test_memory_service.py index 19e7286e513..99414cec7f0 100644 --- a/tests/unittests/integrations/mongodb/test_memory_service.py +++ b/tests/unittests/integrations/mongodb/test_memory_service.py @@ -17,6 +17,9 @@ from __future__ import annotations import copy +from datetime import datetime +from datetime import timedelta +from datetime import timezone from typing import Any from unittest import mock @@ -50,17 +53,31 @@ def __init__(self) -> None: self.create_index_error: Exception | None = None self.bulk_write_calls: list[tuple[list[Any], bool]] = [] self.last_pipeline: list[dict[str, Any]] | None = None + self.search_index_models: list[Any] = [] def create_index( - self, keys: list[tuple[str, int]], *, name: str | None = None + self, + keys: list[tuple[str, int]], + *, + name: str | None = None, + expireAfterSeconds: int | None = None, ) -> str: self.create_index_calls += 1 if self.create_index_error is not None: raise self.create_index_error index_name = name or "_".join(f"{k}_{v}" for k, v in keys) - self._indexes[index_name] = {"key": list(keys)} + spec: dict[str, Any] = {"key": list(keys)} + if expireAfterSeconds is not None: + spec["expireAfterSeconds"] = expireAfterSeconds + self._indexes[index_name] = spec return index_name + def create_search_index(self, *, model: Any) -> str: + if self.create_index_error is not None: + raise self.create_index_error + self.search_index_models.append(model) + return model.document["name"] + def index_information(self) -> dict[str, dict[str, Any]]: return copy.deepcopy(self._indexes) @@ -110,6 +127,16 @@ def aggregate(self, pipeline: list[dict[str, Any]]) -> list[dict[str, Any]]: ) elif "$limit" in stage: docs = docs[: stage["$limit"]] + elif "$vectorSearch" in stage: + spec = stage["$vectorSearch"] + if "filter" in spec: + docs = [doc for doc in docs if _matches(doc, spec["filter"])] + docs = docs[: spec.get("limit", len(docs))] + elif "$project" in stage: + excluded = {f for f, v in stage["$project"].items() if v == 0} + docs = [ + {k: v for k, v in doc.items() if k not in excluded} for doc in docs + ] else: raise NotImplementedError(f"Unsupported stage: {stage}") return docs @@ -158,6 +185,16 @@ def __getitem__(self, name: str) -> _FakeCollection: self._collections[name] = _FakeCollection() return self._collections[name] + def list_collection_names(self) -> list[str]: + return list(self._collections) + + def create_collection(self, name: str) -> _FakeCollection: + if name in self._collections: + error = _OperationFailure("collection already exists", code=48) + raise error + self._collections[name] = _FakeCollection() + return self._collections[name] + class _FakeMongoClient: """In-memory stand-in for pymongo.MongoClient.""" @@ -207,7 +244,7 @@ def __init__(self, message: str, code: int | None = None) -> None: async def test_recall_index_created_once_by_default() -> None: - """The app/user/keywords index is created once without opting in.""" + """The recall and TTL indexes are created once without opting in.""" client: Any = _FakeMongoClient() svc = MongoDbMemoryService(mongo_client=client, database_name="test_db") await svc.add_session_to_memory(_session_with_texts("hello world")) @@ -216,7 +253,10 @@ async def test_recall_index_created_once_by_default() -> None: collection = client["test_db"]["memories"] indexes = collection.index_information() assert "app_user_keywords" in indexes - assert collection.create_index_calls == 1 + assert "memory_ttl" in indexes + assert indexes["memory_ttl"]["expireAfterSeconds"] == 0 + # Two create_index calls (recall + TTL), and first-use guard keeps it at 2. + assert collection.create_index_calls == 2 async def test_search_without_keywords_skips_index_creation() -> None: @@ -267,9 +307,10 @@ async def test_index_privilege_error_warns_once_and_falls_back( ) assert len(response.memories) == 1 - assert collection.create_index_calls == 1 + # Recall and TTL index are each attempted once; each denied call warns once. + assert collection.create_index_calls == 2 warnings = [r for r in caplog.records if "createIndex" in r.getMessage()] - assert len(warnings) == 1 + assert len(warnings) == 2 async def test_existing_index_under_other_name_is_accepted( @@ -289,7 +330,8 @@ async def test_existing_index_under_other_name_is_accepted( ) assert len(response.memories) == 1 - assert collection.create_index_calls == 1 + # Both create_index calls (recall + TTL) hit code 85 and are tolerated. + assert collection.create_index_calls == 2 assert not [r for r in caplog.records if "createIndex" in r.getMessage()] @@ -302,11 +344,14 @@ async def test_index_creation_errors_other_than_privilege_propagate() -> None: with pytest.raises(_OperationFailure, match="connection dropped"): await svc.add_session_to_memory(_session_with_texts("hello world")) - # The failure is not cached: the next operation retries index creation. + # The failure is not cached: the next operation retries index creation, + # creating the recall and TTL indexes (the failed first call plus two + # successful ones). collection.create_index_error = None await svc.add_session_to_memory(_session_with_texts("hello world")) - assert collection.create_index_calls == 2 + assert collection.create_index_calls == 3 assert "app_user_keywords" in collection.index_information() + assert "memory_ttl" in collection.index_information() async def test_add_session_then_search_by_keyword( @@ -828,6 +873,269 @@ async def test_close_only_closes_owned_client() -> None: assert owned_client.closed +def _memory_doc(text: str, doc_id: Any = None) -> dict[str, Any]: + """Builds a stored memory document as the service would write it.""" + if doc_id is None: + doc_id = { + "app_name": APP, + "user_id": USER, + "session_id": "s1", + "event_id": "e0", + } + return { + "_id": doc_id, + "app_name": APP, + "user_id": USER, + "author": "user", + "text": text, + "content": ( + types.Content(role="user", parts=[types.Part(text=text)]).model_dump( + exclude_none=True, mode="json" + ) + ), + "timestamp": 1.0, + } + + +def _vector_service(**kwargs: Any) -> MongoDbMemoryService: + defaults: dict[str, Any] = dict( + database_name="test_db", + ensure_indexes=False, + retrieval_mode="vector", + genai_client=mock.MagicMock(), + ) + defaults.update(kwargs) + return MongoDbMemoryService(**defaults) + + +async def test_vector_mode_ingest_embeds_in_one_batch() -> None: + """Vector mode embeds all event texts in one call and stores the vectors.""" + client: Any = _FakeMongoClient() + genai_client = mock.MagicMock() + embedding = mock.MagicMock() + embedding.values = [0.1, 0.2] + genai_client.aio.models.embed_content = mock.AsyncMock( + return_value=mock.MagicMock(embeddings=[embedding, embedding]) + ) + svc = _vector_service(mongo_client=client, genai_client=genai_client) + + await svc.add_session_to_memory( + _session_with_texts("first memory", "second memory") + ) + + embed_call = genai_client.aio.models.embed_content.await_args + assert embed_call.kwargs["contents"] == ["first memory", "second memory"] + assert embed_call.kwargs["model"] == "text-embedding-005" + docs = list(client["test_db"]["memories"]._docs.values()) + assert len(docs) == 2 + assert all(doc["embedding"] == [0.1, 0.2] for doc in docs) + assert all("text" in doc and "keywords" not in doc for doc in docs) + + +async def test_vector_mode_recall_runs_vector_search( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Vector mode recalls via $vectorSearch with an embedded query.""" + embed_query = mock.AsyncMock(return_value=[0.1, 0.2]) + monkeypatch.setattr( + "google.adk.integrations.mongodb._search_tool._embed_query", embed_query + ) + client: Any = _FakeMongoClient() + collection = client["test_db"]["memories"] + collection._docs[_doc_key("m1")] = _memory_doc("nearest memory", "m1") + svc = _vector_service(mongo_client=client) + + response = await svc.search_memory( + app_name=APP, user_id=USER, query="coffee machine" + ) + + assert embed_query.await_args.kwargs["query"] == "coffee machine" + stage = collection.last_pipeline[0]["$vectorSearch"] + assert stage["queryVector"] == [0.1, 0.2] + assert stage["path"] == "embedding" + assert stage["filter"] == {"app_name": APP, "user_id": USER} + assert stage["limit"] == svc._max_results + texts = [ + " ".join(p.text for p in m.content.parts if p.text) + for m in response.memories + ] + assert texts == ["nearest memory"] + + +async def test_auto_mode_ingest_stores_plain_text_without_embeddings() -> None: + """Auto mode stores plain text; Atlas embeds it, not the genai client.""" + client: Any = _FakeMongoClient() + genai_client = mock.MagicMock() + svc = MongoDbMemoryService( + mongo_client=client, + database_name="test_db", + ensure_indexes=False, + retrieval_mode="auto", + genai_client=genai_client, + ) + + await svc.add_session_to_memory(_session_with_texts("just text")) + + doc = next(iter(client["test_db"]["memories"]._docs.values())) + assert doc["text"] == "just text" + assert "embedding" not in doc and "keywords" not in doc + genai_client.aio.models.embed_content.assert_not_called() + + +async def test_auto_mode_recall_sends_query_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Auto mode sends query.text to the autoEmbed index; no embed call.""" + embed_query = mock.AsyncMock(return_value=[0.1]) + monkeypatch.setattr( + "google.adk.integrations.mongodb._search_tool._embed_query", embed_query + ) + client: Any = _FakeMongoClient() + collection = client["test_db"]["memories"] + collection._docs[_doc_key("m1")] = _memory_doc("atlas recall", "m1") + svc = MongoDbMemoryService( + mongo_client=client, + database_name="test_db", + ensure_indexes=False, + retrieval_mode="auto", + ) + + response = await svc.search_memory( + app_name=APP, user_id=USER, query="coffee machine" + ) + + embed_query.assert_not_awaited() + stage = collection.last_pipeline[0]["$vectorSearch"] + assert stage["query"] == {"text": "coffee machine"} + assert stage["path"] == "text" + assert stage["filter"] == {"app_name": APP, "user_id": USER} + assert len(response.memories) == 1 + + +def test_vector_mode_creates_vector_index() -> None: + """ensure_indexes in vector mode creates a vectorSearch index.""" + client: Any = _FakeMongoClient() + svc = _vector_service(mongo_client=client, ensure_indexes=True) + + svc._ensure_indexes_once() + + collection = client["test_db"]["memories"] + assert len(collection.search_index_models) == 1 + model = collection.search_index_models[0] + fields = model.document["definition"]["fields"] + assert { + "type": "vector", + "path": "embedding", + "numDimensions": 768, + "similarity": "cosine", + } in fields + assert {"type": "filter", "path": "app_name"} in fields + assert model.document["name"] == "memory_vector_index" + # The TTL index is created alongside the vector index (default TTL). + assert "memory_ttl" in collection.index_information() + + +def test_auto_mode_creates_autoembed_index() -> None: + """ensure_indexes in auto mode creates an autoEmbed index on `text`.""" + client: Any = _FakeMongoClient() + svc = MongoDbMemoryService( + mongo_client=client, + database_name="test_db", + ensure_indexes=True, + retrieval_mode="auto", + ) + + svc._ensure_indexes_once() + + collection = client["test_db"]["memories"] + model = collection.search_index_models[0] + fields = model.document["definition"]["fields"] + assert { + "type": "autoEmbed", + "modality": "text", + "path": "text", + "model": "voyage-4", + } in fields + + +def test_retrieval_mode_must_be_known() -> None: + """An unknown retrieval mode is a config error.""" + with pytest.raises(ValueError, match="retrieval_mode"): + MongoDbMemoryService( + mongo_client=_FakeMongoClient(), + database_name="test_db", + retrieval_mode="semantic", + ) + + +def test_memory_ttl_days_must_be_positive_or_none() -> None: + """memory_ttl_days of zero or negative is a config error; None is valid.""" + with pytest.raises(ValueError, match="memory_ttl_days"): + MongoDbMemoryService( + mongo_client=_FakeMongoClient(), + database_name="test_db", + memory_ttl_days=0, + ) + # None disables expiry and is accepted. + MongoDbMemoryService( + mongo_client=_FakeMongoClient(), + database_name="test_db", + memory_ttl_days=None, + ) + + +async def test_ingest_writes_expire_at_default_seven_days() -> None: + """Each memory document stores an expire_at ~7 days in the future.""" + client: Any = _FakeMongoClient() + svc = MongoDbMemoryService(mongo_client=client, database_name="test_db") + before = datetime.now(timezone.utc) + await svc.add_session_to_memory( + _session_with_texts("first memory", "second memory") + ) + after = datetime.now(timezone.utc) + + docs = list(client["test_db"]["memories"]._docs.values()) + assert len(docs) == 2 + for doc in docs: + expire_at = doc["expire_at"] + if expire_at.tzinfo is None: + expire_at = expire_at.replace(tzinfo=timezone.utc) + assert expire_at - before >= timedelta(days=7) - timedelta(seconds=1) + assert expire_at - after <= timedelta(days=7) + + +async def test_ttl_days_configurable() -> None: + """A custom memory_ttl_days shifts the stored expire_at accordingly.""" + client: Any = _FakeMongoClient() + svc = MongoDbMemoryService( + mongo_client=client, database_name="test_db", memory_ttl_days=30 + ) + before = datetime.now(timezone.utc) + await svc.add_session_to_memory(_session_with_texts("hello world")) + + doc = next(iter(client["test_db"]["memories"]._docs.values())) + expire_at = doc["expire_at"] + if expire_at.tzinfo is None: + expire_at = expire_at.replace(tzinfo=timezone.utc) + assert expire_at - before >= timedelta(days=30) - timedelta(seconds=1) + + +async def test_ttl_disabled_omits_expire_at_and_index() -> None: + """memory_ttl_days=None stores no expire_at and creates no TTL index.""" + client: Any = _FakeMongoClient() + svc = MongoDbMemoryService( + mongo_client=client, database_name="test_db", memory_ttl_days=None + ) + await svc.add_session_to_memory(_session_with_texts("hello world")) + + doc = next(iter(client["test_db"]["memories"]._docs.values())) + assert "expire_at" not in doc + indexes = client["test_db"]["memories"].index_information() + assert "memory_ttl" not in indexes + # The recall index is still created regardless of the TTL setting. + assert "app_user_keywords" in indexes + + def test_constructor_validates_client_args() -> None: client: Any = _FakeMongoClient() with pytest.raises(ValueError, match="Either `connection_string`"):