feat(memory): add Vertex AI load_profiles tool

Add VertexAiLoadProfilesTool for explicit agent access to structured
user profiles from Vertex AI Memory Bank, backed by a new
VertexAiMemoryBankService.retrieve_profiles method. Profiles are a
Vertex backend capability (a scope-keyed lookup), not a memory-interface
concept, so BaseMemoryService is unchanged.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 938264248
This commit is contained in:
George Weale
2026-06-25 16:49:51 -07:00
committed by Copybara-Service
parent 7b87f910cd
commit fb2b3afea1
5 changed files with 290 additions and 27 deletions
@@ -33,6 +33,7 @@ from .memory_entry import MemoryEntry
if TYPE_CHECKING:
import vertexai
from vertexai import types as vertex_types
from ..events.event import Event
from ..sessions.session import Session
@@ -107,22 +108,21 @@ _MAX_DIRECT_MEMORIES_PER_GENERATE_CALL = 5
def _supports_generate_memories_metadata() -> bool:
"""Returns whether installed Vertex SDK supports config.metadata."""
try:
from vertexai._genai.types import common as vertex_common_types
from vertexai import types as vertex_types
except ImportError:
return False
return (
'metadata'
in vertex_common_types.GenerateAgentEngineMemoriesConfig.model_fields
'metadata' in vertex_types.GenerateAgentEngineMemoriesConfig.model_fields
)
def _supports_create_memory_metadata() -> bool:
"""Returns whether installed Vertex SDK supports create config.metadata."""
try:
from vertexai._genai.types import common as vertex_common_types
from vertexai import types as vertex_types
except ImportError:
return False
return 'metadata' in vertex_common_types.AgentEngineMemoryConfig.model_fields
return 'metadata' in vertex_types.AgentEngineMemoryConfig.model_fields
@lru_cache(maxsize=1)
@@ -133,14 +133,12 @@ def _get_generate_memories_config_keys() -> frozenset[str]:
allowlist to preserve compatibility when introspection is unavailable.
"""
try:
from vertexai._genai.types import common as vertex_common_types
from vertexai import types as vertex_types
except ImportError:
return _GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS
try:
model_fields = (
vertex_common_types.GenerateAgentEngineMemoriesConfig.model_fields
)
model_fields = vertex_types.GenerateAgentEngineMemoriesConfig.model_fields
except AttributeError:
return _GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS
@@ -157,12 +155,12 @@ def _get_create_memory_config_keys() -> frozenset[str]:
allowlist to preserve compatibility when introspection is unavailable.
"""
try:
from vertexai._genai.types import common as vertex_common_types
from vertexai import types as vertex_types
except ImportError:
return _CREATE_MEMORY_CONFIG_FALLBACK_KEYS
try:
model_fields = vertex_common_types.AgentEngineMemoryConfig.model_fields
model_fields = vertex_types.AgentEngineMemoryConfig.model_fields
except AttributeError:
return _CREATE_MEMORY_CONFIG_FALLBACK_KEYS
@@ -574,6 +572,39 @@ class VertexAiMemoryBankService(BaseMemoryService):
)
return SearchMemoryResponse(memories=memory_events)
async def retrieve_profiles(
self,
*,
app_name: str,
user_id: str,
) -> list[vertex_types.MemoryProfile]:
"""Retrieves structured user profiles for the scope, one per schema.
Profiles are a Vertex Memory Bank capability distinct from memory search:
a scope-keyed lookup, not a semantic query.
Args:
app_name: The application name for the profile scope.
user_id: The user ID for the profile scope.
Returns:
The structured profiles for the scope, one per registered schema.
"""
api_client = self._get_api_client()
response = await api_client.agent_engines.memories.retrieve_profiles(
name='reasoningEngines/' + self._agent_engine_id,
scope={
'app_name': app_name,
'user_id': user_id,
},
)
profiles = list((response.profiles or {}).values())
if profiles:
logger.info('Retrieved %d memory profiles.', len(profiles))
else:
logger.info('Retrieved no memory profiles.')
return profiles
def _get_api_client(self) -> vertexai.AsyncClient:
"""Instantiates an API client for the given project and location.
+5
View File
@@ -42,6 +42,7 @@ if TYPE_CHECKING:
from .transfer_to_agent_tool import transfer_to_agent
from .transfer_to_agent_tool import TransferToAgentTool
from .url_context_tool import url_context
from .vertex_ai_load_profiles_tool import VertexAiLoadProfilesTool
from .vertex_ai_search_tool import VertexAiSearchTool
# If you are adding a new tool to this file, please make sure you add it to the
@@ -89,6 +90,10 @@ _LAZY_MAPPING = {
'TransferToAgentTool',
),
'url_context': ('.url_context_tool', 'url_context'),
'VertexAiLoadProfilesTool': (
'.vertex_ai_load_profiles_tool',
'VertexAiLoadProfilesTool',
),
'VertexAiSearchTool': ('.vertex_ai_search_tool', 'VertexAiSearchTool'),
'MCPToolset': ('.mcp_tool.mcp_toolset', 'MCPToolset'),
'McpToolset': ('.mcp_tool.mcp_toolset', 'McpToolset'),
@@ -0,0 +1,67 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Any
from typing import TYPE_CHECKING
from google.genai import types
from typing_extensions import override
from ..features import FeatureName
from ..features import is_feature_enabled
from .function_tool import FunctionTool
from .tool_context import ToolContext
if TYPE_CHECKING:
from ..memory.vertex_ai_memory_bank_service import VertexAiMemoryBankService
class VertexAiLoadProfilesTool(FunctionTool):
"""A tool that loads a user's structured profiles from Vertex Memory Bank."""
def __init__(self, memory_service: VertexAiMemoryBankService):
super().__init__(self.load_profiles)
self._memory_service = memory_service
async def load_profiles(self, tool_context: ToolContext) -> dict[str, Any]:
"""Loads structured user profiles for the current user."""
profiles = await self._memory_service.retrieve_profiles(
app_name=tool_context.session.app_name,
user_id=tool_context.user_id,
)
return {
'profiles': [profile.profile for profile in profiles if profile.profile]
}
@override
def _get_declaration(self) -> types.FunctionDeclaration | None:
if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
return types.FunctionDeclaration(
name=self.name,
description=self.description,
parameters_json_schema={
'type': 'object',
'properties': {},
},
)
return types.FunctionDeclaration(
name=self.name,
description=self.description,
parameters=types.Schema(
type=types.Type.OBJECT,
properties={},
),
)
@@ -14,6 +14,7 @@
import asyncio
import datetime
import logging
from typing import Any
from typing import Iterable
from typing import Optional
@@ -26,7 +27,7 @@ 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
from vertexai import types as vertex_types
MOCK_APP_NAME = 'test-app'
MOCK_USER_ID = 'test-user'
@@ -34,20 +35,16 @@ MOCK_USER_ID = 'test-user'
def _supports_generate_memories_metadata() -> bool:
return (
'metadata'
in vertex_common_types.GenerateAgentEngineMemoriesConfig.model_fields
'metadata' in vertex_types.GenerateAgentEngineMemoriesConfig.model_fields
)
def _supports_create_memory_metadata() -> bool:
return 'metadata' in vertex_common_types.AgentEngineMemoryConfig.model_fields
return 'metadata' in vertex_types.AgentEngineMemoryConfig.model_fields
def _supports_create_memory_revision_labels() -> bool:
return (
'revision_labels'
in vertex_common_types.AgentEngineMemoryConfig.model_fields
)
return 'revision_labels' in vertex_types.AgentEngineMemoryConfig.model_fields
class _AsyncListIterator:
@@ -208,6 +205,9 @@ def mock_vertexai_client():
mock_async_client.agent_engines.memories.generate = mock.AsyncMock()
mock_async_client.agent_engines.memories.create = mock.AsyncMock()
mock_async_client.agent_engines.memories.retrieve = mock.AsyncMock()
mock_async_client.agent_engines.memories.retrieve_profiles = (
mock.AsyncMock()
)
mock_async_client.agent_engines.memories.ingest_events = mock.AsyncMock()
mock_client = mock.MagicMock()
@@ -305,7 +305,7 @@ async def test_add_events_to_memory_with_explicit_events_and_metadata(
source = call_kwargs['direct_contents_source']
assert len(source.events) == 1
assert source.events[0].content.parts[0].text == 'test_content'
vertex_common_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
vertex_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
@pytest.mark.asyncio
@@ -336,7 +336,7 @@ async def test_add_events_to_memory_without_session_id(
source = call_kwargs['direct_contents_source']
assert len(source.events) == 1
assert source.events[0].content.parts[0].text == 'test_content'
vertex_common_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
vertex_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
mock_vertexai_client.agent_engines.memories.create.assert_not_called()
@@ -376,7 +376,7 @@ async def test_add_events_to_memory_merges_metadata_field_and_unknown_keys(
source = call_kwargs['direct_contents_source']
assert len(source.events) == 1
assert source.events[0].content.parts[0].text == 'test_content'
vertex_common_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
vertex_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
@pytest.mark.asyncio
@@ -407,7 +407,7 @@ async def test_add_events_to_memory_none_wait_for_completion_keeps_default(
source = call_kwargs['direct_contents_source']
assert len(source.events) == 1
assert source.events[0].content.parts[0].text == 'test_content'
vertex_common_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
vertex_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
@pytest.mark.asyncio
@@ -442,7 +442,7 @@ async def test_add_events_to_memory_ttl_used_when_revision_ttl_is_none(
source = call_kwargs['direct_contents_source']
assert len(source.events) == 1
assert source.events[0].content.parts[0].text == 'test_content'
vertex_common_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
vertex_types.GenerateAgentEngineMemoriesConfig(**call_kwargs['config'])
@pytest.mark.asyncio
@@ -587,7 +587,7 @@ async def test_add_memory_calls_create(
'config'
]
)
vertex_common_types.AgentEngineMemoryConfig(**create_config)
vertex_types.AgentEngineMemoryConfig(**create_config)
@pytest.mark.asyncio
@@ -634,7 +634,7 @@ async def test_add_memory_enable_consolidation_calls_generate_direct_source(
'config'
]
)
vertex_common_types.GenerateAgentEngineMemoriesConfig(**generate_config)
vertex_types.GenerateAgentEngineMemoriesConfig(**generate_config)
@pytest.mark.asyncio
@@ -768,7 +768,7 @@ async def test_add_memory_calls_create_with_memory_entry_metadata(
'config'
]
)
vertex_common_types.AgentEngineMemoryConfig(**create_config)
vertex_types.AgentEngineMemoryConfig(**create_config)
@pytest.mark.asyncio
@@ -1009,6 +1009,66 @@ async def test_search_memory_empty_results(mock_vertexai_client):
assert len(result.memories) == 0
@pytest.mark.asyncio
async def test_retrieve_profiles(mock_vertexai_client, caplog):
"""Returns the structured profiles for the scope as a list."""
retrieve_profiles_response = vertex_types.RetrieveProfilesResponse(
profiles={
'user-profile': vertex_types.MemoryProfile(
schema_id='user-profile',
profile={'name': 'Kim'},
)
}
)
mock_vertexai_client.agent_engines.memories.retrieve_profiles.return_value = (
retrieve_profiles_response
)
memory_service = mock_vertex_ai_memory_bank_service()
with caplog.at_level(logging.INFO):
result = await memory_service.retrieve_profiles(
app_name=MOCK_APP_NAME,
user_id=MOCK_USER_ID,
)
mock_vertexai_client.agent_engines.memories.retrieve_profiles.assert_awaited_once_with(
name='reasoningEngines/123',
scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
)
assert 'Retrieved 1 memory profiles.' in caplog.text
assert result == [
vertex_types.MemoryProfile(
schema_id='user-profile',
profile={'name': 'Kim'},
)
]
@pytest.mark.asyncio
async def test_retrieve_profiles_empty_results(mock_vertexai_client, caplog):
"""Returns an empty list when the scope has no profiles."""
retrieve_profiles_response = vertex_types.RetrieveProfilesResponse(
profiles=None
)
mock_vertexai_client.agent_engines.memories.retrieve_profiles.return_value = (
retrieve_profiles_response
)
memory_service = mock_vertex_ai_memory_bank_service()
with caplog.at_level(logging.INFO):
result = await memory_service.retrieve_profiles(
app_name=MOCK_APP_NAME,
user_id=MOCK_USER_ID,
)
mock_vertexai_client.agent_engines.memories.retrieve_profiles.assert_awaited_once_with(
name='reasoningEngines/123',
scope={'app_name': MOCK_APP_NAME, 'user_id': MOCK_USER_ID},
)
assert 'Retrieved no memory profiles.' in caplog.text
assert not result
async def test_search_memory_uses_async_client_path():
sync_client = mock.MagicMock()
sync_client.agent_engines.memories.retrieve.side_effect = AssertionError(
@@ -0,0 +1,100 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from types import SimpleNamespace
from google.adk.features import FeatureName
from google.adk.features._feature_registry import temporary_feature_override
from google.adk.models.llm_request import LlmRequest
from google.adk.tools.vertex_ai_load_profiles_tool import VertexAiLoadProfilesTool
from pytest import mark
from vertexai import types as vertex_types
class _FakeMemoryService:
"""Minimal profile-providing service for VertexAiLoadProfilesTool tests."""
def __init__(self, profiles):
self._profiles = profiles
self.calls = []
async def retrieve_profiles(self, *, app_name, user_id):
self.calls.append((app_name, user_id))
return self._profiles
class _StubToolContext:
"""Minimal ToolContext stub exposing only the scope the tool reads."""
def __init__(self, *, app_name='test-app', user_id='test-user'):
self.session = SimpleNamespace(app_name=app_name)
self.user_id = user_id
@mark.asyncio
async def test_load_profiles_returns_profile_payloads():
memory_service = _FakeMemoryService([
vertex_types.MemoryProfile(
schema_id='user-profile', profile={'name': 'Kim'}
),
vertex_types.MemoryProfile(schema_id='empty', profile={}),
])
tool = VertexAiLoadProfilesTool(memory_service=memory_service)
result = await tool.load_profiles(_StubToolContext())
assert result == {'profiles': [{'name': 'Kim'}]}
assert memory_service.calls == [('test-app', 'test-user')]
def test_get_declaration_with_json_schema_feature_disabled():
tool = VertexAiLoadProfilesTool(memory_service=_FakeMemoryService([]))
with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False):
declaration = tool._get_declaration()
assert declaration.name == 'load_profiles'
assert declaration.parameters_json_schema is None
assert declaration.parameters.properties == {}
def test_get_declaration_with_json_schema_feature_enabled():
tool = VertexAiLoadProfilesTool(memory_service=_FakeMemoryService([]))
with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True):
declaration = tool._get_declaration()
assert declaration.name == 'load_profiles'
assert declaration.parameters is None
assert declaration.parameters_json_schema == {
'type': 'object',
'properties': {},
}
@mark.asyncio
async def test_process_llm_request_registers_tool_only():
tool = VertexAiLoadProfilesTool(memory_service=_FakeMemoryService([]))
llm_request = LlmRequest()
await tool.process_llm_request(
tool_context=_StubToolContext(),
llm_request=llm_request,
)
assert llm_request.config.system_instruction is None
assert llm_request.config.tools is not None
assert llm_request.config.tools[0].function_declarations is not None
assert llm_request.config.tools[0].function_declarations[0].name == (
'load_profiles'
)
assert 'load_profiles' in llm_request.tools_dict