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:
committed by
Copybara-Service
parent
7b87f910cd
commit
fb2b3afea1
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user