feat: add error handling to VertexAiMemoryBankService.search_memory

Author: vietnamesekid

src/google/adk/memory/vertex_ai_memory_bank_service.py

@@ -122,19 +122,38 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str):
 
     logger.info('Search memory response received.')
 
-    memory_events = []
-    for retrieved_memory in retrieved_memories_iterator:
-      # TODO: add more complex error handling
-      logger.debug('Retrieved memory: %s', retrieved_memory)
-      memory_events.append(
-          MemoryEntry(
-              author='user',
-              content=types.Content(
-                  parts=[types.Part(text=retrieved_memory.memory.fact)],
-                  role='user',
-              ),
-              timestamp=retrieved_memory.memory.update_time.isoformat(),
+    memory_events: list[MemoryEntry] = []
+    try:
+      for retrieved_memory in retrieved_memories_iterator:
+        try:
+          memory = retrieved_memory.memory
+          if memory is None:
+            logger.warning('Skipping memory entry with missing memory object.')
+            continue
+          fact = memory.fact
+          if not fact:
+            logger.warning('Skipping memory entry with empty or missing fact.')
+            continue
+
+          update_time = memory.update_time
+          memory_events.append(
+              MemoryEntry(
+                  author='user',
+                  content=types.Content(
+                      parts=[types.Part(text=fact)],
+                      role='user',
+                  ),
+                  timestamp=update_time.isoformat() if update_time else None,
+              )
           )
+        except AttributeError:
+          logger.warning(
+              'Skipping malformed memory entry: %s', retrieved_memory
+          )
+    except Exception:  # pylint: disable=broad-except
+      logger.exception(
+          'Error while iterating memory results. Returning %d partial results.',
+          len(memory_events),
       )
     return SearchMemoryResponse(memories=memory_events)
 

tests/unittests/memory/test_vertex_ai_memory_bank_service.py

@@ -182,3 +182,110 @@ async def test_search_memory_empty_results(mock_vertexai_client):
   )
 
   assert len(result.memories) == 0
+
+
+@pytest.mark.asyncio
+async def test_search_memory_skips_entry_with_none_memory(mock_vertexai_client):
+  bad_entry = mock.MagicMock()
+  bad_entry.memory = None
+
+  good_entry = mock.MagicMock()
+  good_entry.memory.fact = 'good fact'
+  good_entry.memory.update_time = datetime(2024, 1, 1)
+
+  mock_vertexai_client.agent_engines.memories.retrieve.return_value = [
+      bad_entry,
+      good_entry,
+  ]
+  memory_service = mock_vertex_ai_memory_bank_service()
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='query'
+  )
+
+  assert len(result.memories) == 1
+  assert result.memories[0].content.parts[0].text == 'good fact'
+
+
+@pytest.mark.asyncio
+async def test_search_memory_skips_entry_with_empty_fact(mock_vertexai_client):
+  for empty_fact in [None, '']:
+    bad_entry = mock.MagicMock()
+    bad_entry.memory.fact = empty_fact
+    bad_entry.memory.update_time = datetime(2024, 1, 1)
+
+    mock_vertexai_client.agent_engines.memories.retrieve.return_value = [
+        bad_entry
+    ]
+    memory_service = mock_vertex_ai_memory_bank_service()
+
+    result = await memory_service.search_memory(
+        app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='query'
+    )
+
+    assert len(result.memories) == 0
+
+
+@pytest.mark.asyncio
+async def test_search_memory_handles_missing_update_time(mock_vertexai_client):
+  entry = mock.MagicMock()
+  entry.memory.fact = 'some fact'
+  entry.memory.update_time = None
+
+  mock_vertexai_client.agent_engines.memories.retrieve.return_value = [entry]
+  memory_service = mock_vertex_ai_memory_bank_service()
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='query'
+  )
+
+  assert len(result.memories) == 1
+  assert result.memories[0].content.parts[0].text == 'some fact'
+  assert result.memories[0].timestamp is None
+
+
+@pytest.mark.asyncio
+async def test_search_memory_skips_malformed_entry(mock_vertexai_client):
+  malformed = mock.MagicMock(spec=[])  # no attributes → AttributeError
+
+  good_entry = mock.MagicMock()
+  good_entry.memory.fact = 'good fact'
+  good_entry.memory.update_time = datetime(2024, 1, 1)
+
+  mock_vertexai_client.agent_engines.memories.retrieve.return_value = [
+      malformed,
+      good_entry,
+  ]
+  memory_service = mock_vertex_ai_memory_bank_service()
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='query'
+  )
+
+  assert len(result.memories) == 1
+  assert result.memories[0].content.parts[0].text == 'good fact'
+
+
+@pytest.mark.asyncio
+async def test_search_memory_returns_partial_results_on_iterator_error(
+    mock_vertexai_client,
+):
+  good_entry = mock.MagicMock()
+  good_entry.memory.fact = 'good fact'
+  good_entry.memory.update_time = datetime(2024, 1, 1)
+
+  def failing_iterator():
+    yield good_entry
+    raise RuntimeError('API stream error')
+
+  mock_vertexai_client.agent_engines.memories.retrieve.return_value = (
+      failing_iterator()
+  )
+  memory_service = mock_vertex_ai_memory_bank_service()
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='query'
+  )
+
+  assert len(result.memories) == 1
+  assert result.memories[0].content.parts[0].text == 'good fact'