diff --git a/agentplatform/_genai/memory_banks.py b/agentplatform/_genai/memory_banks.py index 1ad1b6e231..bddc4ea5ca 100644 --- a/agentplatform/_genai/memory_banks.py +++ b/agentplatform/_genai/memory_banks.py @@ -590,6 +590,44 @@ def _StructuredMemorySchemaConfig_to_vertex( return to_object +def _UpdateMemoryBankConfig_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, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + return to_object + + +def _UpdateMemoryBankRequestParameters_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, ["memory_bank_config"]) is not None: + setv( + to_object, + ["context_spec", "memory_bank_config"], + _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + getv(from_object, ["memory_bank_config"]), to_object + ), + ) + + if getv(from_object, ["config"]) is not None: + _UpdateMemoryBankConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + class MemoryBanks(_api_module.BaseModule): def _create( @@ -669,6 +707,85 @@ def _create( self._api_client._verify_response(return_value) return return_value + def _update( + self, + *, + name: str, + memory_bank_config: Optional[ + types.ReasoningEngineContextSpecMemoryBankConfigOrDict + ] = None, + config: Optional[types.UpdateMemoryBankConfigOrDict] = None, + ) -> types.MemoryBankOperation: + """ + Updates a Memory Bank. + """ + + parameter_model = types._UpdateMemoryBankRequestParameters( + name=name, + memory_bank_config=memory_bank_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateMemoryBankRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + 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("patch", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _MemoryBankOperation_from_vertex(response_dict) + + return_value = types.MemoryBankOperation._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 _delete( self, *, @@ -1082,6 +1199,41 @@ def create( memory_bank = types.MemoryBank(name=operation.response.name) return memory_bank + def update( + self, + *, + name: str, + managed_semantic_memory_config: Optional[ + types.ManagedSemanticMemoryConfigOrDict + ] = None, + config: Optional[types.UpdateMemoryBankConfigOrDict] = None, + ) -> types.MemoryBank: + """Updates a Memory Bank resource.""" + memory_bank_config = ( + _memory_bank_utils._managed_semantic_memory_config_to_memory_bank_config( + managed_semantic_memory_config + ) + ) + + operation = self._update( + name=name, + memory_bank_config=memory_bank_config, + config=config, + ) + + operation = _memory_bank_utils._await_operation( + operation_name=operation.name, + get_operation_fn=self._get_memory_bank_operation, + ) + if operation.error: + raise RuntimeError(operation.error) + if not operation.response: + raise RuntimeError( + f"Memory Bank update completed without a response: {operation}" + ) + memory_bank = types.MemoryBank(name=operation.response.name) + return memory_bank + def delete( self, *, @@ -1332,6 +1484,87 @@ async def _create( self._api_client._verify_response(return_value) return return_value + async def _update( + self, + *, + name: str, + memory_bank_config: Optional[ + types.ReasoningEngineContextSpecMemoryBankConfigOrDict + ] = None, + config: Optional[types.UpdateMemoryBankConfigOrDict] = None, + ) -> types.MemoryBankOperation: + """ + Updates a Memory Bank. + """ + + parameter_model = types._UpdateMemoryBankRequestParameters( + name=name, + memory_bank_config=memory_bank_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateMemoryBankRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + 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( + "patch", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _MemoryBankOperation_from_vertex(response_dict) + + return_value = types.MemoryBankOperation._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 _delete( self, *, @@ -1755,6 +1988,41 @@ async def create( memory_bank = types.MemoryBank(name=operation.response.name) return memory_bank + async def update( + self, + *, + name: str, + managed_semantic_memory_config: Optional[ + types.ManagedSemanticMemoryConfigOrDict + ] = None, + config: Optional[types.UpdateMemoryBankConfigOrDict] = None, + ) -> types.MemoryBank: + """Updates a Memory Bank.""" + memory_bank_config = ( + _memory_bank_utils._managed_semantic_memory_config_to_memory_bank_config( + managed_semantic_memory_config + ) + ) + + operation = await self._update( + name=name, + memory_bank_config=memory_bank_config, + config=config, + ) + + operation = await _memory_bank_utils._await_async_operation( + operation_name=operation.name, + get_operation_fn=self._get_memory_bank_operation, + ) + if operation.error: + raise RuntimeError(operation.error) + if not operation.response: + raise RuntimeError( + f"Memory Bank update completed without a response: {operation}" + ) + memory_bank = types.MemoryBank(name=operation.response.name) + return memory_bank + async def delete( self, *, diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 1c122ccd3a..4c0d33c86e 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -183,6 +183,7 @@ from .common import _UndeployModelRequestParameters from .common import _UpdateDatasetParameters from .common import _UpdateEvaluationExperimentParameters +from .common import _UpdateMemoryBankRequestParameters from .common import _UpdateMemoryRequestParameters from .common import _UpdateMultimodalDatasetParameters from .common import _UpdateRagConfigRequestParameters @@ -2265,6 +2266,9 @@ from .common import UpdateEvaluationExperimentConfig from .common import UpdateEvaluationExperimentConfigDict from .common import UpdateEvaluationExperimentConfigOrDict +from .common import UpdateMemoryBankConfig +from .common import UpdateMemoryBankConfigDict +from .common import UpdateMemoryBankConfigOrDict from .common import UpdateMemoryConfig from .common import UpdateMemoryConfigDict from .common import UpdateMemoryConfigOrDict @@ -3024,6 +3028,9 @@ "MemoryBankOperation", "MemoryBankOperationDict", "MemoryBankOperationOrDict", + "UpdateMemoryBankConfig", + "UpdateMemoryBankConfigDict", + "UpdateMemoryBankConfigOrDict", "DeleteMemoryBankConfig", "DeleteMemoryBankConfigDict", "DeleteMemoryBankConfigOrDict", @@ -4493,6 +4500,7 @@ "_GetDeleteRuntimeRevisionOperationParameters", "_QueryRuntimeRevisionRequestParameters", "_CreateMemoryBankRequestParameters", + "_UpdateMemoryBankRequestParameters", "_DeleteMemoryBankRequestParameters", "_GetMemoryBankRequestParameters", "_IngestEventsRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 9cac27746a..8895b9feef 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -10965,6 +10965,75 @@ class MemoryBankOperationDict(TypedDict, total=False): MemoryBankOperationOrDict = Union[MemoryBankOperation, MemoryBankOperationDict] +class UpdateMemoryBankConfig(_common.BaseModel): + """Config for updating a memory bank.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Memory Bank. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, + ) + description: Optional[str] = Field( + default=None, description="""The description of the Memory Bank.""" + ) + + +class UpdateMemoryBankConfigDict(TypedDict, total=False): + """Config for updating a memory bank.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + display_name: Optional[str] + """The user-defined name of the Memory Bank. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ + + description: Optional[str] + """The description of the Memory Bank.""" + + +UpdateMemoryBankConfigOrDict = Union[UpdateMemoryBankConfig, UpdateMemoryBankConfigDict] + + +class _UpdateMemoryBankRequestParameters(_common.BaseModel): + """Parameters for updating memory banks.""" + + name: Optional[str] = Field( + default=None, description="""Name of the memory bank.""" + ) + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( + default=None, description="""""" + ) + config: Optional[UpdateMemoryBankConfig] = Field(default=None, description="""""") + + +class _UpdateMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for updating memory banks.""" + + name: Optional[str] + """Name of the memory bank.""" + + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + """""" + + config: Optional[UpdateMemoryBankConfigDict] + """""" + + +_UpdateMemoryBankRequestParametersOrDict = Union[ + _UpdateMemoryBankRequestParameters, _UpdateMemoryBankRequestParametersDict +] + + class DeleteMemoryBankConfig(_common.BaseModel): """Config for delete memory bank.""" diff --git a/tests/unit/agentplatform/genai/replays/test_memory_banks_private_update.py b/tests/unit/agentplatform/genai/replays/test_memory_banks_private_update.py new file mode 100644 index 0000000000..1f673ee964 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memory_banks_private_update.py @@ -0,0 +1,223 @@ +"""Tests for updating a Memory Bank using the private update method.""" + +# 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,g-bad-import-order + +import pytest +import time + +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper + +LOCATION = "test-location" +PROJECT = "test-project" +LOCATION_NAME = f"projects/{PROJECT}/locations/{LOCATION}" +OPERATION_NAME = "projects/test-project/locations/test-location/operations/op-123" +MEMORY_BANK_NAME = "projects/test-project/locations/test-location/reasoningEngines/123" +DISPLAY_NAME = "My Test Memory Bank" +DESCRIPTION = "My Test Memory Bank Description" +PENDING_OP = { + "name": OPERATION_NAME, + "done": False, +} +FINISHED_OP = { + "name": OPERATION_NAME, + "done": True, + "response": {"name": MEMORY_BANK_NAME}, +} + +GENERATION_CONFIG = { + "model": ( + "projects/test-project/locations/test-location/publishers/google/models/gemini-3.5-flash" + ), +} +UPDATED_GENERATION_CONFIG = { + "model": ( + "projects/test-project/locations/test-location/publishers/google/models/gemini-3.8-flash" + ), +} +SIMILARITY_SEARCH_CONFIG = { + "embedding_model": ( + "projects/test-project/locations/test-location/" + "publishers/google/models/gemini-embedding-2" + ) +} +UNSTRUCTURED_MEMORY_CONFIGS = [ + { + "memory_topics": [ + {"managed_memory_topic": {"managed_topic_enum": "USER_PERSONAL_INFO"}} + ], + "consolidation_config": {"revisions_per_candidate_count": 1}, + } +] +TTL_CONFIG = {"memory_revision_default_ttl": f"{365 * 24 * 60 * 60}s"} +UPDATED_TTL_CONFIG = {"memory_revision_default_ttl": f"{180 * 24 * 60 * 60}s"} +MEMORY_SCHEMA = { + "properties": { + "name": { + "description": "User's name", + "type": "string", + } + }, + "type": "object", +} +# The SDK uses the alias `memory_schema` while the API uses `schema`. The SDK +# inner workings will convert between the two. +STRUCTURED_MEMORY_SCHEMA_CONFIGS = [ + { + "scopeKeys": ["user_id"], + "schemaConfigs": [ + { + "id": "user-profile", + "memory_schema": MEMORY_SCHEMA, + } + ], + } +] + + +def test_update_memory_bank(client): + memory_bank = client.memory_banks.create( + managed_semantic_memory_config={ + "generation_config": GENERATION_CONFIG, + "similarity_search_config": SIMILARITY_SEARCH_CONFIG, + "unstructured_memory_configs": UNSTRUCTURED_MEMORY_CONFIGS, + "structured_memory_configs": STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttl_config": TTL_CONFIG, + "disable_memory_revisions": True, + }, + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + }, + ) + try: + memory_bank = client.memory_banks.get(name=memory_bank.name) + assert memory_bank.name == memory_bank.name + assert memory_bank.display_name == DISPLAY_NAME + assert memory_bank.description == DESCRIPTION + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=GENERATION_CONFIG, + similarity_search_config=SIMILARITY_SEARCH_CONFIG, + unstructured_memory_configs=UNSTRUCTURED_MEMORY_CONFIGS, + structured_memory_configs=STRUCTURED_MEMORY_SCHEMA_CONFIGS, + ttl_config=TTL_CONFIG, + disable_memory_revisions=True, + ) + ) + # Update the memory bank. + client.memory_banks._update( + name=memory_bank.name, + config={ + "display_name": "Updated Display Name", + "description": "Updated Description", + }, + memory_bank_config={ + "ttl_config": UPDATED_TTL_CONFIG, + "generation_config": UPDATED_GENERATION_CONFIG, + }, + ) + # Give time for the operation to complete. + time.sleep(10) + # Fetch the updated memory bank. + memory_bank = client.memory_banks.get(name=memory_bank.name) + assert memory_bank.display_name == "Updated Display Name" + assert memory_bank.description == "Updated Description" + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=UPDATED_GENERATION_CONFIG, + ttl_config=UPDATED_TTL_CONFIG, + ) + ) + + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks._update", +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_create_memory_bank_async(client): + memory_bank = await client.aio.memory_banks.create( + managed_semantic_memory_config={ + "generation_config": GENERATION_CONFIG, + "similarity_search_config": SIMILARITY_SEARCH_CONFIG, + "unstructured_memory_configs": UNSTRUCTURED_MEMORY_CONFIGS, + "structured_memory_configs": STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttl_config": TTL_CONFIG, + "disable_memory_revisions": True, + }, + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + }, + ) + try: + memory_bank = await client.aio.memory_banks.get(name=memory_bank.name) + assert memory_bank.name == memory_bank.name + assert memory_bank.display_name == DISPLAY_NAME + assert memory_bank.description == DESCRIPTION + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=GENERATION_CONFIG, + similarity_search_config=SIMILARITY_SEARCH_CONFIG, + unstructured_memory_configs=UNSTRUCTURED_MEMORY_CONFIGS, + structured_memory_configs=STRUCTURED_MEMORY_SCHEMA_CONFIGS, + ttl_config=TTL_CONFIG, + disable_memory_revisions=True, + ) + ) + # Update the memory bank. + await client.aio.memory_banks._update( + name=memory_bank.name, + memory_bank_config={ + "ttl_config": UPDATED_TTL_CONFIG, + "generation_config": UPDATED_GENERATION_CONFIG, + }, + config={ + "display_name": "Updated Display Name", + "description": "Updated Description", + }, + ) + # Give time for the operation to complete. + time.sleep(10) + # Fetch the updated memory bank. + memory_bank = await client.aio.memory_banks.get(name=memory_bank.name) + assert memory_bank.display_name == "Updated Display Name" + assert memory_bank.description == "Updated Description" + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=UPDATED_GENERATION_CONFIG, + ttl_config=UPDATED_TTL_CONFIG, + ) + ) + + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_memory_banks_update.py b/tests/unit/agentplatform/genai/replays/test_memory_banks_update.py new file mode 100644 index 0000000000..1ad185a105 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memory_banks_update.py @@ -0,0 +1,218 @@ +"""Tests for updating a Memory Bank.""" + +# 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,g-bad-import-order + +import pytest + +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper + +LOCATION = "test-location" +PROJECT = "test-project" +LOCATION_NAME = f"projects/{PROJECT}/locations/{LOCATION}" +OPERATION_NAME = "projects/test-project/locations/test-location/operations/op-123" +MEMORY_BANK_NAME = "projects/test-project/locations/test-location/reasoningEngines/123" +DISPLAY_NAME = "My Test Memory Bank" +DESCRIPTION = "My Test Memory Bank Description" +PENDING_OP = { + "name": OPERATION_NAME, + "done": False, +} +FINISHED_OP = { + "name": OPERATION_NAME, + "done": True, + "response": {"name": MEMORY_BANK_NAME}, +} + +GENERATION_CONFIG = { + "model": ( + "projects/test-project/locations/test-location/publishers/google/models/gemini-3.5-flash" + ), +} +UPDATED_GENERATION_CONFIG = { + "model": ( + "projects/test-project/locations/test-location/publishers/google/models/gemini-3.8-flash" + ), +} +SIMILARITY_SEARCH_CONFIG = { + "embedding_model": ( + "projects/test-project/locations/test-location/" + "publishers/google/models/gemini-embedding-2" + ) +} +UNSTRUCTURED_MEMORY_CONFIGS = [ + { + "memory_topics": [ + {"managed_memory_topic": {"managed_topic_enum": "USER_PERSONAL_INFO"}} + ], + "consolidation_config": {"revisions_per_candidate_count": 1}, + } +] +TTL_CONFIG = {"memory_revision_default_ttl": f"{365 * 24 * 60 * 60}s"} +UPDATED_TTL_CONFIG = {"memory_revision_default_ttl": f"{180 * 24 * 60 * 60}s"} +MEMORY_SCHEMA = { + "properties": { + "name": { + "description": "User's name", + "type": "string", + } + }, + "type": "object", +} +# The SDK uses the alias `memory_schema` while the API uses `schema`. The SDK +# inner workings will convert between the two. +STRUCTURED_MEMORY_SCHEMA_CONFIGS = [ + { + "scopeKeys": ["user_id"], + "schemaConfigs": [ + { + "id": "user-profile", + "memory_schema": MEMORY_SCHEMA, + } + ], + } +] + + +def test_update_memory_bank(client): + memory_bank = client.memory_banks.create( + managed_semantic_memory_config={ + "generation_config": GENERATION_CONFIG, + "similarity_search_config": SIMILARITY_SEARCH_CONFIG, + "unstructured_memory_configs": UNSTRUCTURED_MEMORY_CONFIGS, + "structured_memory_configs": STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttl_config": TTL_CONFIG, + "disable_memory_revisions": True, + }, + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + }, + ) + try: + memory_bank = client.memory_banks.get(name=memory_bank.name) + assert memory_bank.name == memory_bank.name + assert memory_bank.display_name == DISPLAY_NAME + assert memory_bank.description == DESCRIPTION + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=GENERATION_CONFIG, + similarity_search_config=SIMILARITY_SEARCH_CONFIG, + unstructured_memory_configs=UNSTRUCTURED_MEMORY_CONFIGS, + structured_memory_configs=STRUCTURED_MEMORY_SCHEMA_CONFIGS, + ttl_config=TTL_CONFIG, + disable_memory_revisions=True, + ) + ) + # Update the memory bank. + memory_bank = client.memory_banks.update( + name=memory_bank.name, + managed_semantic_memory_config={ + "ttl_config": UPDATED_TTL_CONFIG, + "generation_config": UPDATED_GENERATION_CONFIG, + }, + config={ + "display_name": "Updated Display Name", + "description": "Updated Description", + }, + ) + # Fetch the updated memory bank. + memory_bank = client.memory_banks.get(name=memory_bank.name) + assert memory_bank.display_name == "Updated Display Name" + assert memory_bank.description == "Updated Description" + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=UPDATED_GENERATION_CONFIG, + ttl_config=UPDATED_TTL_CONFIG, + ) + ) + + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.update", +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_create_memory_bank_async(client): + memory_bank = await client.aio.memory_banks.create( + managed_semantic_memory_config={ + "generation_config": GENERATION_CONFIG, + "similarity_search_config": SIMILARITY_SEARCH_CONFIG, + "unstructured_memory_configs": UNSTRUCTURED_MEMORY_CONFIGS, + "structured_memory_configs": STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttl_config": TTL_CONFIG, + "disable_memory_revisions": True, + }, + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + }, + ) + try: + memory_bank = await client.aio.memory_banks.get(name=memory_bank.name) + assert memory_bank.name == memory_bank.name + assert memory_bank.display_name == DISPLAY_NAME + assert memory_bank.description == DESCRIPTION + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=GENERATION_CONFIG, + similarity_search_config=SIMILARITY_SEARCH_CONFIG, + unstructured_memory_configs=UNSTRUCTURED_MEMORY_CONFIGS, + structured_memory_configs=STRUCTURED_MEMORY_SCHEMA_CONFIGS, + ttl_config=TTL_CONFIG, + disable_memory_revisions=True, + ) + ) + # Update the memory bank. + memory_bank = await client.aio.memory_banks.update( + name=memory_bank.name, + managed_semantic_memory_config={ + "ttl_config": UPDATED_TTL_CONFIG, + "generation_config": UPDATED_GENERATION_CONFIG, + }, + config={ + "display_name": "Updated Display Name", + "description": "Updated Description", + }, + ) + # Fetch the updated memory bank. + memory_bank = await client.aio.memory_banks.get(name=memory_bank.name) + assert memory_bank.display_name == "Updated Display Name" + assert memory_bank.description == "Updated Description" + assert ( + memory_bank.managed_semantic_memory_config + == types.ManagedSemanticMemoryConfig( + generation_config=UPDATED_GENERATION_CONFIG, + ttl_config=UPDATED_TTL_CONFIG, + ) + ) + + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank.name, force=True)