diff --git a/tests/unit/vertexai/genai/replays/test_structured_memories.py b/tests/unit/vertexai/genai/replays/test_structured_memories.py new file mode 100644 index 0000000000..ce5df9fa45 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_structured_memories.py @@ -0,0 +1,100 @@ +# Copyright 2025 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + + +def test_generate_and_retrieve_profile(client): + # TODO: Use prod once available. + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com" + ) + customization_config = {"disable_natural_language_memories": True} + memory_bank_customization_config = types.MemoryBankCustomizationConfig( + **customization_config + ) + structured_memory_config = { + "scope_keys": ["user_id"], + "schema_configs": [ + { + "id": "user-profile", + "memory_schema": { + "properties": { + "name": {"description": "User's name", "type": "string"} + }, + "type": "object", + }, + } + ], + } + structured_memory_config_obj = types.StructuredMemoryConfig( + **structured_memory_config + ) + agent_engine = client.agent_engines.create( + config={ + "context_spec": { + "memory_bank_config": { + "customization_configs": [memory_bank_customization_config], + "structured_memory_configs": [structured_memory_config_obj], + }, + }, + "http_options": {"api_version": "v1beta1"}, + }, + ) + try: + agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name) + memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config + assert memory_bank_config.customization_configs == [ + memory_bank_customization_config + ] + assert memory_bank_config.structured_memory_configs == [ + structured_memory_config_obj + ] + + scope = {"user_id": "123"} + client.agent_engines.memories.generate( + name=agent_engine.api_resource.name, + scope=scope, + direct_contents_source={ + "events": [{"content": {"parts": [{"text": "My name is Kim."}]}}] + }, + ) + memories = list( + client.agent_engines.memories.retrieve( + name=agent_engine.api_resource.name, + scope=scope, + config={"memory_types": ["STRUCTURED_PROFILE"]}, + ) + ) + assert len(memories) >= 1 + assert memories[0].memory.structured_content is not None + + response = client.agent_engines.memories.retrieve_profiles( + name=agent_engine.api_resource.name, scope=scope + ) + assert len(response.profiles) == 1 + + finally: + # Clean up resources. + client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.retrieve_profiles", +) diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index b72880ce22..ea0a00b9e9 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -49,6 +49,33 @@ logger.setLevel(logging.INFO) +def _AgentEngineOperation_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["metadata"]) is not None: + setv(to_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["done"]) is not None: + setv(to_object, ["done"], getv(from_object, ["done"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + if getv(from_object, ["response"]) is not None: + setv( + to_object, + ["response"], + _ReasoningEngine_from_vertex(getv(from_object, ["response"]), to_object), + ) + + return to_object + + def _CheckQueryJobAgentEngineConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -118,7 +145,13 @@ def _CreateAgentEngineConfig_to_vertex( setv(parent_object, ["spec"], getv(from_object, ["spec"])) if getv(from_object, ["context_spec"]) is not None: - setv(parent_object, ["contextSpec"], getv(from_object, ["context_spec"])) + setv( + parent_object, + ["contextSpec"], + _ReasoningEngineContextSpec_to_vertex( + getv(from_object, ["context_spec"]), to_object + ), + ) if getv(from_object, ["psc_interface_config"]) is not None: setv( @@ -251,6 +284,30 @@ def _ListAgentEngineRequestParameters_to_vertex( return to_object +def _ListReasoningEnginesResponse_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["sdkHttpResponse"]) is not None: + setv(to_object, ["sdk_http_response"], getv(from_object, ["sdkHttpResponse"])) + + if getv(from_object, ["nextPageToken"]) is not None: + setv(to_object, ["next_page_token"], getv(from_object, ["nextPageToken"])) + + if getv(from_object, ["reasoningEngines"]) is not None: + setv( + to_object, + ["reasoning_engines"], + [ + _ReasoningEngine_from_vertex(item, to_object) + for item in getv(from_object, ["reasoningEngines"]) + ], + ) + + return to_object + + def _QueryAgentEngineConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -283,6 +340,177 @@ def _QueryAgentEngineRequestParameters_to_vertex( return to_object +def _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["customizationConfigs"]) is not None: + setv( + to_object, + ["customization_configs"], + [item for item in getv(from_object, ["customizationConfigs"])], + ) + + if getv(from_object, ["disableMemoryRevisions"]) is not None: + setv( + to_object, + ["disable_memory_revisions"], + getv(from_object, ["disableMemoryRevisions"]), + ) + + if getv(from_object, ["generationConfig"]) is not None: + setv(to_object, ["generation_config"], getv(from_object, ["generationConfig"])) + + if getv(from_object, ["similaritySearchConfig"]) is not None: + setv( + to_object, + ["similarity_search_config"], + getv(from_object, ["similaritySearchConfig"]), + ) + + if getv(from_object, ["ttlConfig"]) is not None: + setv(to_object, ["ttl_config"], getv(from_object, ["ttlConfig"])) + + if getv(from_object, ["structuredMemoryConfigs"]) is not None: + setv( + to_object, + ["structured_memory_configs"], + [ + _StructuredMemoryConfig_from_vertex(item, to_object) + for item in getv(from_object, ["structuredMemoryConfigs"]) + ], + ) + + return to_object + + +def _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["customization_configs"]) is not None: + setv( + to_object, + ["customizationConfigs"], + [item for item in getv(from_object, ["customization_configs"])], + ) + + if getv(from_object, ["disable_memory_revisions"]) is not None: + setv( + to_object, + ["disableMemoryRevisions"], + getv(from_object, ["disable_memory_revisions"]), + ) + + if getv(from_object, ["generation_config"]) is not None: + setv(to_object, ["generationConfig"], getv(from_object, ["generation_config"])) + + if getv(from_object, ["similarity_search_config"]) is not None: + setv( + to_object, + ["similaritySearchConfig"], + getv(from_object, ["similarity_search_config"]), + ) + + if getv(from_object, ["ttl_config"]) is not None: + setv(to_object, ["ttlConfig"], getv(from_object, ["ttl_config"])) + + if getv(from_object, ["structured_memory_configs"]) is not None: + setv( + to_object, + ["structuredMemoryConfigs"], + [ + _StructuredMemoryConfig_to_vertex(item, to_object) + for item in getv(from_object, ["structured_memory_configs"]) + ], + ) + + return to_object + + +def _ReasoningEngineContextSpec_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memoryBankConfig"]) is not None: + setv( + to_object, + ["memory_bank_config"], + _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( + getv(from_object, ["memoryBankConfig"]), to_object + ), + ) + + return to_object + + +def _ReasoningEngineContextSpec_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memory_bank_config"]) is not None: + setv( + to_object, + ["memoryBankConfig"], + _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + getv(from_object, ["memory_bank_config"]), to_object + ), + ) + + return to_object + + +def _ReasoningEngine_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["encryptionSpec"]) is not None: + setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"])) + + if getv(from_object, ["contextSpec"]) is not None: + setv( + to_object, + ["context_spec"], + _ReasoningEngineContextSpec_from_vertex( + getv(from_object, ["contextSpec"]), to_object + ), + ) + + if getv(from_object, ["createTime"]) is not None: + setv(to_object, ["create_time"], getv(from_object, ["createTime"])) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["displayName"]) is not None: + setv(to_object, ["display_name"], getv(from_object, ["displayName"])) + + if getv(from_object, ["etag"]) is not None: + setv(to_object, ["etag"], getv(from_object, ["etag"])) + + if getv(from_object, ["labels"]) is not None: + setv(to_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["spec"]) is not None: + setv(to_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["updateTime"]) is not None: + setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) + + if getv(from_object, ["trafficConfig"]) is not None: + setv(to_object, ["traffic_config"], getv(from_object, ["trafficConfig"])) + + return to_object + + def _RunQueryJobAgentEngineConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -318,6 +546,82 @@ def _RunQueryJobAgentEngineRequestParameters_to_vertex( return to_object +def _StructuredMemoryConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schemaConfigs"]) is not None: + setv( + to_object, + ["schema_configs"], + [ + _StructuredMemorySchemaConfig_from_vertex(item, to_object) + for item in getv(from_object, ["schemaConfigs"]) + ], + ) + + if getv(from_object, ["scopeKeys"]) is not None: + setv(to_object, ["scope_keys"], getv(from_object, ["scopeKeys"])) + + return to_object + + +def _StructuredMemoryConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema_configs"]) is not None: + setv( + to_object, + ["schemaConfigs"], + [ + _StructuredMemorySchemaConfig_to_vertex(item, to_object) + for item in getv(from_object, ["schema_configs"]) + ], + ) + + if getv(from_object, ["scope_keys"]) is not None: + setv(to_object, ["scopeKeys"], getv(from_object, ["scope_keys"])) + + return to_object + + +def _StructuredMemorySchemaConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema"]) is not None: + setv(to_object, ["memory_schema"], getv(from_object, ["schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memoryType"]) is not None: + setv(to_object, ["memory_type"], getv(from_object, ["memoryType"])) + + return to_object + + +def _StructuredMemorySchemaConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memory_schema"]) is not None: + setv(to_object, ["schema"], getv(from_object, ["memory_schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memory_type"]) is not None: + setv(to_object, ["memoryType"], getv(from_object, ["memory_type"])) + + return to_object + + def _UpdateAgentEngineConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -334,7 +638,13 @@ def _UpdateAgentEngineConfig_to_vertex( setv(parent_object, ["spec"], getv(from_object, ["spec"])) if getv(from_object, ["context_spec"]) is not None: - setv(parent_object, ["contextSpec"], getv(from_object, ["context_spec"])) + setv( + parent_object, + ["contextSpec"], + _ReasoningEngineContextSpec_to_vertex( + getv(from_object, ["context_spec"]), to_object + ), + ) if getv(from_object, ["psc_interface_config"]) is not None: setv( @@ -533,6 +843,9 @@ def _run_query_job( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( @@ -601,6 +914,9 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( @@ -755,6 +1071,9 @@ def _get( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex(response_dict) + return_value = types.ReasoningEngine._from_response( response=response_dict, kwargs=( @@ -821,6 +1140,9 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + return_value = types.ListReasoningEnginesResponse._from_response( response=response_dict, kwargs=( @@ -887,6 +1209,9 @@ def _get_agent_operation( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( @@ -1023,6 +1348,9 @@ def _update( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( @@ -2770,6 +3098,9 @@ async def _run_query_job( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( @@ -2840,6 +3171,9 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( @@ -2998,6 +3332,9 @@ async def _get( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex(response_dict) + return_value = types.ReasoningEngine._from_response( response=response_dict, kwargs=( @@ -3066,6 +3403,9 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + return_value = types.ListReasoningEnginesResponse._from_response( response=response_dict, kwargs=( @@ -3134,6 +3474,9 @@ async def _get_agent_operation( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( @@ -3274,6 +3617,9 @@ async def _update( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _AgentEngineOperation_from_vertex(response_dict) + return_value = types.AgentEngineOperation._from_response( response=response_dict, kwargs=( diff --git a/vertexai/_genai/memories.py b/vertexai/_genai/memories.py index 4a48350272..1291b04827 100644 --- a/vertexai/_genai/memories.py +++ b/vertexai/_genai/memories.py @@ -336,6 +336,9 @@ def _RetrieveAgentEngineMemoriesConfig_to_vertex( [item for item in getv(from_object, ["filter_groups"])], ) + if getv(from_object, ["memory_types"]) is not None: + setv(parent_object, ["memoryTypes"], getv(from_object, ["memory_types"])) + return to_object @@ -372,6 +375,20 @@ def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( return to_object +def _RetrieveMemoryProfilesRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["scope"]) is not None: + setv(to_object, ["scope"], getv(from_object, ["scope"])) + + return to_object + + def _RollbackAgentEngineMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -1078,6 +1095,89 @@ def _retrieve( self._api_client._verify_response(return_value) return return_value + def retrieve_profiles( + self, + *, + name: str, + scope: dict[str, str], + config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, + ) -> types.RetrieveProfilesResponse: + """ + Retrieves memory profiles for an Agent Engine. + + Args: + name (str): Required. A fully-qualified resource name or ID such as + "projects/123/locations/us-central1/reasoningEngines/456". + scope (dict[str, str]): Required. The scope of the memories to retrieve. + A memory must have exactly the same scope as the scope provided here + to be retrieved (i.e. same keys and values). Order does not matter, + but it is case-sensitive. + + """ + + parameter_model = types._RetrieveMemoryProfilesRequestParameters( + name=name, + scope=scope, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _RetrieveMemoryProfilesRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/memories:retrieveProfiles".format_map(request_url_dict) + else: + path = "{name}/memories:retrieveProfiles" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.RetrieveProfilesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + def _rollback( self, *, @@ -2239,6 +2339,91 @@ async def _retrieve( self._api_client._verify_response(return_value) return return_value + async def retrieve_profiles( + self, + *, + name: str, + scope: dict[str, str], + config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, + ) -> types.RetrieveProfilesResponse: + """ + Retrieves memory profiles for an Agent Engine. + + Args: + name (str): Required. A fully-qualified resource name or ID such as + "projects/123/locations/us-central1/reasoningEngines/456". + scope (dict[str, str]): Required. The scope of the memories to retrieve. + A memory must have exactly the same scope as the scope provided here + to be retrieved (i.e. same keys and values). Order does not matter, + but it is case-sensitive. + + """ + + parameter_model = types._RetrieveMemoryProfilesRequestParameters( + name=name, + scope=scope, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _RetrieveMemoryProfilesRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/memories:retrieveProfiles".format_map(request_url_dict) + else: + path = "{name}/memories:retrieveProfiles" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.RetrieveProfilesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + async def _rollback( self, *, diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index f293f22bc7..e39f46ae40 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -95,6 +95,7 @@ from .common import _QueryAgentEngineRequestParameters from .common import _RestoreVersionRequestParameters from .common import _RetrieveAgentEngineMemoriesRequestParameters +from .common import _RetrieveMemoryProfilesRequestParameters from .common import _RollbackAgentEngineMemoryRequestParameters from .common import _RunQueryJobAgentEngineConfig from .common import _RunQueryJobAgentEngineConfigDict @@ -699,12 +700,19 @@ from .common import MemoryMetadataValueDict from .common import MemoryMetadataValueOrDict from .common import MemoryOrDict +from .common import MemoryProfile +from .common import MemoryProfileDict +from .common import MemoryProfileOrDict from .common import MemoryRevision from .common import MemoryRevisionDict from .common import MemoryRevisionOrDict +from .common import MemoryStructuredContent +from .common import MemoryStructuredContentDict +from .common import MemoryStructuredContentOrDict from .common import MemoryTopicId from .common import MemoryTopicIdDict from .common import MemoryTopicIdOrDict +from .common import MemoryType from .common import Message from .common import MessageDict from .common import Metadata @@ -929,6 +937,12 @@ from .common import RetrieveMemoriesResponseRetrievedMemory from .common import RetrieveMemoriesResponseRetrievedMemoryDict from .common import RetrieveMemoriesResponseRetrievedMemoryOrDict +from .common import RetrieveMemoryProfilesConfig +from .common import RetrieveMemoryProfilesConfigDict +from .common import RetrieveMemoryProfilesConfigOrDict +from .common import RetrieveProfilesResponse +from .common import RetrieveProfilesResponseDict +from .common import RetrieveProfilesResponseOrDict from .common import RollbackAgentEngineMemoryConfig from .common import RollbackAgentEngineMemoryConfigDict from .common import RollbackAgentEngineMemoryConfigOrDict @@ -1084,6 +1098,12 @@ from .common import SessionOrDict from .common import State from .common import Strategy +from .common import StructuredMemoryConfig +from .common import StructuredMemoryConfigDict +from .common import StructuredMemoryConfigOrDict +from .common import StructuredMemorySchemaConfig +from .common import StructuredMemorySchemaConfigDict +from .common import StructuredMemorySchemaConfigOrDict from .common import SummaryMetric from .common import SummaryMetricDict from .common import SummaryMetricOrDict @@ -1699,6 +1719,12 @@ "ReasoningEngineContextSpecMemoryBankConfigTtlConfig", "ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict", "ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict", + "StructuredMemorySchemaConfig", + "StructuredMemorySchemaConfigDict", + "StructuredMemorySchemaConfigOrDict", + "StructuredMemoryConfig", + "StructuredMemoryConfigDict", + "StructuredMemoryConfigOrDict", "ReasoningEngineContextSpecMemoryBankConfig", "ReasoningEngineContextSpecMemoryBankConfigDict", "ReasoningEngineContextSpecMemoryBankConfigOrDict", @@ -1816,6 +1842,9 @@ "AgentEngineMemoryConfig", "AgentEngineMemoryConfigDict", "AgentEngineMemoryConfigOrDict", + "MemoryStructuredContent", + "MemoryStructuredContentDict", + "MemoryStructuredContentOrDict", "Memory", "MemoryDict", "MemoryOrDict", @@ -1885,6 +1914,15 @@ "RetrieveMemoriesResponse", "RetrieveMemoriesResponseDict", "RetrieveMemoriesResponseOrDict", + "RetrieveMemoryProfilesConfig", + "RetrieveMemoryProfilesConfigDict", + "RetrieveMemoryProfilesConfigOrDict", + "MemoryProfile", + "MemoryProfileDict", + "MemoryProfileOrDict", + "RetrieveProfilesResponse", + "RetrieveProfilesResponseDict", + "RetrieveProfilesResponseOrDict", "RollbackAgentEngineMemoryConfig", "RollbackAgentEngineMemoryConfigDict", "RollbackAgentEngineMemoryConfigOrDict", @@ -2272,6 +2310,7 @@ "ManagedTopicEnum", "IdentityType", "AgentServerMode", + "MemoryType", "Operator", "Language", "MachineConfig", @@ -2345,6 +2384,7 @@ "_GetAgentEngineMemoryOperationParameters", "_GetAgentEngineGenerateMemoriesOperationParameters", "_RetrieveAgentEngineMemoriesRequestParameters", + "_RetrieveMemoryProfilesRequestParameters", "_RollbackAgentEngineMemoryRequestParameters", "_UpdateAgentEngineMemoryRequestParameters", "_PurgeAgentEngineMemoriesRequestParameters", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 97ecc9c3eb..f94476e49c 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -280,6 +280,17 @@ class AgentServerMode(_common.CaseInSensitiveEnum): """Experimental agent server mode. This mode contains experimental features.""" +class MemoryType(_common.CaseInSensitiveEnum): + """The type of the memory.""" + + MEMORY_TYPE_UNSPECIFIED = "MEMORY_TYPE_UNSPECIFIED" + """Represents an unspecified memory type. This value should not be used.""" + NATURAL_LANGUAGE_COLLECTION = "NATURAL_LANGUAGE_COLLECTION" + """Indicates belonging to a collection of natural language memories.""" + STRUCTURED_PROFILE = "STRUCTURED_PROFILE" + """Indicates belonging to a structured profile.""" + + class Operator(_common.CaseInSensitiveEnum): """Represents the operator to apply to the filter. If not set, then EQUAL will be used.""" @@ -6683,6 +6694,10 @@ class MemoryBankCustomizationConfig(_common.BaseModel): description="""Optional. Represents configuration for customizing how memories are consolidated together.""", ) ) + disable_natural_language_memories: Optional[bool] = Field( + default=None, + description="""Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""", + ) class MemoryBankCustomizationConfigDict(TypedDict, total=False): @@ -6705,6 +6720,9 @@ class MemoryBankCustomizationConfigDict(TypedDict, total=False): consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfigDict] """Optional. Represents configuration for customizing how memories are consolidated together.""" + disable_natural_language_memories: Optional[bool] + """Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""" + MemoryBankCustomizationConfigOrDict = Union[ MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict @@ -6841,6 +6859,67 @@ class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=F ] +class StructuredMemorySchemaConfig(_common.BaseModel): + """Represents the OpenAPI schema of the structured memories.""" + + memory_schema: Optional[genai_types.Schema] = Field( + default=None, + description="""Required. Represents the OpenAPI schema of the structured memories.""", + ) + id: Optional[str] = Field( + default=None, + description="""Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""", + ) + memory_type: Optional[MemoryType] = Field( + default=None, + description="""Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""", + ) + + +class StructuredMemorySchemaConfigDict(TypedDict, total=False): + """Represents the OpenAPI schema of the structured memories.""" + + memory_schema: Optional[genai_types.SchemaDict] + """Required. Represents the OpenAPI schema of the structured memories.""" + + id: Optional[str] + """Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""" + + memory_type: Optional[MemoryType] + """Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""" + + +StructuredMemorySchemaConfigOrDict = Union[ + StructuredMemorySchemaConfig, StructuredMemorySchemaConfigDict +] + + +class StructuredMemoryConfig(_common.BaseModel): + """Configuration for organizing structured memories within a scope.""" + + schema_configs: Optional[list[StructuredMemorySchemaConfig]] = Field( + default=None, + description="""Optional. Represents configuration of the structured memories' schemas.""", + ) + scope_keys: Optional[list[str]] = Field( + default=None, + description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", + ) + + +class StructuredMemoryConfigDict(TypedDict, total=False): + """Configuration for organizing structured memories within a scope.""" + + schema_configs: Optional[list[StructuredMemorySchemaConfigDict]] + """Optional. Represents configuration of the structured memories' schemas.""" + + scope_keys: Optional[list[str]] + """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" + + +StructuredMemoryConfigOrDict = Union[StructuredMemoryConfig, StructuredMemoryConfigDict] + + class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): """Specification for a Memory Bank.""" @@ -6868,6 +6947,10 @@ class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): default=None, description="""Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""", ) + structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + default=None, + description="""Optional. Configuration for organizing structured memories for a particular scope.""", + ) class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): @@ -6892,6 +6975,9 @@ class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict] """Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""" + structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] + """Optional. Configuration for organizing structured memories for a particular scope.""" + ReasoningEngineContextSpecMemoryBankConfigOrDict = Union[ ReasoningEngineContextSpecMemoryBankConfig, @@ -9019,6 +9105,34 @@ class _CreateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): ] +class MemoryStructuredContent(_common.BaseModel): + """Represents the structured value of the memory.""" + + data: Optional[dict[str, Any]] = Field( + default=None, + description="""Required. Represents the structured value of the memory.""", + ) + schema_id: Optional[str] = Field( + default=None, + description="""Required. Represents the schema ID for which this structured memory belongs to.""", + ) + + +class MemoryStructuredContentDict(TypedDict, total=False): + """Represents the structured value of the memory.""" + + data: Optional[dict[str, Any]] + """Required. Represents the structured value of the memory.""" + + schema_id: Optional[str] + """Required. Represents the schema ID for which this structured memory belongs to.""" + + +MemoryStructuredContentOrDict = Union[ + MemoryStructuredContent, MemoryStructuredContentDict +] + + class Memory(_common.BaseModel): """A memory.""" @@ -9081,6 +9195,14 @@ class Memory(_common.BaseModel): default=None, description="""Output only. Represents the timestamp when this Memory was most recently updated.""", ) + memory_type: Optional[MemoryType] = Field( + default=None, + description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", + ) + structured_content: Optional[MemoryStructuredContent] = Field( + default=None, + description="""Optional. Represents the structured content of the memory.""", + ) class MemoryDict(TypedDict, total=False): @@ -9131,6 +9253,12 @@ class MemoryDict(TypedDict, total=False): update_time: Optional[datetime.datetime] """Output only. Represents the timestamp when this Memory was most recently updated.""" + memory_type: Optional[MemoryType] + """Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""" + + structured_content: Optional[MemoryStructuredContentDict] + """Optional. Represents the structured content of the memory.""" + MemoryOrDict = Union[Memory, MemoryDict] @@ -10044,6 +10172,13 @@ class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): metadata.author = "agent 321"))`. """, ) + memory_types: Optional[list[MemoryType]] = Field( + default=None, + description="""Specifies the types of memories to retrieve. If this field is empty + or not provided, the request will default to retrieving only memories of + type `NATURAL_LANGUAGE_COLLECTION`. If populated, the request will + retrieve memories matching any of the specified `MemoryType` values.""", + ) class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): @@ -10078,6 +10213,12 @@ class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): metadata.author = "agent 321"))`. """ + memory_types: Optional[list[MemoryType]] + """Specifies the types of memories to retrieve. If this field is empty + or not provided, the request will default to retrieving only memories of + type `NATURAL_LANGUAGE_COLLECTION`. If populated, the request will + retrieve memories matching any of the specified `MemoryType` values.""" + RetrieveAgentEngineMemoriesConfigOrDict = Union[ RetrieveAgentEngineMemoriesConfig, RetrieveAgentEngineMemoriesConfigDict @@ -10201,6 +10342,121 @@ class RetrieveMemoriesResponseDict(TypedDict, total=False): ] +class RetrieveMemoryProfilesConfig(_common.BaseModel): + """Config for retrieving memory profiles.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): + """Config for retrieving memory profiles.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +RetrieveMemoryProfilesConfigOrDict = Union[ + RetrieveMemoryProfilesConfig, RetrieveMemoryProfilesConfigDict +] + + +class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): + """Parameters for retrieving agent engine memory profiles.""" + + name: Optional[str] = Field( + default=None, + description="""Name of the agent engine to retrieve memory profiles from.""", + ) + scope: Optional[dict[str, str]] = Field( + default=None, + description="""The scope of the memories to retrieve. + + A memory must have exactly the same scope as the scope provided here to be + retrieved (i.e. same keys and values). Order does not matter, but it is + case-sensitive.""", + ) + config: Optional[RetrieveMemoryProfilesConfig] = Field( + default=None, description="""""" + ) + + +class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving agent engine memory profiles.""" + + name: Optional[str] + """Name of the agent engine to retrieve memory profiles from.""" + + scope: Optional[dict[str, str]] + """The scope of the memories to retrieve. + + A memory must have exactly the same scope as the scope provided here to be + retrieved (i.e. same keys and values). Order does not matter, but it is + case-sensitive.""" + + config: Optional[RetrieveMemoryProfilesConfigDict] + """""" + + +_RetrieveMemoryProfilesRequestParametersOrDict = Union[ + _RetrieveMemoryProfilesRequestParameters, + _RetrieveMemoryProfilesRequestParametersDict, +] + + +class MemoryProfile(_common.BaseModel): + """A memory profile.""" + + schema_id: Optional[str] = Field( + default=None, + description="""Represents the ID of the schema. This ID corresponds to the `schema_id` defined inside the SchemaConfig, under StructuredMemoryCustomizationConfig.""", + ) + profile: Optional[dict[str, Any]] = Field( + default=None, description="""Represents the profile data.""" + ) + + +class MemoryProfileDict(TypedDict, total=False): + """A memory profile.""" + + schema_id: Optional[str] + """Represents the ID of the schema. This ID corresponds to the `schema_id` defined inside the SchemaConfig, under StructuredMemoryCustomizationConfig.""" + + profile: Optional[dict[str, Any]] + """Represents the profile data.""" + + +MemoryProfileOrDict = Union[MemoryProfile, MemoryProfileDict] + + +class RetrieveProfilesResponse(_common.BaseModel): + """The response for retrieving memory profiles.""" + + memory_profiles: Optional[dict[str, MemoryProfile]] = Field( + default=None, description="""The memory profiles.""" + ) + profiles: Optional[dict[str, MemoryProfile]] = Field( + default=None, + description="""The retrieved structured profiles, which match the schemas under the requested scope. The key is the ID of the schema that the profile is linked with, which corresponds to the `schema_id` defined inside the `SchemaConfig`, under `StructuredMemoryCustomizationConfig`.""", + ) + + +class RetrieveProfilesResponseDict(TypedDict, total=False): + """The response for retrieving memory profiles.""" + + memory_profiles: Optional[dict[str, MemoryProfileDict]] + """The memory profiles.""" + + profiles: Optional[dict[str, MemoryProfileDict]] + """The retrieved structured profiles, which match the schemas under the requested scope. The key is the ID of the schema that the profile is linked with, which corresponds to the `schema_id` defined inside the `SchemaConfig`, under `StructuredMemoryCustomizationConfig`.""" + + +RetrieveProfilesResponseOrDict = Union[ + RetrieveProfilesResponse, RetrieveProfilesResponseDict +] + + class RollbackAgentEngineMemoryConfig(_common.BaseModel): """Config for rolling back a memory.""" @@ -10686,6 +10942,14 @@ class IntermediateExtractedMemory(_common.BaseModel): default=None, description="""Output only. Represents the fact of the extracted memory.""", ) + context: Optional[str] = Field( + default=None, + description="""Output only. Represents the explanation of why the information was extracted from the source content.""", + ) + structured_data: Optional[dict[str, Any]] = Field( + default=None, + description="""Output only. Represents the structured value of the extracted memory.""", + ) class IntermediateExtractedMemoryDict(TypedDict, total=False): @@ -10694,6 +10958,12 @@ class IntermediateExtractedMemoryDict(TypedDict, total=False): fact: Optional[str] """Output only. Represents the fact of the extracted memory.""" + context: Optional[str] + """Output only. Represents the explanation of why the information was extracted from the source content.""" + + structured_data: Optional[dict[str, Any]] + """Output only. Represents the structured value of the extracted memory.""" + IntermediateExtractedMemoryOrDict = Union[ IntermediateExtractedMemory, IntermediateExtractedMemoryDict @@ -10727,6 +10997,10 @@ class MemoryRevision(_common.BaseModel): default=None, description="""Identifier. Represents the resource name of the Memory Revision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{memory_revision}`""", ) + structured_data: Optional[dict[str, Any]] = Field( + default=None, + description="""Output only. Represents the structured value of the memory at the time of revision creation.""", + ) class MemoryRevisionDict(TypedDict, total=False): @@ -10750,6 +11024,9 @@ class MemoryRevisionDict(TypedDict, total=False): name: Optional[str] """Identifier. Represents the resource name of the Memory Revision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{memory_revision}`""" + structured_data: Optional[dict[str, Any]] + """Output only. Represents the structured value of the memory at the time of revision creation.""" + MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict]