feat: add add_events_to_memory facade for event-delta

Adds BaseMemoryService.add_events_to_memory(session, events=..., custom_metadata=...) and CallbackContext.add_events_to_memory(events=..., custom_metadata=...) so callers can add memories from an explicit subset of ADK events. Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 868261578
2026-07-09 18:19:28 -07:00 · 2026-02-10 12:14:52 -08:00
parent de79bf12b5
commit 59e88972ae
7 changed files with 661 additions and 14 deletions
@@ -118,6 +118,116 @@ async def test_add_session_to_memory():
  assert session_memory[MOCK_SESSION_1.id][1].id == 'event-1c'


+@pytest.mark.asyncio
+async def test_add_events_to_memory_with_explicit_events():
+  """Tests that add_events_to_memory can ingest an explicit event list."""
+  memory_service = InMemoryMemoryService()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION_1.app_name,
+      user_id=MOCK_SESSION_1.user_id,
+      session_id=MOCK_SESSION_1.id,
+      events=[MOCK_SESSION_1.events[0]],
+  )
+
+  user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}'
+  session_memory = memory_service._session_events[user_key]
+  assert len(session_memory[MOCK_SESSION_1.id]) == 1
+  assert session_memory[MOCK_SESSION_1.id][0].id == 'event-1a'
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_without_session_id_uses_default_bucket():
+  """Tests add_events_to_memory when no session_id is provided."""
+  memory_service = InMemoryMemoryService()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION_1.app_name,
+      user_id=MOCK_SESSION_1.user_id,
+      events=[MOCK_SESSION_1.events[0]],
+  )
+
+  user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}'
+  session_memory = memory_service._session_events[user_key]
+  assert len(session_memory) == 1
+  unknown_session_events = next(iter(session_memory.values()))
+  assert len(unknown_session_events) == 1
+  assert unknown_session_events[0].id == 'event-1a'
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_alias_is_supported():
+  """Tests that add_events_to_memory remains a compatibility alias."""
+  memory_service = InMemoryMemoryService()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION_1.app_name,
+      user_id=MOCK_SESSION_1.user_id,
+      session_id=MOCK_SESSION_1.id,
+      events=[MOCK_SESSION_1.events[0]],
+  )
+
+  user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}'
+  session_memory = memory_service._session_events[user_key]
+  assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [
+      'event-1a'
+  ]
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_appends_without_replacing():
+  """Tests that add_events_to_memory appends events rather than replacing."""
+  memory_service = InMemoryMemoryService()
+  await memory_service.add_session_to_memory(MOCK_SESSION_1)
+
+  new_event = Event(
+      id='event-1d',
+      invocation_id='inv-6',
+      author='user',
+      timestamp=12348,
+      content=types.Content(parts=[types.Part(text='A new fact.')]),
+  )
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION_1.app_name,
+      user_id=MOCK_SESSION_1.user_id,
+      session_id=MOCK_SESSION_1.id,
+      events=[new_event],
+  )
+
+  user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}'
+  session_memory = memory_service._session_events[user_key]
+  assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [
+      'event-1a',
+      'event-1c',
+      'event-1d',
+  ]
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_deduplicates_event_ids():
+  """Tests that duplicate event IDs are not appended multiple times."""
+  memory_service = InMemoryMemoryService()
+  await memory_service.add_session_to_memory(MOCK_SESSION_1)
+
+  duplicate_event = Event(
+      id='event-1a',
+      invocation_id='inv-7',
+      author='user',
+      timestamp=12349,
+      content=types.Content(parts=[types.Part(text='Updated duplicate text.')]),
+  )
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION_1.app_name,
+      user_id=MOCK_SESSION_1.user_id,
+      session_id=MOCK_SESSION_1.id,
+      events=[duplicate_event],
+  )
+
+  user_key = f'{MOCK_APP_NAME}/{MOCK_USER_ID}'
+  session_memory = memory_service._session_events[user_key]
+  assert [event.id for event in session_memory[MOCK_SESSION_1.id]] == [
+      'event-1a',
+      'event-1c',
+  ]
+
+
@pytest.mark.asyncio
 async def test_add_session_with_no_events_to_memory():
  """Tests that adding a session with no events does not cause an error."""
@@ -23,11 +23,19 @@ from google.adk.memory.vertex_ai_memory_bank_service import VertexAiMemoryBankSe
 from google.adk.sessions.session import Session
 from google.genai import types
 import pytest
+from vertexai._genai.types import common as vertex_common_types

 MOCK_APP_NAME = 'test-app'
 MOCK_USER_ID = 'test-user'


+def _supports_generate_memories_metadata() -> bool:
+  return (
+      'metadata'
+      in vertex_common_types.GenerateAgentEngineMemoriesConfig.model_fields
+  )
+
+
 class _AsyncListIterator:
  """Minimal async iterator wrapper for list-like results."""

@@ -156,6 +164,217 @@ async def test_add_session_to_memory(mock_vertexai_client):
  )


+@pytest.mark.asyncio
+async def test_add_events_to_memory_with_explicit_events_and_metadata(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION.app_name,
+      user_id=MOCK_SESSION.user_id,
+      session_id=MOCK_SESSION.id,
+      events=[MOCK_SESSION.events[0]],
+      custom_metadata={'ttl': '6000s', 'source': 'agent'},
+  )
+
+  expected_config = {
+      'wait_for_completion': False,
+      'revision_ttl': '6000s',
+  }
+  if _supports_generate_memories_metadata():
+    expected_config['metadata'] = {'source': {'string_value': 'agent'}}
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_called_once_with(
+      name='reasoningEngines/123',
+      direct_contents_source={
+          'events': [
+              {
+                  'content': {
+                      'parts': [{'text': 'test_content'}],
+                  }
+              }
+          ]
+      },
+      scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
+      config=expected_config,
+  )
+  generate_config = (
+      mock_vertexai_client.agent_engines.memories.generate.call_args.kwargs[
+          'config'
+      ]
+  )
+  vertex_common_types.GenerateAgentEngineMemoriesConfig(**generate_config)
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_without_session_id(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION.app_name,
+      user_id=MOCK_SESSION.user_id,
+      events=[MOCK_SESSION.events[0]],
+  )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_called_once_with(
+      name='reasoningEngines/123',
+      direct_contents_source={
+          'events': [
+              {
+                  'content': {
+                      'parts': [{'text': 'test_content'}],
+                  }
+              }
+          ]
+      },
+      scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
+      config={'wait_for_completion': False},
+  )
+  generate_config = (
+      mock_vertexai_client.agent_engines.memories.generate.call_args.kwargs[
+          'config'
+      ]
+  )
+  vertex_common_types.GenerateAgentEngineMemoriesConfig(**generate_config)
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_merges_metadata_field_and_unknown_keys(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION.app_name,
+      user_id=MOCK_SESSION.user_id,
+      session_id=MOCK_SESSION.id,
+      events=[MOCK_SESSION.events[0]],
+      custom_metadata={
+          'metadata': {'origin': 'unit-test'},
+          'source': 'agent',
+      },
+  )
+
+  expected_config = {'wait_for_completion': False}
+  if _supports_generate_memories_metadata():
+    expected_config['metadata'] = {
+        'origin': {'string_value': 'unit-test'},
+        'source': {'string_value': 'agent'},
+    }
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_called_once_with(
+      name='reasoningEngines/123',
+      direct_contents_source={
+          'events': [
+              {
+                  'content': {
+                      'parts': [{'text': 'test_content'}],
+                  }
+              }
+          ]
+      },
+      scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
+      config=expected_config,
+  )
+  generate_config = (
+      mock_vertexai_client.agent_engines.memories.generate.call_args.kwargs[
+          'config'
+      ]
+  )
+  vertex_common_types.GenerateAgentEngineMemoriesConfig(**generate_config)
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_none_wait_for_completion_keeps_default(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION.app_name,
+      user_id=MOCK_SESSION.user_id,
+      session_id=MOCK_SESSION.id,
+      events=[MOCK_SESSION.events[0]],
+      custom_metadata={'wait_for_completion': None},
+  )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_called_once_with(
+      name='reasoningEngines/123',
+      direct_contents_source={
+          'events': [
+              {
+                  'content': {
+                      'parts': [{'text': 'test_content'}],
+                  }
+              }
+          ]
+      },
+      scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
+      config={'wait_for_completion': False},
+  )
+  generate_config = (
+      mock_vertexai_client.agent_engines.memories.generate.call_args.kwargs[
+          'config'
+      ]
+  )
+  vertex_common_types.GenerateAgentEngineMemoriesConfig(**generate_config)
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_ttl_used_when_revision_ttl_is_none(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION.app_name,
+      user_id=MOCK_SESSION.user_id,
+      session_id=MOCK_SESSION.id,
+      events=[MOCK_SESSION.events[0]],
+      custom_metadata={
+          'ttl': '6000s',
+          'revision_ttl': None,
+      },
+  )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_called_once_with(
+      name='reasoningEngines/123',
+      direct_contents_source={
+          'events': [
+              {
+                  'content': {
+                      'parts': [{'text': 'test_content'}],
+                  }
+              }
+          ]
+      },
+      scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
+      config={
+          'wait_for_completion': False,
+          'revision_ttl': '6000s',
+      },
+  )
+  generate_config = (
+      mock_vertexai_client.agent_engines.memories.generate.call_args.kwargs[
+          'config'
+      ]
+  )
+  vertex_common_types.GenerateAgentEngineMemoriesConfig(**generate_config)
+
+
+@pytest.mark.asyncio
+async def test_add_events_to_memory_with_filtered_events_skips_rpc(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  await memory_service.add_events_to_memory(
+      app_name=MOCK_SESSION.app_name,
+      user_id=MOCK_SESSION.user_id,
+      session_id=MOCK_SESSION.id,
+      events=[MOCK_SESSION.events[1], MOCK_SESSION.events[2]],
+  )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_not_called()
+
+
@pytest.mark.asyncio
 async def test_add_empty_session_to_memory(mock_vertexai_client):
  memory_service = mock_vertex_ai_memory_bank_service()