feat: Add JSON file memory service used in e2e testing#20
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hello @ipv1337, 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!
Summary of Changes
This pull request introduces a JsonFileMemoryService to persist agent sessions to a JSON file, primarily for use in end-to-end (E2E) testing. This addresses the limitation of the InMemoryMemoryService, which does not persist data between separate adk run invocations in E2E scripts. The PR also includes modifications to the software engineer agent to integrate the built-in load_memory tool and provides placeholder tools for manual memory persistence. Additionally, the E2E testing script is updated to utilize the new JsonFileMemoryService and includes memory recall tests.
Highlights
- New Feature: Adds a
JsonFileMemoryServicethat persists agent sessions to a JSON file, enabling cross-session memory recall in E2E tests. - Tool Integration: Integrates the built-in
load_memorytool into the software engineer agent and introduces placeholder tools for manual memory persistence. - E2E Testing: Updates the E2E testing script to utilize the
JsonFileMemoryServiceand includes tests for memory recall. - Documentation: Updates the architecture and testing strategies documentation to reflect the new memory service and testing approaches.
Changelog
Click here to see the changelog
- code_agent/adk/json_memory_service.py
- Introduces the
JsonFileMemoryServiceclass, which inherits fromBaseMemoryService. - Implements methods for loading sessions from and saving sessions to a JSON file.
- Includes
add_session_to_memory,load_memory,search_memory, andget_memory_service_infomethods. - Uses
ast.literal_evalfor safe evaluation of keys when loading from JSON. - Implements basic substring search for
load_memory.
- Introduces the
- code_agent/agent/software_engineer/software_engineer/agent.py
- Imports the built-in
load_memorytool fromgoogle.adk.tools. - Adds
load_memoryto the list of tools available to the agent. - Adds placeholder tools
save_current_session_to_file_toolandload_memory_from_file_tool. - Introduces
initialize_session_memoryfunction to initialize session memory in the tool context.
- Imports the built-in
- code_agent/agent/software_engineer/software_engineer/prompt.py
- Adds instructions for the agent to use the
load_memorytool to answer questions about past interactions. - Includes a description of the
load_memorytool and its usage. - Adds a placeholder section for manual memory persistence tools, indicating they are not yet implemented.
- Adds instructions for the agent to use the
- code_agent/agent/software_engineer/software_engineer/tools/init.py
- Exports the filesystem, shell command, search, system info, and persistent memory tools for easier imports.
- Defines
__all__to specify the tools available for import.
- code_agent/agent/software_engineer/software_engineer/tools/persistent_memory_tool.py
- Introduces placeholder tools for manually saving and loading session memory to a file.
- Includes
save_current_session_to_file_toolandload_memory_from_file_tool, which are not yet fully implemented. - Provides commented-out example implementations for these tools.
- docs/getting_started_architecture.md
- Adds a section on E2E testing with
adk run. - Includes diagrams illustrating the limitations of
InMemoryMemoryServicein E2E testing and the future state with a persistentMemoryService.
- Adds a section on E2E testing with
- docs/plan_adk_memory.md
- Adds a plan for integrating session memory into the Software Engineer Agent, including steps for defining the memory schema, implementing read/write operations, and developing test cases.
- docs/testing_strategies.md
- Adds documentation on testing strategies for the Code Agent, including unit, integration, and E2E tests.
- Discusses the limitations of
InMemoryMemoryServicein E2E testing and the need for a persistentMemoryService.
- scripts/run_e2e.py
- Modifies the E2E test script to allow injection of specific
SessionServiceandMemoryServiceimplementations. - Adds support for using
VertexAiRagMemoryServiceandJsonFileMemoryServicein E2E tests. - Includes logic to explicitly save the memory state after the agent run.
- Modifies the E2E test script to allow injection of specific
- scripts/run_e2e_tests.sh
- Updates the E2E testing script to utilize the
JsonFileMemoryService. - Adds memory recall tests that require a persistent
MemoryService. - Includes helper functions for asserting output and configuring the memory service.
- Updates the E2E testing script to utilize the
- tests/integration/test_memory_integration.py
- Adds a new test,
test_agent_load_memory_e2e, to test the agent's ability to recall information using theload_memorytool. - Marks the new test as expected to fail because the agent does not reliably call
load_memoryin the programmaticRunnercontext.
- Adds a new test,
Using Gemini Code Assist
The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and 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 to provide feedback.
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.
A JSON file, a memory's keep,
Where sessions slumber, secrets sleep.
In testing's realm, its worth we find,
A persistent agent, of a different kind.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a JSON file memory service for use in end-to-end testing, which is a valuable addition for enabling cross-session memory recall. The implementation appears well-structured, with clear separation of concerns and comprehensive error handling. However, there are a few areas that could benefit from further refinement to enhance robustness and maintainability.
Summary of Findings
- Error Handling in JSON Loading: The
_load_from_jsonmethod includes multiple exception handling blocks, but some exceptions are logged as warnings instead of errors. Consider promoting these to errors to ensure proper visibility of critical issues during memory loading. - Tool Implementation Completeness: The
save_current_session_to_file_toolandload_memory_from_file_toolare implemented as placeholders. While this is acceptable for the current scope, ensure that the agent's behavior gracefully handles the absence of these tools and that there are clear plans for future implementation. - Test Coverage for Memory Recall: The new test
test_agent_load_memory_e2eis marked asxfail. It's important to address the underlying issue causing the test to fail, as it directly relates to the core functionality of memory recall. Investigate why the agent is not reliably callingload_memoryin the programmatic Runner context.
Merge Readiness
The pull request introduces significant new functionality and is a step in the right direction. However, the test_agent_load_memory_e2e test failure and the placeholder tools indicate that further work is needed before merging. I recommend addressing the test failure and outlining a plan for the placeholder tools before considering this pull request ready for merging. I am unable to approve this pull request, and users should have others review and approve this code before merging.
| # Import your custom/wrapper service getter | ||
|
|
||
|
|
||
| @pytest.mark.xfail(reason="Agent currently does not reliably call load_memory in programmatic Runner context.") |
There was a problem hiding this comment.
The test is marked as xfail with the reason that the agent does not reliably call load_memory. It's crucial to investigate and address this issue, as it directly impacts the core functionality of memory recall. What steps have been taken to debug why the agent is not calling load_memory as expected?
| logger.warning(f"Failed to parse key {key_str}: {e}. Skipping.") | ||
| except Exception as e: | ||
| logger.warning(f"Failed to validate session data for key {key_str}: {e}. Skipping.") |
There was a problem hiding this comment.
These exceptions are currently logged as warnings. Given that these indicate failures to parse or validate session data, should these be logged as errors instead to ensure they are more visible? This would help in debugging issues with corrupted or malformed session data in the JSON file.
| # --- END DETAILED LOGGING --- | ||
|
|
||
| # logger.debug(f"Retrieved session object content: {completed_session}") # Keep debug for full object | ||
| if completed_session and isinstance(memory_service, JsonFileMemoryService): # Check type before accessing private members |
There was a problem hiding this comment.
Consider adding a check to ensure completed_session is not None before accessing its attributes. This will prevent potential AttributeError exceptions if the session retrieval fails.
| if completed_session and isinstance(memory_service, JsonFileMemoryService): # Check type before accessing private members | |
| if completed_session and isinstance(memory_service, JsonFileMemoryService): | |
Test Coverage ReportOverall Coverage: 82.33%
|
|
1 participant
Summary of Changes
This pull request introduces a
JsonFileMemoryServiceto persist agent sessions to a JSON file, primarily for use in end-to-end (E2E) testing. This addresses the limitation of theInMemoryMemoryService, which does not persist data between separateadk runinvocations in E2E scripts. The PR also includes modifications to the software engineer agent to integrate the built-inload_memorytool and provides placeholder tools for manual memory persistence. Additionally, the E2E testing script is updated to utilize the newJsonFileMemoryServiceand includes memory recall tests.Highlights
JsonFileMemoryServicethat persists agent sessions to a JSON file, enabling cross-session memory recall in E2E tests.load_memorytool into the software engineer agent and introduces placeholder tools for manual memory persistence.JsonFileMemoryServiceand includes tests for memory recall.Changelog
Click here to see the changelog
JsonFileMemoryServiceclass, which inherits fromBaseMemoryService.add_session_to_memory,load_memory,search_memory, andget_memory_service_infomethods.ast.literal_evalfor safe evaluation of keys when loading from JSON.load_memory.load_memorytool fromgoogle.adk.tools.load_memoryto the list of tools available to the agent.save_current_session_to_file_toolandload_memory_from_file_tool.initialize_session_memoryfunction to initialize session memory in the tool context.load_memorytool to answer questions about past interactions.load_memorytool and its usage.__all__to specify the tools available for import.save_current_session_to_file_toolandload_memory_from_file_tool, which are not yet fully implemented.adk run.InMemoryMemoryServicein E2E testing and the future state with a persistentMemoryService.InMemoryMemoryServicein E2E testing and the need for a persistentMemoryService.SessionServiceandMemoryServiceimplementations.VertexAiRagMemoryServiceandJsonFileMemoryServicein E2E tests.JsonFileMemoryService.MemoryService.test_agent_load_memory_e2e, to test the agent's ability to recall information using theload_memorytool.load_memoryin the programmaticRunnercontext.