Skip to content

Commit e32b2e8

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add retrieve profiles
PiperOrigin-RevId: 893649750
1 parent 09794ba commit e32b2e8

File tree

4 files changed

+568
-0
lines changed

4 files changed

+568
-0
lines changed

‎tests/unit/vertexai/genai/replays/test_structured_memories.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
16+
17+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
20+
21+
def test_generate_and_retrieve_profile(client):
22+
# TODO: Use prod once available.
23+
client._api_client._http_options.base_url = (
24+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com"
25+
)
26+
customization_config = {"disable_natural_language_memories": True}
27+
memory_bank_customization_config = types.MemoryBankCustomizationConfig(
28+
**customization_config
29+
)
30+
structured_memory_config = {
31+
"scope_keys": ["user_id"],
32+
"schema_configs": [
33+
{
34+
"id": "user-profile",
35+
"schema": {
36+
"properties": {
37+
"name": {"description": "User's name", "type": "string"}
38+
},
39+
"type": "object",
40+
},
41+
}
42+
],
43+
}
44+
structured_memory_config_obj = types.StructuredMemoryConfig(
45+
**structured_memory_config
46+
)
47+
agent_engine = client.agent_engines.create(
48+
config={
49+
"context_spec": {
50+
"memory_bank_config": {
51+
"customization_configs": [memory_bank_customization_config],
52+
"structured_memory_configs": [structured_memory_config_obj],
53+
},
54+
},
55+
"http_options": {"api_version": "v1beta1"},
56+
},
57+
)
58+
try:
59+
agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name)
60+
memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config
61+
assert memory_bank_config.customization_configs == [
62+
memory_bank_customization_config
63+
]
64+
assert memory_bank_config.structured_memory_configs == [
65+
structured_memory_config_obj
66+
]
67+
68+
scope = {"user_id": "123"}
69+
client.agent_engines.memories.generate(
70+
name=agent_engine.api_resource.name,
71+
scope=scope,
72+
direct_contents_source={
73+
"events": [{"content": {"parts": [{"text": "My name is Kim."}]}}]
74+
},
75+
)
76+
memories = list(
77+
client.agent_engines.memories.retrieve(
78+
name=agent_engine.api_resource.name,
79+
scope=scope,
80+
config={"memory_types": ["STRUCTURED_PROFILE"]},
81+
)
82+
)
83+
assert len(memories) >= 1
84+
assert memories[0].memory.structured_content is not None
85+
86+
response = client.agent_engines.memories.retrieve_profiles(
87+
name=agent_engine.api_resource.name, scope=scope
88+
)
89+
assert len(response.profiles) == 1
90+
91+
finally:
92+
# Clean up resources.
93+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
94+
95+
96+
pytestmark = pytest_helper.setup(
97+
file=__file__,
98+
globals_for_file=globals(),
99+
test_method="agent_engines.retrieve_profiles",
100+
)

‎vertexai/_genai/memories.py‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,9 @@ def _RetrieveAgentEngineMemoriesConfig_to_vertex(
329329
[item for item in getv(from_object, ["filter_groups"])],
330330
)
331331

332+
if getv(from_object, ["memory_types"]) is not None:
333+
setv(parent_object, ["memoryTypes"], getv(from_object, ["memory_types"]))
334+
332335
return to_object
333336

334337

@@ -365,6 +368,20 @@ def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex(
365368
return to_object
366369

367370

371+
def _RetrieveMemoryProfilesRequestParameters_to_vertex(
372+
from_object: Union[dict[str, Any], object],
373+
parent_object: Optional[dict[str, Any]] = None,
374+
) -> dict[str, Any]:
375+
to_object: dict[str, Any] = {}
376+
if getv(from_object, ["name"]) is not None:
377+
setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
378+
379+
if getv(from_object, ["scope"]) is not None:
380+
setv(to_object, ["scope"], getv(from_object, ["scope"]))
381+
382+
return to_object
383+
384+
368385
def _RollbackAgentEngineMemoryRequestParameters_to_vertex(
369386
from_object: Union[dict[str, Any], object],
370387
parent_object: Optional[dict[str, Any]] = None,
@@ -935,6 +952,72 @@ def _retrieve(
935952
self._api_client._verify_response(return_value)
936953
return return_value
937954

955+
def retrieve_profiles(
956+
self,
957+
*,
958+
name: str,
959+
scope: dict[str, str],
960+
config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None,
961+
) -> types.RetrieveProfilesResponse:
962+
"""
963+
Retrieves memory profiles for an Agent Engine.
964+
965+
Args:
966+
name (str): Required. A fully-qualified resource name or ID such as
967+
"projects/123/locations/us-central1/reasoningEngines/456".
968+
scope (dict[str, str]): Required. The scope of the memories to retrieve.
969+
A memory must have exactly the same scope as the scope provided here
970+
to be retrieved (i.e. same keys and values). Order does not matter,
971+
but it is case-sensitive.
972+
973+
"""
974+
975+
parameter_model = types._RetrieveMemoryProfilesRequestParameters(
976+
name=name,
977+
scope=scope,
978+
config=config,
979+
)
980+
981+
request_url_dict: Optional[dict[str, str]]
982+
if not self._api_client.vertexai:
983+
raise ValueError("This method is only supported in the Vertex AI client.")
984+
else:
985+
request_dict = _RetrieveMemoryProfilesRequestParameters_to_vertex(
986+
parameter_model
987+
)
988+
request_url_dict = request_dict.get("_url")
989+
if request_url_dict:
990+
path = "{name}/memories:retrieveProfiles".format_map(request_url_dict)
991+
else:
992+
path = "{name}/memories:retrieveProfiles"
993+
994+
query_params = request_dict.get("_query")
995+
if query_params:
996+
path = f"{path}?{urlencode(query_params)}"
997+
# TODO: remove the hack that pops config.
998+
request_dict.pop("config", None)
999+
1000+
http_options: Optional[types.HttpOptions] = None
1001+
if (
1002+
parameter_model.config is not None
1003+
and parameter_model.config.http_options is not None
1004+
):
1005+
http_options = parameter_model.config.http_options
1006+
1007+
request_dict = _common.convert_to_dict(request_dict)
1008+
request_dict = _common.encode_unserializable_types(request_dict)
1009+
1010+
response = self._api_client.request("post", path, request_dict, http_options)
1011+
1012+
response_dict = {} if not response.body else json.loads(response.body)
1013+
1014+
return_value = types.RetrieveProfilesResponse._from_response(
1015+
response=response_dict, kwargs=parameter_model.model_dump()
1016+
)
1017+
1018+
self._api_client._verify_response(return_value)
1019+
return return_value
1020+
9381021
def _rollback(
9391022
self,
9401023
*,
@@ -1909,6 +1992,74 @@ async def _retrieve(
19091992
self._api_client._verify_response(return_value)
19101993
return return_value
19111994

1995+
async def retrieve_profiles(
1996+
self,
1997+
*,
1998+
name: str,
1999+
scope: dict[str, str],
2000+
config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None,
2001+
) -> types.RetrieveProfilesResponse:
2002+
"""
2003+
Retrieves memory profiles for an Agent Engine.
2004+
2005+
Args:
2006+
name (str): Required. A fully-qualified resource name or ID such as
2007+
"projects/123/locations/us-central1/reasoningEngines/456".
2008+
scope (dict[str, str]): Required. The scope of the memories to retrieve.
2009+
A memory must have exactly the same scope as the scope provided here
2010+
to be retrieved (i.e. same keys and values). Order does not matter,
2011+
but it is case-sensitive.
2012+
2013+
"""
2014+
2015+
parameter_model = types._RetrieveMemoryProfilesRequestParameters(
2016+
name=name,
2017+
scope=scope,
2018+
config=config,
2019+
)
2020+
2021+
request_url_dict: Optional[dict[str, str]]
2022+
if not self._api_client.vertexai:
2023+
raise ValueError("This method is only supported in the Vertex AI client.")
2024+
else:
2025+
request_dict = _RetrieveMemoryProfilesRequestParameters_to_vertex(
2026+
parameter_model
2027+
)
2028+
request_url_dict = request_dict.get("_url")
2029+
if request_url_dict:
2030+
path = "{name}/memories:retrieveProfiles".format_map(request_url_dict)
2031+
else:
2032+
path = "{name}/memories:retrieveProfiles"
2033+
2034+
query_params = request_dict.get("_query")
2035+
if query_params:
2036+
path = f"{path}?{urlencode(query_params)}"
2037+
# TODO: remove the hack that pops config.
2038+
request_dict.pop("config", None)
2039+
2040+
http_options: Optional[types.HttpOptions] = None
2041+
if (
2042+
parameter_model.config is not None
2043+
and parameter_model.config.http_options is not None
2044+
):
2045+
http_options = parameter_model.config.http_options
2046+
2047+
request_dict = _common.convert_to_dict(request_dict)
2048+
request_dict = _common.encode_unserializable_types(request_dict)
2049+
2050+
response = await self._api_client.async_request(
2051+
"post", path, request_dict, http_options
2052+
)
2053+
2054+
response_dict = {} if not response.body else json.loads(response.body)
2055+
2056+
return_value = types.RetrieveProfilesResponse._from_response(
2057+
response=response_dict, kwargs=parameter_model.model_dump()
2058+
)
2059+
2060+
self._api_client._verify_response(return_value)
2061+
return return_value
2062+
19122063
async def _rollback(
19132064
self,
19142065
*,

‎vertexai/_genai/types/__init__.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@
9494
from .common import _QueryAgentEngineRequestParameters
9595
from .common import _RestoreVersionRequestParameters
9696
from .common import _RetrieveAgentEngineMemoriesRequestParameters
97+
from .common import _RetrieveMemoryProfilesRequestParameters
9798
from .common import _RollbackAgentEngineMemoryRequestParameters
9899
from .common import _RunQueryJobAgentEngineConfig
99100
from .common import _RunQueryJobAgentEngineConfigDict
@@ -689,12 +690,19 @@
689690
from .common import MemoryMetadataValueDict
690691
from .common import MemoryMetadataValueOrDict
691692
from .common import MemoryOrDict
693+
from .common import MemoryProfile
694+
from .common import MemoryProfileDict
695+
from .common import MemoryProfileOrDict
692696
from .common import MemoryRevision
693697
from .common import MemoryRevisionDict
694698
from .common import MemoryRevisionOrDict
699+
from .common import MemoryStructuredContent
700+
from .common import MemoryStructuredContentDict
701+
from .common import MemoryStructuredContentOrDict
695702
from .common import MemoryTopicId
696703
from .common import MemoryTopicIdDict
697704
from .common import MemoryTopicIdOrDict
705+
from .common import MemoryType
698706
from .common import Message
699707
from .common import MessageDict
700708
from .common import Metadata
@@ -886,6 +894,12 @@
886894
from .common import RetrieveMemoriesResponseRetrievedMemory
887895
from .common import RetrieveMemoriesResponseRetrievedMemoryDict
888896
from .common import RetrieveMemoriesResponseRetrievedMemoryOrDict
897+
from .common import RetrieveMemoryProfilesConfig
898+
from .common import RetrieveMemoryProfilesConfigDict
899+
from .common import RetrieveMemoryProfilesConfigOrDict
900+
from .common import RetrieveProfilesResponse
901+
from .common import RetrieveProfilesResponseDict
902+
from .common import RetrieveProfilesResponseOrDict
889903
from .common import RollbackAgentEngineMemoryConfig
890904
from .common import RollbackAgentEngineMemoryConfigDict
891905
from .common import RollbackAgentEngineMemoryConfigOrDict
@@ -1041,6 +1055,12 @@
10411055
from .common import SessionOrDict
10421056
from .common import State
10431057
from .common import Strategy
1058+
from .common import StructuredMemoryConfig
1059+
from .common import StructuredMemoryConfigDict
1060+
from .common import StructuredMemoryConfigOrDict
1061+
from .common import StructuredMemoryConfigSchemaConfig
1062+
from .common import StructuredMemoryConfigSchemaConfigDict
1063+
from .common import StructuredMemoryConfigSchemaConfigOrDict
10441064
from .common import SummaryMetric
10451065
from .common import SummaryMetricDict
10461066
from .common import SummaryMetricOrDict
@@ -1653,6 +1673,12 @@
16531673
"ReasoningEngineContextSpecMemoryBankConfigTtlConfig",
16541674
"ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict",
16551675
"ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict",
1676+
"StructuredMemoryConfigSchemaConfig",
1677+
"StructuredMemoryConfigSchemaConfigDict",
1678+
"StructuredMemoryConfigSchemaConfigOrDict",
1679+
"StructuredMemoryConfig",
1680+
"StructuredMemoryConfigDict",
1681+
"StructuredMemoryConfigOrDict",
16561682
"ReasoningEngineContextSpecMemoryBankConfig",
16571683
"ReasoningEngineContextSpecMemoryBankConfigDict",
16581684
"ReasoningEngineContextSpecMemoryBankConfigOrDict",
@@ -1743,6 +1769,9 @@
17431769
"AgentEngineMemoryConfig",
17441770
"AgentEngineMemoryConfigDict",
17451771
"AgentEngineMemoryConfigOrDict",
1772+
"MemoryStructuredContent",
1773+
"MemoryStructuredContentDict",
1774+
"MemoryStructuredContentOrDict",
17461775
"Memory",
17471776
"MemoryDict",
17481777
"MemoryOrDict",
@@ -1812,6 +1841,15 @@
18121841
"RetrieveMemoriesResponse",
18131842
"RetrieveMemoriesResponseDict",
18141843
"RetrieveMemoriesResponseOrDict",
1844+
"RetrieveMemoryProfilesConfig",
1845+
"RetrieveMemoryProfilesConfigDict",
1846+
"RetrieveMemoryProfilesConfigOrDict",
1847+
"MemoryProfile",
1848+
"MemoryProfileDict",
1849+
"MemoryProfileOrDict",
1850+
"RetrieveProfilesResponse",
1851+
"RetrieveProfilesResponseDict",
1852+
"RetrieveProfilesResponseOrDict",
18151853
"RollbackAgentEngineMemoryConfig",
18161854
"RollbackAgentEngineMemoryConfigDict",
18171855
"RollbackAgentEngineMemoryConfigOrDict",
@@ -2197,6 +2235,7 @@
21972235
"Type",
21982236
"JobState",
21992237
"ManagedTopicEnum",
2238+
"MemoryType",
22002239
"IdentityType",
22012240
"AgentServerMode",
22022241
"Operator",
@@ -2271,6 +2310,7 @@
22712310
"_GetAgentEngineMemoryOperationParameters",
22722311
"_GetAgentEngineGenerateMemoriesOperationParameters",
22732312
"_RetrieveAgentEngineMemoriesRequestParameters",
2313+
"_RetrieveMemoryProfilesRequestParameters",
22742314
"_RollbackAgentEngineMemoryRequestParameters",
22752315
"_UpdateAgentEngineMemoryRequestParameters",
22762316
"_PurgeAgentEngineMemoriesRequestParameters",

0 commit comments

Comments
 (0)