Skip to content

Commit 4fae7b7

Browse files
hyeonsang010716copybara-github
authored andcommitted
fix: keep thought parts out of Memory Bank events
Merge #7851 VertexAiMemoryBankService no longer sends model thought parts to Memory Bank, and skips events whose only content is thoughts. Fixes #7850 PiperOrigin-RevId: 1010580092
1 parent 87d9dcd commit 4fae7b7

2 files changed

Lines changed: 118 additions & 3 deletions

File tree

‎src/google/adk/memory/vertex_ai_memory_bank_service.py‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -360,7 +360,7 @@ async def _add_events_to_memory_from_events(
360360
if event.content:
361361
direct_events.append(
362362
agentplatform.types.GenerateMemoriesRequestDirectContentsSourceEvent(
363-
content=event.content
363+
content=_without_thoughts(event.content)
364364
)
365365
)
366366
if direct_events:
@@ -431,7 +431,7 @@ async def _add_events_to_memory_via_ingest(
431431
)
432432
direct_events.append(
433433
agentplatform.types.IngestionDirectContentsSourceEvent(
434-
content=event.content,
434+
content=_without_thoughts(event.content),
435435
event_id=event.id,
436436
event_time=event_time,
437437
)
@@ -732,7 +732,23 @@ def _should_filter_out_event(event: Event) -> bool:
732732
# Deferred: a module-level import would load the agent and flow stack too.
733733
from ..flows.llm_flows.context._contents import _is_auth_event
734734

735-
return _is_empty_content(event.content) or _is_auth_event(event)
735+
content = _without_thoughts(event.content)
736+
return _is_empty_content(content) or _is_auth_event(event)
737+
738+
739+
def _without_thoughts(content: types.Content | None) -> types.Content | None:
740+
"""Returns the content without the model's thought parts.
741+
742+
Thoughts are the model's reasoning, not part of the conversation, so they
743+
are kept out of the events that memories are generated from.
744+
"""
745+
if not content or not content.parts:
746+
return content
747+
if not any(part.thought for part in content.parts):
748+
return content
749+
return content.model_copy(
750+
update={'parts': [part for part in content.parts if not part.thought]}
751+
)
736752

737753

738754
def _is_empty_content(content: types.Content | None) -> bool:

‎tests/unittests/memory/test_vertex_ai_memory_bank_service.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,35 @@ async def __anext__(self) -> Any:
142142
),
143143
]
144144

145+
THOUGHT_TEXT = 'the user sounds anxious, keep the reply short'
146+
147+
THOUGHT_EVENTS = [
148+
Event(
149+
id='901',
150+
invocation_id='900',
151+
author='agent',
152+
timestamp=45678,
153+
content=types.Content(
154+
role='model',
155+
parts=[
156+
types.Part(text=THOUGHT_TEXT, thought=True),
157+
types.Part(text='Here is the answer.'),
158+
],
159+
),
160+
),
161+
# Thought-only event, should be ignored
162+
Event(
163+
id='902',
164+
invocation_id='900',
165+
author='agent',
166+
timestamp=45679,
167+
content=types.Content(
168+
role='model',
169+
parts=[types.Part(text=THOUGHT_TEXT, thought=True)],
170+
),
171+
),
172+
]
173+
145174
MOCK_SESSION_WITH_EMPTY_EVENTS = Session(
146175
app_name=MOCK_APP_NAME,
147176
user_id=MOCK_USER_ID,
@@ -582,6 +611,76 @@ async def test_add_events_to_memory_via_generate_skips_credential_events(
582611
assert CREDENTIAL_TOKEN not in str(call_kwargs)
583612

584613

614+
@pytest.mark.asyncio
615+
async def test_add_session_to_memory_skips_thought_parts(
616+
mock_agentplatform_client,
617+
):
618+
memory_service = mock_vertex_ai_memory_bank_service()
619+
session = Session(
620+
app_name=MOCK_APP_NAME,
621+
user_id=MOCK_USER_ID,
622+
id='333',
623+
last_update_time=22333,
624+
events=THOUGHT_EVENTS,
625+
)
626+
627+
await memory_service.add_session_to_memory(session)
628+
629+
# Allow the fire-and-forget task to complete.
630+
await asyncio.sleep(0)
631+
632+
call_kwargs = (
633+
mock_agentplatform_client.memory_banks.ingest_events.call_args.kwargs
634+
)
635+
source = call_kwargs['direct_contents_source']
636+
assert [event.event_id for event in source.events] == ['901']
637+
assert source.events[0].content.role == 'model'
638+
assert [part.text for part in source.events[0].content.parts] == [
639+
'Here is the answer.'
640+
]
641+
assert THOUGHT_TEXT not in str(call_kwargs)
642+
# The session's own events are left unchanged.
643+
assert len(THOUGHT_EVENTS[0].content.parts) == 2
644+
645+
646+
@pytest.mark.asyncio
647+
async def test_add_events_to_memory_via_generate_skips_thought_parts(
648+
mock_agentplatform_client,
649+
):
650+
memory_service = mock_vertex_ai_memory_bank_service()
651+
await memory_service.add_events_to_memory(
652+
app_name=MOCK_APP_NAME,
653+
user_id=MOCK_USER_ID,
654+
events=THOUGHT_EVENTS,
655+
custom_metadata={'revision_ttl': '3600s'},
656+
)
657+
658+
call_kwargs = (
659+
mock_agentplatform_client.memory_banks.memories.generate.call_args.kwargs
660+
)
661+
source = call_kwargs['direct_contents_source']
662+
assert len(source.events) == 1
663+
assert [part.text for part in source.events[0].content.parts] == [
664+
'Here is the answer.'
665+
]
666+
assert THOUGHT_TEXT not in str(call_kwargs)
667+
668+
669+
@pytest.mark.asyncio
670+
async def test_add_events_to_memory_via_generate_with_only_thoughts_skips_rpc(
671+
mock_agentplatform_client,
672+
):
673+
memory_service = mock_vertex_ai_memory_bank_service()
674+
await memory_service.add_events_to_memory(
675+
app_name=MOCK_APP_NAME,
676+
user_id=MOCK_USER_ID,
677+
events=[THOUGHT_EVENTS[1]],
678+
custom_metadata={'revision_ttl': '3600s'},
679+
)
680+
681+
mock_agentplatform_client.memory_banks.memories.generate.assert_not_called()
682+
683+
585684
@pytest.mark.asyncio
586685
async def test_add_events_to_memory_with_explicit_events_and_metadata(
587686
mock_agentplatform_client,

0 commit comments

Comments
 (0)