feat(memory): add DatabaseMemoryService with SQL backend and agent scratchpad#4736
feat(memory): add DatabaseMemoryService with SQL backend and agent scratchpad#4736Raman369AI wants to merge 2 commits intogoogle:mainfrom
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a robust, persistent memory solution for agents, addressing the current limitation of volatile in-memory storage. By integrating a Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Response from ADK Triaging Agent Hello @Raman369AI, thank you for your contribution! Before we can review your pull request, you'll need to sign our Contributor License Agreement (CLA). It seems the CLA check has failed. You can review and sign the agreement here: https://cla.developers.google.com/ Once the CLA is signed, we can proceed with the review. Thanks! |
Adds a durable, RDBMS-backed memory service that works with any SQLAlchemy-supported database (SQLite, PostgreSQL, MySQL, MariaDB) as an alternative to the volatile InMemoryMemoryService. Key additions: - DatabaseMemoryService: implements BaseMemoryService with lazy table creation, idempotent session ingest, and delta event ingestion - MemorySearchBackend ABC + KeywordSearchBackend: LIKE/ILIKE search with AND-first → OR-fallback tokenization strategy - Scratchpad KV store and append-only log for intermediate agent state - Four agent-callable BaseTool subclasses: scratchpad_get_tool, scratchpad_set_tool, scratchpad_append_log_tool, scratchpad_get_log_tool - 38 unit tests covering all methods, tool happy-paths, wrong-service errors, multi-user isolation, and session scoping
5c69b8f to
7e9c8e9
Compare
There was a problem hiding this comment.
Code Review
This pull request introduces a DatabaseMemoryService that provides a durable, SQL-backed memory and scratchpad for agents, including a search backend abstraction, SQLAlchemy ORM schemas, and agent-facing tools. While the implementation demonstrates high code quality and strong patterns, a security issue was identified: sensitive database credentials could be leaked in error messages if the database initialization fails. This should be addressed by sanitizing the database URL in exception messages. Additionally, there are a couple of suggestions aimed at enhancing code readability and simplifying logic in a few areas.
| except ArgumentError as exc: | ||
| raise ValueError( | ||
| f"Invalid database URL format or argument '{db_url}'." | ||
| ) from exc |
There was a problem hiding this comment.
The DatabaseMemoryService.__init__ method catches ArgumentError and ImportError during database engine creation and re-raises them as a ValueError. The error message includes the full db_url, which may contain sensitive credentials such as database passwords. If this exception is logged or displayed to a user, it could lead to the exposure of these credentials.
To remediate this, avoid including the full db_url in the error message. If the URL must be included for debugging purposes, ensure that any passwords or sensitive tokens are sanitized or masked.
| except ImportError as exc: | ||
| raise ValueError( | ||
| f"Database-related module not found for URL '{db_url}'." | ||
| ) from exc |
There was a problem hiding this comment.
The DatabaseMemoryService.__init__ method catches ArgumentError and ImportError during database engine creation and re-raises them as a ValueError. The error message includes the full db_url, which may contain sensitive credentials such as database passwords. If this exception is logged or displayed to a user, it could lead to the exposure of these credentials.
To remediate this, avoid including the full db_url in the error message. If the URL must be included for debugging purposes, ensure that any passwords or sensitive tokens are sanitized or masked.
| """Return True if the event has no usable text content.""" | ||
| if not event.content or not event.content.parts: | ||
| return True | ||
| return not any(part.text for part in event.content.parts if part.text) |
There was a problem hiding this comment.
The logic to check for usable text content can be simplified. The if part.text condition within the generator expression is redundant because any() on an iterator of strings will correctly evaluate to False only if all strings are empty.
| return not any(part.text for part in event.content.parts if part.text) | |
| return not any(part.text for part in event.content.parts) |
| tokens = [ | ||
| cleaned | ||
| for raw in query.split() | ||
| if raw.strip() | ||
| for cleaned in [re.sub(r"[^\w]", "", raw).lower()] | ||
| if cleaned | ||
| ] |
There was a problem hiding this comment.
This list comprehension is quite complex and can be difficult to read due to the nested for and the use of a single-element list to assign a variable. For better maintainability and readability, I suggest refactoring this into a simple for loop.
Note that query.split() (with no arguments) handles splitting by any whitespace and discards empty strings, so the if raw.strip() check becomes unnecessary.
| tokens = [ | |
| cleaned | |
| for raw in query.split() | |
| if raw.strip() | |
| for cleaned in [re.sub(r"[^\w]", "", raw).lower()] | |
| if cleaned | |
| ] | |
| tokens = [] | |
| for raw_token in query.split(): | |
| # Clean the token by removing non-alphanumeric characters and lowercasing. | |
| cleaned_token = re.sub(r"[^\w]", "", raw_token).lower() | |
| # Add the token only if it's not empty after cleaning. | |
| if cleaned_token: | |
| tokens.append(cleaned_token) |
Closes #4735
Summary
DatabaseMemoryService, aBaseMemoryServicebacked by any SQLAlchemy async-compatible database (SQLite, PostgreSQL, MySQL, MariaDB)MemorySearchBackendABC andKeywordSearchBackend(LIKE/ILIKE, AND-first → OR-fallback tokenization)BaseToolsubclassesMotivation
The existing
InMemoryMemoryServiceis volatile and explicitly test-only. Developers not using Vertex AI have no durable, self-hosted memory option.DatabaseMemoryServicefills that gap using the same SQLAlchemy async pattern already established byDatabaseSessionService.Changes
src/google/adk/memory/schemas/__init__.pysrc/google/adk/memory/schemas/memory_schema.pyadk_memory_entries,adk_scratchpad_kv,adk_scratchpad_log) with a separateDeclarativeBasesrc/google/adk/memory/memory_search_backend.pyMemorySearchBackendABC +KeywordSearchBackendsrc/google/adk/memory/database_memory_service.pysrc/google/adk/tools/scratchpad_tool.pyScratchpadGetTool,ScratchpadSetTool,ScratchpadAppendLogTool,ScratchpadGetLogTool+ singleton instancessrc/google/adk/memory/__init__.pytry/except ImportErrortests/unittests/memory/test_database_memory_service.pyDesign notes
DeclarativeBasefrom sessions schema — avoids table coupling while allowing the same DBsession_id=''sentinel in scratchpad tables — enables user-level (non-session) scope without nullable PK columnsasyncio.Lock(double-checked) — noawait svc.initialize()call requiredtry/except ImportErrorin__init__.py— SQLAlchemy + async driver are optional; users without them are unaffectedTYPE_CHECKINGguard forEvent/Sessionimports — correct pattern withfrom __future__ import annotationsTest plan
All tests run with
sqlite+aiosqlite:///:memory:— no external database required.Scenarios covered:
add_session_to_memory— filters empty events, persists content/author/timestampadd_events_to_memory— delta, skips duplicateevent_idadd_memory— directMemoryEntrypersist, auto-UUIDsearch_memory— AND match, OR fallback, empty query, no matchMemorySearchBackendis honouredValueErrorValueErrorFormatting
isort+pyink --config pyproject.tomlrun on all changed files — no remaining violations.