diff --git a/src/google/adk/memory/vertex_ai_memory_bank_service.py b/src/google/adk/memory/vertex_ai_memory_bank_service.py index e52ef99d..af949fb2 100644 --- a/src/google/adk/memory/vertex_ai_memory_bank_service.py +++ b/src/google/adk/memory/vertex_ai_memory_bank_service.py @@ -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. diff --git a/src/google/adk/tools/__init__.py b/src/google/adk/tools/__init__.py index 770a2cb3..641c5c20 100644 --- a/src/google/adk/tools/__init__.py +++ b/src/google/adk/tools/__init__.py @@ -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'), diff --git a/src/google/adk/tools/vertex_ai_load_profiles_tool.py b/src/google/adk/tools/vertex_ai_load_profiles_tool.py new file mode 100644 index 00000000..26d8b87b --- /dev/null +++ b/src/google/adk/tools/vertex_ai_load_profiles_tool.py @@ -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={}, + ), + ) diff --git a/tests/unittests/memory/test_vertex_ai_memory_bank_service.py b/tests/unittests/memory/test_vertex_ai_memory_bank_service.py index e41e797f..529c3aec 100644 --- a/tests/unittests/memory/test_vertex_ai_memory_bank_service.py +++ b/tests/unittests/memory/test_vertex_ai_memory_bank_service.py @@ -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( diff --git a/tests/unittests/tools/test_vertex_ai_load_profiles_tool.py b/tests/unittests/tools/test_vertex_ai_load_profiles_tool.py new file mode 100644 index 00000000..81e4e308 --- /dev/null +++ b/tests/unittests/tools/test_vertex_ai_load_profiles_tool.py @@ -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