Repository navigation
feat(integrations): add vector search retrieval to MongoDbMemoryService - #7573
Open
theshanbhag wants to merge 27 commits into
Open
theshanbhag wants to merge 27 commits into
theshanbhag wants to merge 27 commits into
Conversation
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.
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
9 tasks done
…b-session-service # Conflicts: # docs/guides/integrations/mongodb/mongodb_toolset/index.md # src/google/adk/integrations/mongodb/README.md # tests/unittests/integrations/mongodb/test_search_tool.py
These local-only sample/utility files (mongodb_samples/, .codegraph/, setup_mongodb_sample.py) were never meant to be committed; a git add -A in the merge commit included them by mistake. They stay on disk, just untracked again.
# Conflicts: # src/google/adk/integrations/mongodb/README.md
# Conflicts: # src/google/adk/integrations/mongodb/_client.py # src/google/adk/integrations/mongodb/_mongodb_toolset.py # tests/unittests/integrations/mongodb/test_mongodb_toolset.py
…ault 7 days) Memories previously lived forever. Each ingested memory document now stores an expire_at date memory_ttl_days (default 7) past ingest, and a memory_ttl TTL index on that field (created when ensure_indexes=True) lets MongoDB's TTL monitor delete expired memories. Pass memory_ttl_days=None to keep memories indefinitely. Re-ingesting a session refreshes expire_at.
Contributor
Author
|
Added a configurable TTL for stored memories, since memory retention was not previously handled.
Testing: |
Resolve conflicts from the stacked base (google#7509/google#7510/google#7533) having landed on main since this branch was cut: - MongoDbMemoryService: adopt upstream's rewritten base (compound _id keys, add_events_to_memory, unicode-aware tokenization, thought-part filtering, max_results=50, code-85 tolerance) and layer the vector/auto retrieval modes and configurable memory TTL (memory_ttl_days, default 7) on top. Index creation is now per-index so a conflict/privilege warning on one index does not skip the others. - Session service, toolset, their tests, constraints, and README: take upstream's reviewed versions, re-adding the memory vector/auto retrieval and TTL documentation to the README. - Extend the memory-service test fakes to cover TTL index options, vector search index creation, and $vectorSearch aggregation; add vector/auto/TTL test coverage (38 tests).
Contributor
Author
|
Resolved the merge conflicts with How conflicts were resolved:
Testing: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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. #7509 and #7510 have landed on main; stacked on #7533 (which carries the memory service, max_results, and the review fixes). The vector-memory commits are the last ones.Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
Problem:
MongoDbMemoryServicerecalls 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_modeconstructor argument with three modes; keyword stays the default so existing deployments are unaffected:"keyword"(default)$match+ overlap rank + limit (unchanged)"vector"embed_contentcall per session stores anembeddingper memory$vectorSearchwith the embedded query and anapp_name/user_idpre-filter"auto"textonly$vectorSearchwithquery.text— Atlas embeds with the index's Voyage AI model (Automated Embedding, Preview)Details:
max_results(vector modes as the$vectorSearchlimit, withnumCandidates= 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 avectorSearchindex — a cosinevectorfield onembeddingsized fromembedding_output_dimensionality(default 768), or anautoEmbedfield overtextwith the configurable Voyage model — both withapp_name/user_idfilter fields.genai_client,embedding_model_name,embedding_output_dimensionality,auto_embedding_model_name,vector_index_name.Bug caught by the live E2E run:
createSearchIndexesrejects 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:
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"):
retrieval_mode="vector"— service self-created the collection and vector index; after index READY, recall ranked the espresso memory first.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).Checklist
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.