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
This commit is contained in:
George Weale
2026-02-18 17:06:17 -08:00
committed by Copybara-Service
parent 2d8b6a2f5b
commit f27a9cfb87
6 changed files with 479 additions and 60 deletions
+33 -2
View File
@@ -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")])
)
]
)