fix: Expand add_memory to accept MemoryEntry

The `add_memory` methods in `Context` and `BaseMemoryService` now accept `MemoryEntry` objects in addition to strings. The Vertex AI Memory Bank service implementation is updated to handle these new types Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 872108561
2026-07-09 18:19:28 -07:00 · 2026-02-18 17:06:17 -08:00
parent 2d8b6a2f5b
commit f27a9cfb87
6 changed files with 479 additions and 60 deletions
@@ -22,7 +22,9 @@ from google.adk.agents.callback_context import CallbackContext
 from google.adk.auth.auth_credential import AuthCredential
 from google.adk.auth.auth_credential import AuthCredentialTypes
 from google.adk.auth.auth_tool import AuthConfig
+from google.adk.memory.memory_entry import MemoryEntry
 from google.adk.tools.tool_context import ToolContext
+from google.genai import types
 from google.genai.types import Part
 import pytest

@@ -417,7 +419,9 @@ class TestCallbackContextAddEventsToMemory:
    """Tests that add_memory forwards memories and metadata."""
    memory_service = AsyncMock()
    mock_invocation_context.memory_service = memory_service
-    memories = ["fact one"]
+    memories = [
+        MemoryEntry(content=types.Content(parts=[types.Part(text="fact one")]))
+    ]
    metadata = {"ttl": "6000s"}

    context = CallbackContext(mock_invocation_context)
@@ -430,6 +434,27 @@ class TestCallbackContextAddEventsToMemory:
        custom_metadata=metadata,
    )

+  @pytest.mark.asyncio
+  async def test_add_memory_accepts_memory_entries(
+      self, mock_invocation_context
+  ):
+    """Tests that add_memory forwards MemoryEntry inputs unchanged."""
+    memory_service = AsyncMock()
+    mock_invocation_context.memory_service = memory_service
+    memory_entry = MemoryEntry(
+        content=types.Content(parts=[types.Part(text="fact one")])
+    )
+
+    context = CallbackContext(mock_invocation_context)
+    await context.add_memory(memories=[memory_entry])
+
+    memory_service.add_memory.assert_called_once_with(
+        app_name=mock_invocation_context.session.app_name,
+        user_id=mock_invocation_context.session.user_id,
+        memories=[memory_entry],
+        custom_metadata=None,
+    )
+
  @pytest.mark.asyncio
  async def test_add_memory_no_service_raises(self, mock_invocation_context):
    """Tests that add_memory raises ValueError with no service."""
@@ -441,7 +466,13 @@ class TestCallbackContextAddEventsToMemory:
        ValueError,
        match=r"Cannot add memory: memory service is not available\.",
    ):
-      await context.add_memory(memories=["fact one"])
+      await context.add_memory(
+          memories=[
+              MemoryEntry(
+                  content=types.Content(parts=[types.Part(text="fact one")])
+              )
+          ]
+      )


 class TestToolContextAddSessionToMemory:
@@ -22,7 +22,9 @@ from google.adk.auth.auth_credential import AuthCredential
 from google.adk.auth.auth_credential import AuthCredentialTypes
 from google.adk.auth.auth_tool import AuthConfig
 from google.adk.memory.base_memory_service import SearchMemoryResponse
+from google.adk.memory.memory_entry import MemoryEntry
 from google.adk.tools.tool_confirmation import ToolConfirmation
+from google.genai import types
 from google.genai.types import Part
 import pytest

@@ -492,7 +494,9 @@ class TestContextMemoryMethods:
    """Tests that add_memory forwards memories and metadata."""
    memory_service = AsyncMock()
    mock_invocation_context.memory_service = memory_service
-    memories = ["fact one"]
+    memories = [
+        MemoryEntry(content=types.Content(parts=[types.Part(text="fact one")]))
+    ]
    metadata = {"ttl": "6000s"}

    context = Context(mock_invocation_context)
@@ -505,6 +509,27 @@ class TestContextMemoryMethods:
        custom_metadata=metadata,
    )

+  @pytest.mark.asyncio
+  async def test_add_memory_accepts_memory_entries(
+      self, mock_invocation_context
+  ):
+    """Tests that add_memory forwards MemoryEntry inputs unchanged."""
+    memory_service = AsyncMock()
+    mock_invocation_context.memory_service = memory_service
+    memory_entry = MemoryEntry(
+        content=types.Content(parts=[types.Part(text="fact one")])
+    )
+
+    context = Context(mock_invocation_context)
+    await context.add_memory(memories=[memory_entry])
+
+    memory_service.add_memory.assert_called_once_with(
+        app_name=mock_invocation_context.session.app_name,
+        user_id=mock_invocation_context.session.user_id,
+        memories=[memory_entry],
+        custom_metadata=None,
+    )
+
  async def test_add_memory_no_service_raises(self, mock_invocation_context):
    """Test that add_memory raises ValueError when no service."""
    mock_invocation_context.memory_service = None
@@ -515,4 +540,10 @@ class TestContextMemoryMethods:
        ValueError,
        match=r"Cannot add memory: memory service is not available\.",
    ):
-      await context.add_memory(memories=["fact one"])
+      await context.add_memory(
+          memories=[
+              MemoryEntry(
+                  content=types.Content(parts=[types.Part(text="fact one")])
+              )
+          ]
+      )
@@ -20,6 +20,7 @@ from unittest import mock

 from google.adk.events.event import Event
 from google.adk.memory import vertex_ai_memory_bank_service as memory_service_module
+from google.adk.memory.memory_entry import MemoryEntry
 from google.adk.memory.vertex_ai_memory_bank_service import VertexAiMemoryBankService
 from google.adk.sessions.session import Session
 from google.genai import types
@@ -41,6 +42,13 @@ def _supports_create_memory_metadata() -> bool:
  return 'metadata' in vertex_common_types.AgentEngineMemoryConfig.model_fields


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

@@ -165,6 +173,33 @@ def test_build_create_memory_config_uses_runtime_config_keys():
  }


+def test_build_create_memory_config_merges_revision_labels_when_supported():
+  with (
+      mock.patch.object(
+          memory_service_module,
+          '_get_create_memory_config_keys',
+          return_value=frozenset({'wait_for_completion', 'revision_labels'}),
+      ),
+      mock.patch.object(
+          memory_service_module,
+          '_supports_create_memory_metadata',
+          return_value=False,
+      ),
+  ):
+    config = memory_service_module._build_create_memory_config(
+        {'revision_labels': {'source': 'global'}},
+        memory_revision_labels={'author': 'agent'},
+    )
+
+  assert config == {
+      'wait_for_completion': False,
+      'revision_labels': {
+          'source': 'global',
+          'author': 'agent',
+      },
+  }
+
+
@pytest.fixture
 def mock_vertexai_client():
  with mock.patch('vertexai.Client') as mock_client_constructor:
@@ -437,7 +472,14 @@ async def test_add_memory_calls_create(
  await memory_service.add_memory(
      app_name=MOCK_SESSION.app_name,
      user_id=MOCK_SESSION.user_id,
-      memories=['fact one', 'fact two'],
+      memories=[
+          MemoryEntry(
+              content=types.Content(parts=[types.Part(text='fact one')])
+          ),
+          MemoryEntry(
+              content=types.Content(parts=[types.Part(text='fact two')])
+          ),
+      ],
      custom_metadata={
          'ttl': '6000s',
          'source': 'agent',
@@ -476,6 +518,176 @@ async def test_add_memory_calls_create(
  vertex_common_types.AgentEngineMemoryConfig(**create_config)


+@pytest.mark.asyncio
+async def test_add_memory_calls_create_with_memory_entry_metadata(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  await memory_service.add_memory(
+      app_name=MOCK_SESSION.app_name,
+      user_id=MOCK_SESSION.user_id,
+      memories=[
+          MemoryEntry(
+              author='agent',
+              timestamp='2026-02-13T14:46:21Z',
+              content=types.Content(parts=[types.Part(text='fact one')]),
+              custom_metadata={'source': 'entry'},
+          )
+      ],
+      custom_metadata={'ttl': '6000s', 'source': 'global'},
+  )
+
+  expected_config = {
+      'wait_for_completion': False,
+      'ttl': '6000s',
+  }
+  if _supports_create_memory_metadata():
+    expected_config['metadata'] = {
+        'source': {'string_value': 'entry'},
+    }
+  if _supports_create_memory_revision_labels():
+    expected_config['revision_labels'] = {
+        'author': 'agent',
+        'timestamp': '2026-02-13T14:46:21Z',
+    }
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_not_called()
+  mock_vertexai_client.agent_engines.memories.create.assert_awaited_once_with(
+      name='reasoningEngines/123',
+      fact='fact one',
+      scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
+      config=expected_config,
+  )
+  create_config = (
+      mock_vertexai_client.agent_engines.memories.create.call_args.kwargs[
+          'config'
+      ]
+  )
+  vertex_common_types.AgentEngineMemoryConfig(**create_config)
+
+
+@pytest.mark.asyncio
+async def test_add_memory_calls_create_with_multimodal_content(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  with pytest.raises(
+      ValueError,
+      match=(
+          r'memories\[0\] must include text only; inline_data and file_data '
+          r'are not supported'
+      ),
+  ):
+    await memory_service.add_memory(
+        app_name=MOCK_SESSION.app_name,
+        user_id=MOCK_SESSION.user_id,
+        memories=[
+            MemoryEntry(
+                content=types.Content(
+                    parts=[
+                        types.Part(text='caption'),
+                        types.Part(
+                            file_data=types.FileData(
+                                mime_type='image/png',
+                                file_uri='gs://bucket/image.png',
+                            )
+                        ),
+                    ]
+                )
+            )
+        ],
+    )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_not_called()
+  mock_vertexai_client.agent_engines.memories.create.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_add_memory_with_missing_text_raises(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  with pytest.raises(
+      ValueError,
+      match=r'memories\[0\] must include text',
+  ):
+    await memory_service.add_memory(
+        app_name=MOCK_SESSION.app_name,
+        user_id=MOCK_SESSION.user_id,
+        memories=[
+            MemoryEntry(
+                content=types.Content(
+                    parts=[
+                        types.Part(
+                            function_call=types.FunctionCall(name='tool')
+                        )
+                    ]
+                )
+            )
+        ],
+    )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_not_called()
+  mock_vertexai_client.agent_engines.memories.create.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_add_memory_with_whitespace_only_text_raises(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  with pytest.raises(
+      ValueError,
+      match=r'memories\[0\] must include non-whitespace text',
+  ):
+    await memory_service.add_memory(
+        app_name=MOCK_SESSION.app_name,
+        user_id=MOCK_SESSION.user_id,
+        memories=[
+            MemoryEntry(content=types.Content(parts=[types.Part(text='   ')]))
+        ],
+    )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_not_called()
+  mock_vertexai_client.agent_engines.memories.create.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_add_memory_with_whitespace_and_non_text_parts_raises(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  with pytest.raises(
+      ValueError,
+      match=(
+          r'memories\[0\] must include text only; inline_data and file_data '
+          r'are not supported'
+      ),
+  ):
+    await memory_service.add_memory(
+        app_name=MOCK_SESSION.app_name,
+        user_id=MOCK_SESSION.user_id,
+        memories=[
+            MemoryEntry(
+                content=types.Content(
+                    parts=[
+                        types.Part(text='  '),
+                        types.Part(
+                            inline_data=types.Blob(
+                                mime_type='image/png',
+                                data=b'abc',
+                            )
+                        ),
+                    ]
+                )
+            )
+        ],
+    )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_not_called()
+  mock_vertexai_client.agent_engines.memories.create.assert_not_called()
+
+
@pytest.mark.asyncio
 async def test_add_memory_missing_memories_raises(
    mock_vertexai_client,
@@ -498,7 +710,10 @@ async def test_add_memory_with_invalid_memory_type_raises(
    mock_vertexai_client,
 ):
  memory_service = mock_vertex_ai_memory_bank_service()
-  with pytest.raises(TypeError, match=r'memories\[0\] must be a string'):
+  with pytest.raises(
+      TypeError,
+      match=r'memories\[0\] must be a MemoryEntry',
+  ):
    await memory_service.add_memory(
        app_name=MOCK_SESSION.app_name,
        user_id=MOCK_SESSION.user_id,
@@ -508,6 +723,25 @@ async def test_add_memory_with_invalid_memory_type_raises(
  mock_vertexai_client.agent_engines.memories.create.assert_not_called()


+@pytest.mark.asyncio
+async def test_add_memory_with_content_type_raises(
+    mock_vertexai_client,
+):
+  memory_service = mock_vertex_ai_memory_bank_service()
+  with pytest.raises(
+      TypeError,
+      match=r'memories\[0\] must be a MemoryEntry',
+  ):
+    await memory_service.add_memory(
+        app_name=MOCK_SESSION.app_name,
+        user_id=MOCK_SESSION.user_id,
+        memories=[types.Content(parts=[types.Part(text='fact one')])],
+    )
+
+  mock_vertexai_client.agent_engines.memories.generate.assert_not_called()
+  mock_vertexai_client.agent_engines.memories.create.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()