Repository navigation
Conversation
add_session_to_memory in the five memory services read session.events directly, so the turns Runner.rewind_async took back were ingested, and load_memory / preload_memory returned them in later sessions. Filter the events with _apply_rewinds, as prompt building and compaction already do.
Contributor
There was a problem hiding this comment.
Code Review
This pull request updates various memory services (Firestore, SQLite, In-Memory, Vertex AI Memory Bank, and Vertex AI RAG) to filter out rewound events when adding a session to memory by applying the _apply_rewinds utility function. It also adds corresponding unit tests for each service to ensure that rewound events are correctly skipped and not stored. I have no additional feedback to provide as there are no review comments.
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.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
Problem:
Runner.rewind_asynckeeps every event and appends a rewind event, and_apply_rewindsdecides which events are still live; prompt building and compaction both use it.add_session_to_memoryin all five memory services readsession.eventsdirectly, so the turns a rewind took back were ingested into long-term memory, andload_memory/preload_memoryhanded them back to the model in later sessions. The rewind documentation says the rewound requests are ignored when the next model requests are prepared.Solution:
add_session_to_memoryinInMemoryMemoryService,SqliteMemoryService,VertexAiMemoryBankService,VertexAiRagMemoryServiceandFirestoreMemoryServicepasses_apply_rewinds(session.events)instead ofsession.events, the same filter the prompt and compaction paths use.InMemoryMemoryServicereplaces a session's events on each call, so adding the session again after a rewind also drops turns stored before the rewind.add_events_to_memory, where the caller chooses the events, and memories a remote service (Memory Bank, RAG, Firestore) already ingested before the rewind; the rewind documentation leaves external systems to the caller.Testing Plan
Unit Tests:
New tests, one per service, each with a kept turn, a rewound turn and the rewind event:
test_in_memory_memory_service.py::test_add_session_to_memory_skips_rewound_events: search no longer finds the rewound turn, also when the session was added once before the rewind.test_sqlite_memory_service.py::test_add_session_to_memory_skips_rewound_events(FTS on and off).test_vertex_ai_memory_bank_service.py::test_add_session_to_memory_skips_rewound_events: only the kept event is sent toingest_events.test_vertex_ai_rag_memory_service.py::test_add_session_to_memory_skips_rewound_events: the uploaded file has the kept turn and not the rewound one.test_firestore_memory_service.py::test_add_session_to_memory_skips_rewound_events: only the kept turn's keywords are written.All 6 cases fail with the
src/change reverted and pass with it:The full
tests/unittestsrun has no new failures: 6 more tests pass than onmain, and the tests that fail locally (Windows shell, symlink and path tests) fail onmaintoo.Manual End-to-End (E2E) Tests:
The script in #7570 (no model or credentials needed) rewinds a two-turn session and adds it to
InMemoryMemoryServiceandSqliteMemoryService.Before:
After:
Checklist