Visitar URL original
fix(memory): apply rewinds to incremental events by 1aifanatic · Pull Request #8138 · google/adk-python · GitHub
Skip to content

fix(memory): apply rewinds to incremental events #8138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
1aifanatic wants to merge 1 commit into
base: main
Choose a base branch
from
+122 −15
Draft

fix(memory): apply rewinds to incremental events #8138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 48 additions & 11 deletions src/google/adk/memory/in_memory_memory_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,12 @@

from . import _utils
from ..events._rewind_events import _apply_rewinds
from ..events.event import Event
from .base_memory_service import BaseMemoryService
from .base_memory_service import SearchMemoryResponse
from .memory_entry import MemoryEntry

if TYPE_CHECKING:
from ..events.event import Event
from ..sessions.session import Session

_MAX_SEARCH_RESULTS = 10
Expand Down Expand Up @@ -112,12 +112,25 @@ async def add_session_to_memory(self, session: Session) -> None:
with self._lock:
deleted_ids = self._deleted_memory_ids.get(user_key, set())
self._session_events[user_key] = self._session_events.get(user_key, {})
# Invocation boundaries are needed when a later delta rewinds the session,
# even if that invocation has no searchable content or was deleted.
self._session_events[user_key][session.id] = [
event.model_copy(deep=True)
(
Event(
id=event.id,
invocation_id=event.invocation_id,
author=event.author,
)
if event.id in deleted_ids
else event.model_copy(deep=True)
)
for event in _apply_rewinds(session.events)
if event.content
and event.content.parts
and event.id not in deleted_ids
if event.invocation_id
or (
event.content
and event.content.parts
and event.id not in deleted_ids
)
]

@override
Expand All @@ -134,7 +147,11 @@ async def add_events_to_memory(
user_key = _user_key(app_name, user_id)
scoped_session_id = session_id or None
events_to_add = [
event for event in events if event.content and event.content.parts
event
for event in events
if event.invocation_id
or (event.content and event.content.parts)
or (event.actions and event.actions.rewind_before_invocation_id)
]

with self._lock:
Expand All @@ -145,10 +162,20 @@ async def add_events_to_memory(
)
existing_ids = {event.id for event in existing_events}
for event in events_to_add:
if event.id not in existing_ids and event.id not in deleted_ids:
if event.id in deleted_ids:
if not event.invocation_id:
continue
event = Event(
id=event.id,
invocation_id=event.invocation_id,
author=event.author,
)
if event.id not in existing_ids:
existing_events.append(event.model_copy(deep=True))
existing_ids.add(event.id)
self._session_events[user_key][scoped_session_id] = existing_events
self._session_events[user_key][scoped_session_id] = _apply_rewinds(
existing_events
)

@override
async def search_memory(
Expand Down Expand Up @@ -239,7 +266,17 @@ async def delete_memory(
user_key = _user_key(app_name, user_id)
with self._lock:
for session_id, events in self._session_events.get(user_key, {}).items():
kept = [event for event in events if event.id != memory_id]
if len(kept) != len(events):
self._session_events[user_key][session_id] = kept
if any(event.id == memory_id for event in events):
# Keep only the invocation boundary so a later rewind can find it.
self._session_events[user_key][session_id] = [
Event(
id=event.id,
invocation_id=event.invocation_id,
author=event.author,
)
if event.id == memory_id
else event
for event in events
if event.id != memory_id or event.invocation_id
]
self._deleted_memory_ids.setdefault(user_key, set()).add(memory_id)
78 changes: 74 additions & 4 deletions tests/unittests/memory/test_in_memory_memory_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@
parts=[types.Part(text='The ADK is a great toolkit.')]
),
),
# Event with no content, should be ignored by the service
# An invocation boundary without searchable content.
Event(
id='event-1b',
invocation_id='inv-2',
Expand Down Expand Up @@ -120,10 +120,11 @@ async def test_add_session_to_memory():
assert user_key in memory_service._session_events
session_memory = memory_service._session_events[user_key]
assert MOCK_SESSION_1.id in session_memory
# Check that the event with no content was filtered out
assert len(session_memory[MOCK_SESSION_1.id]) == 2
# Contentless invocation boundaries remain available for later rewinds.
assert len(session_memory[MOCK_SESSION_1.id]) == 3
assert session_memory[MOCK_SESSION_1.id][0].id == 'event-1a'
assert session_memory[MOCK_SESSION_1.id][1].id == 'event-1c'
assert session_memory[MOCK_SESSION_1.id][1].id == 'event-1b'
assert session_memory[MOCK_SESSION_1.id][2].id == 'event-1c'


@pytest.mark.asyncio
Expand Down Expand Up @@ -231,6 +232,7 @@ async def test_add_events_to_memory_appends_without_replacing():
session_memory = memory_service._session_events[user_key]
assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [
'event-1a',
'event-1b',
'event-1c',
'event-1d',
]
Expand Down Expand Up @@ -260,6 +262,7 @@ async def test_add_events_to_memory_deduplicates_event_ids():
session_memory = memory_service._session_events[user_key]
assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [
'event-1a',
'event-1b',
'event-1c',
]

Expand Down Expand Up @@ -1111,6 +1114,42 @@ def _rewind_event(rewind_before_invocation_id: str) -> Event:
)


@pytest.mark.parametrize('separate_delta', [False, True])
async def test_add_events_to_memory_skips_rewound_events(separate_delta):
"""Incremental ingestion removes recalled turns when a rewind arrives."""
service = InMemoryMemoryService()
scope = dict(app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID)
events = [
_text_event('kept', 'trip to paris'),
_text_event('rewound', 'address is baker street'),
]
await service.add_events_to_memory(
**scope,
session_id='other',
events=[_text_event('other', 'baker elsewhere')],
)
if separate_delta:
await service.add_events_to_memory(
**scope, session_id='target', events=events
)
events = []
events += [
_rewind_event('inv-rewound'),
_text_event('replacement', 'address is maple avenue'),
]

await service.add_events_to_memory(
**scope, session_id='target', events=events
)

old = await service.search_memory(**scope, query='baker')
kept = await service.search_memory(**scope, query='paris')
replacement = await service.search_memory(**scope, query='maple')
assert [m.id for m in old.memories] == ['event-other']
assert [m.id for m in kept.memories] == ['event-kept']
assert [m.id for m in replacement.memories] == ['event-replacement']


@pytest.mark.asyncio
async def test_add_session_to_memory_skips_rewound_events():
"""Turns a rewind took back are not stored, also when added before it."""
Expand All @@ -1137,3 +1176,34 @@ async def test_add_session_to_memory_skips_rewound_events():
)
assert not rewound.memories
assert [m.content.parts[0].text for m in kept.memories] == ['trip to paris']


@pytest.mark.parametrize('ingest_session', [False, True])
@pytest.mark.parametrize('delete_anchor', [False, True])
async def test_incremental_rewind_can_target_an_unsearchable_invocation(
ingest_session: bool, delete_anchor: bool
):
"""A contentless or deleted invocation still anchors a later rewind."""
service = InMemoryMemoryService()
scope = dict(app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID)
anchor = _text_event('anchor', 'private fact')
if not delete_anchor:
anchor.content = None
events = [anchor, _text_event('later', 'obsolete information')]
if ingest_session:
await service.add_session_to_memory(
Session(**scope, id='target', events=events)
)
else:
await service.add_events_to_memory(
**scope, session_id='target', events=events
)
if delete_anchor:
await service.delete_memory(**scope, memory_id=anchor.id)

await service.add_events_to_memory(
**scope, session_id='target', events=[_rewind_event('inv-anchor')]
)

assert not (await service.search_memory(**scope, query='obsolete')).memories
assert not (await service.search_memory(**scope, query='private')).memories
Loading