Skip to content

feat(integrations): add vector search retrieval to MongoDbMemoryService - #7573

Open
theshanbhag wants to merge 20 commits into
google:mainfrom
theshanbhag:mongodb-vector-memory
Open

theshanbhag wants to merge 20 commits into
google:mainfrom
theshanbhag:mongodb-vector-memory

Conversation

@theshanbhag

Copy link
Copy Markdown
Contributor

Adds MongoDB Vector Search retrieval to MongoDbMemoryService — the follow-up promised in the #7291 review ("do you plan to use vector search here?"). Keyword retrieval stays the default; two vector modes join it. Stacked on #7533 (which carries the memory service, max_results, and the review fixes); only the last two commits are new.

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

Problem:

MongoDbMemoryService recalls memories by keyword overlap only — synonyms and paraphrases miss ("coffee maker repair" does not recall "my espresso machine stopped working"). The integration already supports vector search for its tools, so memory recall should offer it too.

Solution:

A retrieval_mode constructor argument with three modes; keyword stays the default so existing deployments are unaffected:

Mode Ingest Recall Requirements
"keyword" (default) Extract keywords (unchanged) $match + overlap rank + limit (unchanged) Any MongoDB deployment
"vector" One batch embed_content call per session stores an embedding per memory $vectorSearch with the embedded query and an app_name/user_id pre-filter Vector-search-capable deployment + Google credentials
"auto" Plain text only $vectorSearch with query.text — Atlas embeds with the index's Voyage AI model (Automated Embedding, Preview) Atlas; no Google credentials

Details:

  • All modes honor max_results (vector modes as the $vectorSearch limit, with numCandidates = max(10x, 100)) and the idempotent per-event document keys.
  • ensure_indexes (default on) now creates the index the mode needs: the keyword multikey index, or a vectorSearch index — a cosine vector field on embedding sized from embedding_output_dimensionality (default 768), or an autoEmbed field over text with the configurable Voyage model — both with app_name/user_id filter fields.
  • New knobs: genai_client, embedding_model_name, embedding_output_dimensionality, auto_embedding_model_name, vector_index_name.
  • Mode/ingest consistency is documented: memories ingested in one mode are invisible to the others.

Bug caught by the live E2E run: createSearchIndexes rejects a collection that does not exist, so a fresh deployment's first ingest failed in vector/auto modes (NamespaceNotFound). Index creation now creates the memories collection first (tolerating a cross-process race via the CollectionInvalid code).

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

pytest tests/unittests/integrations/mongodb/ tests/unittests/features/ → 145 passed. New coverage: mode validation, vector ingest batching (one embed call, vectors stored, no keywords), vector recall stage shape (queryVector, path, app/user filter, limit) with parsed results, auto ingest (plain text, genai client never called), auto recall (query.text, no embed call), and index creation for both vector modes (vector vs autoEmbed fields, collection pre-creation). All pre-commit hooks pass.

Manual End-to-End (E2E) Tests:

Against live Atlas 8.0.34, ingesting a session whose relevant event shares no keywords with the query ("my espresso machine stopped working" vs. query "coffee maker repair"):

  1. retrieval_mode="vector" — service self-created the collection and vector index; after index READY, recall ranked the espresso memory first.
  2. retrieval_mode="auto" — service self-created the autoEmbed index; recall returned the espresso memory first with zero Google credentials in the service (after Atlas's embedding lag, the documented caveat).
  3. Both modes ranked an unrelated memory below it, and all test state was cleaned up.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

Auto mode has an inherent freshness lag: Atlas embeds documents asynchronously, so just-ingested memories may take a few seconds to become recallable (the E2E poll observed ~10–60 s on a shared-tier cluster). This is documented in the README and class docstring.

Add `use_mongodb_auto_embedding` and `mongodb_auto_embedding_model` to
MongoDbToolSettings. When enabled, the search tools send the query as
plain text (`query.text`) instead of embedding it through the genai
client first: Atlas generates the embedding itself with the Voyage AI
model configured on the index's `autoEmbed` field (Atlas Automated
Embedding, Preview), so deployments that already store data on Atlas
need no Google embedding credentials. The default mode (Google
embedding models via `queryVector`) is unchanged.

Split from google#7291.
Agent Engine packages the app object with cloudpickle, and a live
pymongo.MongoClient (sockets, locks, background threads) cannot cross
that boundary. Constructed with `connection_string`, the toolset now
drops the client at pickle time and rebuilds it on the runtime; the
query-embedding genai client is likewise rebuilt lazily from the
ambient environment on the destination. Constructing with a
caller-owned `mongo_client` raises TypeError at pickle time with
guidance, since such a client cannot be rebuilt remotely.

Pickling support for the MongoDB session and memory services is
deliberately left to the follow-up that wires those services into
Agent Engine.

Split from google#7291.
MongoDbSessionService persists sessions, events, and app:/user:-scoped
state in MongoDB, exported from google.adk.integrations.mongodb behind
the MONGODB_SESSION_SERVICE experimental feature (default on).

Concurrency model: within a process a per-session lock serializes
append_event; across runner replicas the optimistic-concurrency
revision on the session document is the only guard, and the losing
writer gets StaleSessionError. Multi-document writes run in a
transaction on replica sets / sharded clusters / Atlas; on deployments
without transaction support (standalone mongod) the verdict is cached
and writes fall back to sequential with the revision check. Optional
`ensure_indexes` creates the recommended secondary indexes on first
use. mongomock joins the test dependencies for the new unit tests.

NO_UNIT_GUIDE=The integration README documents the service, including
its document layout, concurrency model, and least-privilege setup; a
per-unit guide will follow when the experimental API stabilizes.

Split from google#7291.
MongoDbMemoryService ingests session events into a memory collection
and recalls them by keyword matching, exported from
google.adk.integrations.mongodb behind the MONGODB_MEMORY_SERVICE
experimental feature (default on). Ingestion is idempotent (memory
documents are keyed by app/user/session/event), and `ensure_indexes`
creates the app/user/keywords index on first use.

Retrieval is intentionally keyword-only: it needs no embedding model,
no Google credentials, and no vector index, so it works on any MongoDB
deployment at zero per-query cost. Vector recall is a possible
follow-up building on this integration's own search tools.

NO_UNIT_GUIDE=The integration README documents the service, including
its storage layout and least-privilege setup; a per-unit guide will
follow when the experimental API stabilizes.

Split from google#7291.
…oDB credentials out of Agent Engine packages

Three hardening fixes from the google#7291 review:

- MongoDbMemoryService.search_memory ranked nothing and returned every
  keyword match. Results are now ranked by keyword overlap (recency
  breaks ties) and capped by a new max_results constructor argument
  (default 20); ranking and limiting run server-side in the
  aggregation pipeline so a common keyword no longer returns
  everything.
- ensure_indexes defaulted to False, so users had to opt in or reads
  degraded into collection scans as data grew. It now defaults to True
  for both services; a deployment without the createIndex privilege
  logs one warning and falls back to unindexed reads instead of
  failing (unexpected index-creation errors still propagate and are
  retried on the next operation).
- A literal connection_string stored on the toolset was pickled into
  the Agent Engine deployment package in plaintext, password included.
  MongoDbToolset now accepts connection_string_env_var: only the
  variable's name is stored, and the value is resolved where the
  client is built — on the runtime after unpickling. The README's new
  Agent Engine section pairs it with Secret Manager-backed env vars
  (SecretRef) so the credential never touches the deployment spec.

Split from google#7291.
…b-hardening

# Conflicts:
#	constraints-3.11.txt
#	constraints-3.12.txt
#	constraints-3.13.txt
#	constraints-3.14.txt
…b-session-service

# Conflicts:
#	constraints-3.11.txt
#	constraints-3.12.txt
#	constraints-3.13.txt
#	constraints-3.14.txt
…n MongoDbSessionService

Review follow-up from google#7533: list_sessions fetched each user's state
with a separate find_one per user (N+1); it now issues a single $in
query against the user_states collection. Index creation on first use
is also serialized with a thread lock (double-checked flag), so
concurrent first calls no longer issue duplicate create_index calls.

Addresses gemini-code-assist review comments on google#7533.
Review follow-up from google#7533: add_session_to_memory issued one
replace_one round-trip per event; it now batches all upserts into a
single bulk_write (ordered=False). First-use index creation is
serialized with a thread lock (double-checked flag) so concurrent
startup calls no longer issue duplicate create_index calls. Tests
carry a shim for mongomock's BulkOperationBuilder, which rejects the
`sort` key pymongo 4.9+ adds to bulk operations.

Addresses gemini-code-assist review comments on google#7533.
# Conflicts:
#	src/google/adk/integrations/mongodb/_memory_service.py
#	src/google/adk/integrations/mongodb/_session_service.py
pyink reflow of the merged _ensure_indexes_once bodies, plus
constraints regenerated against upstream's 2026-10-14 cutoff (now
including the new constraints-3.15.txt) with the branch's mongomock
entries applied consistently across all five Python versions.
Answers the google#7291 review question with the promised follow-up: memory
recall is no longer keyword-only. `retrieval_mode` selects the recall
strategy:

- "keyword" (default): unchanged — dependency-free overlap matching that
  works on any MongoDB deployment.
- "vector": texts are embedded at ingest (one batch call per session)
  with the Google embedding model through the genai client, and recall
  embeds the query and runs $vectorSearch with an app/user pre-filter.
- "auto": Atlas Automated Embedding — documents carry plain text and an
  autoEmbed index makes Atlas embed documents and queries (no Google
  credentials; freshly ingested memories become recallable once Atlas
  embeds them).

ensure_indexes now creates the index the mode needs: the keyword
multikey index, or a vectorSearch index (vector field sized from
embedding_output_dimensionality, or an autoEmbed field over `text`,
both with app_name/user_id filter fields). It also creates the
memories collection first — createSearchIndexes rejects a nonexistent
collection, which a fresh deployment hits on its first ingest (caught
by the live E2E test).

All modes honor max_results; keyword mode stays the default so
existing deployments are unaffected.

Split from google#7291 follow-ups; stacked on google#7533.
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants