Repository navigation
fix(integrations): bound memory recall, default indexes on, keep MongoDB credentials out of Agent Engine packages - #7533
theshanbhag wants to merge 18 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.
There was a problem hiding this comment.
Code Review
This pull request introduces MongoDB integration for ADK, adding MongoDbSessionService for session persistence, MongoDbMemoryService for keyword-based long-term memory, and support for MongoDB Automated Embedding (Atlas Preview) in MongoDbToolset. It also adds pickling support to allow deployment to Agent Engine. The review feedback highlights several optimization opportunities, including resolving an N+1 query pattern in list_sessions, utilizing bulk_write for batch memory ingestion, and introducing thread locks to synchronize index creation on startup in both services.
|
Note on the failing The findings are all in repo-owned workflows this PR doesn't touch: (The |
…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.
Addresses the three remaining review notes from #7291. Stacked on the open split PRs (#7509, #7510, #7511, #7512 — it modifies files those PRs introduce); only the last commit is new.
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
Problem:
Three hardening gaps called out in the #7291 review:
ensure_indexes=False) — users had to opt in or reads degraded into collection scans as data grew.connection_stringis stored on the toolset and pickled into the uploaded package in plaintext.Solution:
search_memorynow ranks matches by keyword overlap (recency breaks ties) and caps results at a newmax_resultsconstructor argument (default 20). Ranking and limiting run server-side in the aggregation pipeline ($match→$addFieldsoverlap via$filter+$size→$sort→$limit), so a common keyword returns the best few memories rather than everything — and the excess never crosses the wire.ensure_indexes=True. A runtime user without thecreateIndexprivilege (the least-privilege setup the README recommends) triggers one warning and falls back to unindexed reads instead of failing every operation; unexpected index-creation errors still propagate, and the flag is only marked done on success or a handled privilege denial, so transient failures retry on the next operation.ensure_indexes=Falseremains as the explicit opt-out.MongoDbToolsetacceptsconnection_string_env_var: only the environment variable's name is stored on the toolset, and the value is resolved where the client is built — locally at construction, or on the Agent Engine runtime after unpickling (__setstate__resolves it there). Pickling therefore carries no credential, and the README's new Agent Engine section pairs it with Secret Manager-backed env vars (env_vars={"MONGODB_URI": SecretRef(secret="mongodb-uri", version="latest")}) so the plaintext never touches the deployment spec either. A unit test asserts the pickled payload contains neither the URI nor the password.Testing Plan
Unit Tests:
pytest tests/unittests/integrations/mongodb/ tests/unittests/features/→ 134 passed (mongomock-backed), including new coverage: recall limit, overlap ranking with recency tiebreak,max_resultsvalidation, indexes created by default, opt-out, privilege-denied fallback (warns once, reads still work), non-privilege index errors propagating, env-var resolution, missing-env-var error, multi-source constructor validation, and a pickle test asserting the credential bytes never enter the payload. All pre-commit hooks pass.Manual End-to-End (E2E) Tests:
Against a live Atlas 8.0.34 cluster: memory ingest + recall returns ranked, capped results; session and memory operations work with the default-on indexes; an agent built with
connection_string_env_var="MONGODB_URI"round-trips throughpickle.dumps/loadswith a working rebuilt client while the pickled bytes contain no URI.Checklist
Additional context
The MongoDB session and memory services still cannot be deployed to Agent Engine; the follow-up that wires them in will extend the same pickling +
connection_string_env_varhandling to them.