From 6dfe67856cd03919113149dd26a1e067eb186523 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 25 Aug 2026 17:07:48 -0700 Subject: [PATCH] feat(agentplatform): Support Memory Bank configs. feat(agentplatform): Support client.memory_banks.get feat(agentplatform): Support client.memory_banks.list PiperOrigin-RevId: 970892415 --- agentplatform/_genai/_memory_bank_utils.py | 64 ++ agentplatform/_genai/memory_banks.py | 980 +++++++++++++++++- agentplatform/_genai/types/__init__.py | 52 + agentplatform/_genai/types/common.py | 397 ++++++- .../replays/test_memories_private_create.py | 2 +- .../test_memories_retrieve_profiles.py | 36 +- .../genai/replays/test_memory_banks_create.py | 164 +++ .../test_memory_banks_ingest_events.py | 30 +- .../genai/replays/test_memory_banks_list.py | 62 ++ .../test_memory_banks_private_create.py | 96 ++ .../agentplatform/genai/test_memory_banks.py | 570 ++++++++++ 11 files changed, 2371 insertions(+), 82 deletions(-) create mode 100644 tests/unit/agentplatform/genai/replays/test_memory_banks_create.py create mode 100644 tests/unit/agentplatform/genai/replays/test_memory_banks_list.py create mode 100644 tests/unit/agentplatform/genai/replays/test_memory_banks_private_create.py create mode 100644 tests/unit/agentplatform/genai/test_memory_banks.py diff --git a/agentplatform/_genai/_memory_bank_utils.py b/agentplatform/_genai/_memory_bank_utils.py index 694fa4e3cf..5f1d6bba50 100644 --- a/agentplatform/_genai/_memory_bank_utils.py +++ b/agentplatform/_genai/_memory_bank_utils.py @@ -15,6 +15,7 @@ """Utility functions for memory banks.""" import asyncio +import json import re import time from typing import ( @@ -140,3 +141,66 @@ async def _await_async_operation( operation = await get_operation_fn(operation_name=operation.name) return operation + + +def _managed_semantic_memory_config_to_memory_bank_config( + semantic_memory_config: genai_types.ManagedSemanticMemoryConfigOrDict, +) -> genai_types.ReasoningEngineContextSpecMemoryBankConfigDict: + """Converts ManagedSemanticMemoryConfig to MemoryBankConfig.""" + if semantic_memory_config is None: + semantic_memory_config = {} + if isinstance(semantic_memory_config, dict): + semantic_memory_config = genai_types.ManagedSemanticMemoryConfig.model_validate( + semantic_memory_config + ) + elif not isinstance( + semantic_memory_config, genai_types.ManagedSemanticMemoryConfig + ): + raise TypeError( + "managed_semantic_memory_config must be a dict or " + "ManagedSemanticMemoryConfig, " + f"but got {type(semantic_memory_config)}." + ) + + memory_bank_config = json.loads(semantic_memory_config.model_dump_json()) + if "unstructured_memory_configs" in memory_bank_config: + memory_bank_config["customization_configs"] = memory_bank_config.pop( + "unstructured_memory_configs" + ) + return memory_bank_config + + +def _memory_bank_config_to_managed_semantic_memories_config( + memory_bank_config: genai_types.ReasoningEngineContextSpecMemoryBankConfig, +) -> genai_types.ManagedSemanticMemoryConfigDict: + """Converts MemoryBankConfig to ManagedSemanticMemoriesConfig.""" + memory_bank_config = json.loads(memory_bank_config.model_dump_json()) + if "customization_configs" in memory_bank_config: + memory_bank_config["unstructured_memory_configs"] = memory_bank_config.pop( + "customization_configs" + ) + return memory_bank_config + + +def _reasoning_engine_to_memory_bank( + reasoning_engine: genai_types.ReasoningEngine, +) -> genai_types.MemoryBank: + """Converts ReasoningEngine to MemoryBank.""" + if reasoning_engine.context_spec is not None: + semantic_memory_config = ( + _memory_bank_config_to_managed_semantic_memories_config( + reasoning_engine.context_spec.memory_bank_config + ) + ) + else: + semantic_memory_config = {} + memory_bank = genai_types.MemoryBank( + name=reasoning_engine.name, + create_time=reasoning_engine.create_time, + update_time=reasoning_engine.update_time, + display_name=reasoning_engine.display_name, + description=reasoning_engine.description, + encryption_spec=reasoning_engine.encryption_spec, + managed_semantic_memory_config=semantic_memory_config, + ) + return memory_bank diff --git a/agentplatform/_genai/memory_banks.py b/agentplatform/_genai/memory_banks.py index 0553178e53..cf04542b64 100644 --- a/agentplatform/_genai/memory_banks.py +++ b/agentplatform/_genai/memory_banks.py @@ -19,13 +19,14 @@ import json import logging import typing -from typing import Any, Optional, Union +from typing import Any, AsyncIterator, Iterator, Optional, Union from urllib.parse import urlencode from google.genai import _api_module from google.genai import _common from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv +from google.genai.pagers import AsyncPager, Pager from . import _memory_bank_utils from . import types @@ -39,11 +40,40 @@ logger = logging.getLogger("agentplatform_genai.memorybanks") +def _CreateMemoryBankConfig_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"])) + + if getv(from_object, ["encryption_spec"]) is not None: + setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) + + return to_object + + def _CreateMemoryBankRequestParameters_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, ["config"]) is not None: + _CreateMemoryBankConfig_to_vertex(getv(from_object, ["config"]), to_object) + + 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 + ), + ) return to_object @@ -75,6 +105,17 @@ def _GetMemoryBankOperationParameters_to_vertex( return to_object +def _GetMemoryBankRequestParameters_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"])) + + return to_object + + def _IngestEventsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -155,10 +196,402 @@ def _IngestEventsRequestParameters_to_vertex( return to_object +def _ListMemoryBanksConfig_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, ["page_size"]) is not None: + setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) + + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) + + return to_object + + +def _ListMemoryBanksRequestParameters_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, ["config"]) is not None: + _ListMemoryBanksConfig_to_vertex(getv(from_object, ["config"]), to_object) + + 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 _ManagedSemanticMemoryConfig_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, ["generationConfig"]) is not None: + setv(to_object, ["generation_config"], getv(from_object, ["generationConfig"])) + + if getv(from_object, ["ttlConfig"]) is not None: + setv(to_object, ["ttl_config"], getv(from_object, ["ttlConfig"])) + + if getv(from_object, ["disableMemoryRevisions"]) is not None: + setv( + to_object, + ["disable_memory_revisions"], + getv(from_object, ["disableMemoryRevisions"]), + ) + + if getv(from_object, ["similaritySearchConfig"]) is not None: + setv( + to_object, + ["similarity_search_config"], + getv(from_object, ["similaritySearchConfig"]), + ) + + if getv(from_object, ["unstructuredMemoryConfigs"]) is not None: + setv( + to_object, + ["unstructured_memory_configs"], + [item for item in getv(from_object, ["unstructuredMemoryConfigs"])], + ) + + 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 _MemoryBankOperation_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"], + _MemoryBank_from_vertex(getv(from_object, ["response"]), to_object), + ) + + return to_object + + +def _MemoryBank_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, ["managedSemanticMemoryConfig"]) is not None: + setv( + to_object, + ["managed_semantic_memory_config"], + _ManagedSemanticMemoryConfig_from_vertex( + getv(from_object, ["managedSemanticMemoryConfig"]), to_object + ), + ) + + if getv(from_object, ["displayName"]) is not None: + setv(to_object, ["display_name"], getv(from_object, ["displayName"])) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["createTime"]) is not None: + setv(to_object, ["create_time"], getv(from_object, ["createTime"])) + + if getv(from_object, ["updateTime"]) is not None: + setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) + + if getv(from_object, ["encryptionSpec"]) is not None: + setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"])) + + 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 _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 _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 + + class MemoryBanks(_api_module.BaseModule): def _create( - self, *, config: Optional[types.CreateMemoryBankConfigOrDict] = None + self, + *, + config: Optional[types.CreateMemoryBankConfigOrDict] = None, + memory_bank_config: Optional[ + types.ReasoningEngineContextSpecMemoryBankConfigOrDict + ] = None, ) -> types.MemoryBankOperation: """ Creates a new Memory Bank. @@ -166,6 +599,7 @@ def _create( parameter_model = types._CreateMemoryBankRequestParameters( config=config, + memory_bank_config=memory_bank_config, ) request_url_dict: Optional[dict[str, str]] @@ -201,6 +635,9 @@ def _create( 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=( @@ -233,12 +670,168 @@ def _delete( config: Optional[types.DeleteMemoryBankConfigOrDict] = None, ) -> types.DeleteMemoryBankOperation: """ - Deletes a memory bank. + Deletes a memory bank. + """ + + parameter_model = types._DeleteMemoryBankRequestParameters( + name=name, + force=force, + 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 = _DeleteMemoryBankRequestParameters_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("delete", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteMemoryBankOperation._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 _get( + self, *, name: str, config: Optional[types.GetMemoryBankConfigOrDict] = None + ) -> types.ReasoningEngine: + """ + Get a Memory Bank instance. + """ + + parameter_model = types._GetMemoryBankRequestParameters( + name=name, + 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 = _GetMemoryBankRequestParameters_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("get", path, request_dict, http_options) + + 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=( + { + "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 _ingest_events( + self, + *, + name: str, + stream_id: Optional[str] = None, + direct_contents_source: Optional[ + types.IngestionDirectContentsSourceOrDict + ] = None, + scope: Optional[dict[str, str]] = None, + generation_trigger_config: Optional[ + types.MemoryGenerationTriggerConfigOrDict + ] = None, + config: Optional[types.IngestEventsConfigOrDict] = None, + ) -> types.MemoryBankIngestEventsOperation: + """ + Ingest events into a Memory Bank. """ - parameter_model = types._DeleteMemoryBankRequestParameters( + parameter_model = types._IngestEventsRequestParameters( name=name, - force=force, + stream_id=stream_id, + direct_contents_source=direct_contents_source, + scope=scope, + generation_trigger_config=generation_trigger_config, config=config, ) @@ -248,12 +841,12 @@ def _delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteMemoryBankRequestParameters_to_vertex(parameter_model) + request_dict = _IngestEventsRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: - path = "{name}".format_map(request_url_dict) + path = "{name}/memories:ingestEvents".format_map(request_url_dict) else: - path = "{name}" + path = "{name}/memories:ingestEvents" query_params = request_dict.get("_query") if query_params: @@ -271,11 +864,11 @@ def _delete( request_dict = _common.convert_to_dict(request_dict) request_dict = _common.encode_unserializable_types(request_dict) - response = self._api_client.request("delete", path, request_dict, http_options) + 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.DeleteMemoryBankOperation._from_response( + return_value = types.MemoryBankIngestEventsOperation._from_response( response=response_dict, kwargs=( { @@ -299,30 +892,14 @@ def _delete( self._api_client._verify_response(return_value) return return_value - def _ingest_events( - self, - *, - name: str, - stream_id: Optional[str] = None, - direct_contents_source: Optional[ - types.IngestionDirectContentsSourceOrDict - ] = None, - scope: Optional[dict[str, str]] = None, - generation_trigger_config: Optional[ - types.MemoryGenerationTriggerConfigOrDict - ] = None, - config: Optional[types.IngestEventsConfigOrDict] = None, - ) -> types.MemoryBankIngestEventsOperation: + def _list( + self, *, config: Optional[types.ListMemoryBanksConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: """ - Ingest events into a Memory Bank. + Lists Memory Banks. """ - parameter_model = types._IngestEventsRequestParameters( - name=name, - stream_id=stream_id, - direct_contents_source=direct_contents_source, - scope=scope, - generation_trigger_config=generation_trigger_config, + parameter_model = types._ListMemoryBanksRequestParameters( config=config, ) @@ -332,12 +909,12 @@ def _ingest_events( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _IngestEventsRequestParameters_to_vertex(parameter_model) + request_dict = _ListMemoryBanksRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: - path = "{name}/memories:ingestEvents".format_map(request_url_dict) + path = "reasoningEngines".format_map(request_url_dict) else: - path = "{name}/memories:ingestEvents" + path = "reasoningEngines" query_params = request_dict.get("_query") if query_params: @@ -355,11 +932,14 @@ def _ingest_events( 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 = self._api_client.request("get", path, request_dict, http_options) response_dict = {} if not response.body else json.loads(response.body) - return_value = types.MemoryBankIngestEventsOperation._from_response( + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + + return_value = types.ListReasoningEnginesResponse._from_response( response=response_dict, kwargs=( { @@ -427,6 +1007,9 @@ def _get_memory_bank_operation( 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=( @@ -459,10 +1042,25 @@ def memories(self) -> "memories_module.Memories": self._memories = importlib.import_module(".memories", __package__) return self._memories.Memories(self._api_client) # type: ignore[no-any-return] - def create(self): + def create( + self, + *, + managed_semantic_memory_config: Optional[ + types.ManagedSemanticMemoryConfigOrDict + ] = None, + config: Optional[types.CreateMemoryBankConfigOrDict] = None, + ) -> types.MemoryBank: """Creates a new Memory Bank.""" + memory_bank_config = ( + _memory_bank_utils._managed_semantic_memory_config_to_memory_bank_config( + managed_semantic_memory_config + ) + ) - operation = self._create() + operation = self._create( + memory_bank_config=memory_bank_config, + config=config, + ) operation = _memory_bank_utils._await_operation( operation_name=operation.name, @@ -587,11 +1185,74 @@ def ingest_events( raise RuntimeError(f"Failed to ingest events: {operation.error}") return operation + def get( + self, + *, + name: str, + config: Optional[types.GetMemoryBankConfigOrDict] = None, + ) -> types.MemoryBank: + """Gets a Memory Bank. + + Args: + name (str): + Required. A fully-qualified resource name or ID such as + "projects/123/locations/us-central1/reasoningEngines/456" or + a shortened name such as "reasoningEngines/456". + """ + api_resource = self._get(name=name, config=config) + memory_bank = _memory_bank_utils._reasoning_engine_to_memory_bank(api_resource) + return memory_bank + + def list( + self, *, config: Optional[types.ListMemoryBanksConfigOrDict] = None + ) -> Iterator[types.MemoryBank]: + """List all instances of Memory Bank matching the filter. + + Example Usage: + + .. code-block:: python + import agentplatform + + client = agentplatform.Client(project="my_project", location="us-central1") + for memory_bank in client.memory_banks.list( + config={"filter": "'display_name="My Custom Memory Bank"'}, + ): + print(memory_bank.name) + + Args: + config (ListMemoryBanksConfig): + Optional. The config for the memory banks to be listed. + + Returns: + Iterable[MemoryBank]: An iterable of Memory Banks matching the filter. + """ + + def transformed_list(*args, **kwargs) -> types.ListMemoryBanksResponse: + res = self._list(*args, **kwargs) + if getattr(res, "reasoning_engines", None): + res.reasoning_engines = [ + _memory_bank_utils._reasoning_engine_to_memory_bank(engine) + for engine in res.reasoning_engines + ] + return res + + return Pager( + "reasoning_engines", + transformed_list, + transformed_list(config=config), + config, + ) + class AsyncMemoryBanks(_api_module.BaseModule): async def _create( - self, *, config: Optional[types.CreateMemoryBankConfigOrDict] = None + self, + *, + config: Optional[types.CreateMemoryBankConfigOrDict] = None, + memory_bank_config: Optional[ + types.ReasoningEngineContextSpecMemoryBankConfigOrDict + ] = None, ) -> types.MemoryBankOperation: """ Creates a new Memory Bank. @@ -599,6 +1260,7 @@ async def _create( parameter_model = types._CreateMemoryBankRequestParameters( config=config, + memory_bank_config=memory_bank_config, ) request_url_dict: Optional[dict[str, str]] @@ -636,6 +1298,9 @@ async def _create( 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=( @@ -736,6 +1401,80 @@ async def _delete( self._api_client._verify_response(return_value) return return_value + async def _get( + self, *, name: str, config: Optional[types.GetMemoryBankConfigOrDict] = None + ) -> types.ReasoningEngine: + """ + Get a Memory Bank instance. + """ + + parameter_model = types._GetMemoryBankRequestParameters( + name=name, + 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 = _GetMemoryBankRequestParameters_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( + "get", path, request_dict, http_options + ) + + 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=( + { + "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 _ingest_events( self, *, @@ -822,6 +1561,79 @@ async def _ingest_events( self._api_client._verify_response(return_value) return return_value + async def _list( + self, *, config: Optional[types.ListMemoryBanksConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """ + Lists Memory Banks. + """ + + parameter_model = types._ListMemoryBanksRequestParameters( + 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 = _ListMemoryBanksRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + 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( + "get", path, request_dict, http_options + ) + + 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=( + { + "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 _get_memory_bank_operation( self, *, @@ -868,6 +1680,9 @@ async def _get_memory_bank_operation( 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=( @@ -900,10 +1715,25 @@ def memories(self) -> "memories_module.AsyncMemories": self._memories = importlib.import_module(".memories", __package__) return self._memories.AsyncMemories(self._api_client) # type: ignore[no-any-return] - async def create(self): + async def create( + self, + *, + managed_semantic_memory_config: Optional[ + types.ManagedSemanticMemoryConfigOrDict + ] = None, + config: Optional[types.CreateMemoryBankConfigOrDict] = None, + ) -> types.MemoryBank: """Creates a new Memory Bank.""" + memory_bank_config = ( + _memory_bank_utils._managed_semantic_memory_config_to_memory_bank_config( + managed_semantic_memory_config + ) + ) - operation = await self._create() + operation = await self._create( + memory_bank_config=memory_bank_config, + config=config, + ) operation = await _memory_bank_utils._await_async_operation( operation_name=operation.name, @@ -1027,3 +1857,71 @@ async def ingest_events( if operation.error: raise RuntimeError(f"Failed to ingest events: {operation.error}") return operation + + async def get( + self, + *, + name: str, + config: Optional[types.GetMemoryBankConfigOrDict] = None, + ) -> types.MemoryBank: + """Gets a Memory Bank. + + Args: + name (str): + Required. A fully-qualified resource name or ID such as + "projects/123/locations/us-central1/reasoningEngines/456" or + a shortened name such as "reasoningEngines/456". + """ + api_resource = await self._get(name=name, config=config) + memory_bank = _memory_bank_utils._reasoning_engine_to_memory_bank(api_resource) + return memory_bank + + async def _list_pager( + self, *, config: Optional[types.ListMemoryBanksConfigOrDict] = None + ) -> AsyncPager[types.ReasoningEngine]: + return AsyncPager( + "reasoning_engines", + self._list, + await self._list(config=config), + config, + ) + + async def list( + self, *, config: Optional[types.ListMemoryBanksConfigOrDict] = None + ) -> AsyncIterator[types.MemoryBank]: + """List all instances of Memory Bank matching the filter. + + Example Usage: + + .. code-block:: python + import agentplatform + + client = agentplatform.Client(project="my_project", location="us-central1") + async for memory_bank in await client.memory_banks.list( + config={"filter": "'display_name="My Custom Memory Bank"'}, + ): + print(memory_bank.name) + + Args: + config (ListMemoryBanksConfig): + Optional. The config for the memory banks to be listed. + + Returns: + Iterable[MemoryBank]: An iterable of Memory Banks matching the filter. + """ + + async def transformed_list(*args, **kwargs): + res = await self._list(*args, **kwargs) + if getattr(res, "reasoning_engines", None): + res.reasoning_engines = [ + _memory_bank_utils._reasoning_engine_to_memory_bank(engine) + for engine in res.reasoning_engines + ] + return res + + return AsyncPager( + "reasoning_engines", + transformed_list, + await transformed_list(config=config), + config, + ) diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 05573aefec..9b1dec968e 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -111,6 +111,7 @@ from .common import _GetGenerateMemoriesOperationParameters from .common import _GetImportFilesOperationParameters from .common import _GetMemoryBankOperationParameters +from .common import _GetMemoryBankRequestParameters from .common import _GetMemoryOperationParameters from .common import _GetMemoryRequestParameters from .common import _GetMemoryRevisionRequestParameters @@ -148,6 +149,7 @@ from .common import _ListEvaluationMetricsParameters from .common import _ListEvaluationSetsParameters from .common import _ListMemoriesRequestParameters +from .common import _ListMemoryBanksRequestParameters from .common import _ListMemoryRevisionsRequestParameters from .common import _ListMultimodalDatasetsRequestParameters from .common import _ListPublisherModelsRequestParameters @@ -965,6 +967,9 @@ from .common import GetImportFilesOperationConfig from .common import GetImportFilesOperationConfigDict from .common import GetImportFilesOperationConfigOrDict +from .common import GetMemoryBankConfig +from .common import GetMemoryBankConfigDict +from .common import GetMemoryBankConfigOrDict from .common import GetMemoryBankOperationConfig from .common import GetMemoryBankOperationConfigDict from .common import GetMemoryBankOperationConfigOrDict @@ -1163,6 +1168,12 @@ from .common import ListMemoriesResponse from .common import ListMemoriesResponseDict from .common import ListMemoriesResponseOrDict +from .common import ListMemoryBanksConfig +from .common import ListMemoryBanksConfigDict +from .common import ListMemoryBanksConfigOrDict +from .common import ListMemoryBanksResponse +from .common import ListMemoryBanksResponseDict +from .common import ListMemoryBanksResponseOrDict from .common import ListMemoryRevisionsConfig from .common import ListMemoryRevisionsConfigDict from .common import ListMemoryRevisionsConfigOrDict @@ -1264,6 +1275,21 @@ from .common import MachineSpec from .common import MachineSpecDict from .common import MachineSpecOrDict +from .common import ManagedSemanticMemoryConfig +from .common import ManagedSemanticMemoryConfigDict +from .common import ManagedSemanticMemoryConfigGenerationConfig +from .common import ManagedSemanticMemoryConfigGenerationConfigDict +from .common import ManagedSemanticMemoryConfigGenerationConfigOrDict +from .common import ManagedSemanticMemoryConfigOrDict +from .common import ManagedSemanticMemoryConfigSimilaritySearchConfig +from .common import ManagedSemanticMemoryConfigSimilaritySearchConfigDict +from .common import ManagedSemanticMemoryConfigSimilaritySearchConfigOrDict +from .common import ManagedSemanticMemoryConfigTtlConfig +from .common import ManagedSemanticMemoryConfigTtlConfigDict +from .common import ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig +from .common import ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict +from .common import ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigOrDict +from .common import ManagedSemanticMemoryConfigTtlConfigOrDict from .common import ManagedTopicEnum from .common import MapInstance from .common import MapInstanceDict @@ -3050,6 +3076,21 @@ "CreateMemoryBankConfig", "CreateMemoryBankConfigDict", "CreateMemoryBankConfigOrDict", + "ManagedSemanticMemoryConfigGenerationConfig", + "ManagedSemanticMemoryConfigGenerationConfigDict", + "ManagedSemanticMemoryConfigGenerationConfigOrDict", + "ManagedSemanticMemoryConfigSimilaritySearchConfig", + "ManagedSemanticMemoryConfigSimilaritySearchConfigDict", + "ManagedSemanticMemoryConfigSimilaritySearchConfigOrDict", + "ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig", + "ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict", + "ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigOrDict", + "ManagedSemanticMemoryConfigTtlConfig", + "ManagedSemanticMemoryConfigTtlConfigDict", + "ManagedSemanticMemoryConfigTtlConfigOrDict", + "ManagedSemanticMemoryConfig", + "ManagedSemanticMemoryConfigDict", + "ManagedSemanticMemoryConfigOrDict", "MemoryBank", "MemoryBankDict", "MemoryBankOrDict", @@ -3062,6 +3103,9 @@ "DeleteMemoryBankOperation", "DeleteMemoryBankOperationDict", "DeleteMemoryBankOperationOrDict", + "GetMemoryBankConfig", + "GetMemoryBankConfigDict", + "GetMemoryBankConfigOrDict", "IngestionDirectContentsSourceEvent", "IngestionDirectContentsSourceEventDict", "IngestionDirectContentsSourceEventOrDict", @@ -3077,6 +3121,9 @@ "MemoryBankIngestEventsOperation", "MemoryBankIngestEventsOperationDict", "MemoryBankIngestEventsOperationOrDict", + "ListMemoryBanksConfig", + "ListMemoryBanksConfigDict", + "ListMemoryBanksConfigOrDict", "GetMemoryBankOperationConfig", "GetMemoryBankOperationConfigDict", "GetMemoryBankOperationConfigOrDict", @@ -4385,6 +4432,9 @@ "DeployOption", "DeployOptionDict", "DeployOptionOrDict", + "ListMemoryBanksResponse", + "ListMemoryBanksResponseDict", + "ListMemoryBanksResponseOrDict", "A2aTaskState", "Role", "State", @@ -4503,7 +4553,9 @@ "_UpdateAgentEngineRequestParameters", "_CreateMemoryBankRequestParameters", "_DeleteMemoryBankRequestParameters", + "_GetMemoryBankRequestParameters", "_IngestEventsRequestParameters", + "_ListMemoryBanksRequestParameters", "_GetMemoryBankOperationParameters", "_CreateMemoryRequestParameters", "_DeleteMemoryRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index e95aa77edf..6d52bee882 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -8778,7 +8778,7 @@ class MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=Fals class MemoryBankCustomizationConfig(_common.BaseModel): - """Represents configuration for organizing natural language memories for a particular scope.""" + """Represents configuration for organizing natural language memories.""" enable_third_person_memories: Optional[bool] = Field( default=None, @@ -8811,7 +8811,7 @@ class MemoryBankCustomizationConfig(_common.BaseModel): class MemoryBankCustomizationConfigDict(TypedDict, total=False): - """Represents configuration for organizing natural language memories for a particular scope.""" + """Represents configuration for organizing natural language memories.""" enable_third_person_memories: Optional[bool] """Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""" @@ -11174,6 +11174,21 @@ class CreateMemoryBankConfig(_common.BaseModel): 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.""" + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""The encryption spec to be used for the Memory Bank.""", + ) class CreateMemoryBankConfigDict(TypedDict, total=False): @@ -11182,6 +11197,19 @@ class CreateMemoryBankConfigDict(TypedDict, total=False): 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.""" + + encryption_spec: Optional[genai_types.EncryptionSpec] + """The encryption spec to be used for the Memory Bank.""" + CreateMemoryBankConfigOrDict = Union[CreateMemoryBankConfig, CreateMemoryBankConfigDict] @@ -11190,6 +11218,9 @@ class _CreateMemoryBankRequestParameters(_common.BaseModel): """Parameters for creating memory banks.""" config: Optional[CreateMemoryBankConfig] = Field(default=None, description="""""") + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( + default=None, description="""""" + ) class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): @@ -11198,12 +11229,229 @@ class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): config: Optional[CreateMemoryBankConfigDict] """""" + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + """""" + _CreateMemoryBankRequestParametersOrDict = Union[ _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict ] +class ManagedSemanticMemoryConfigGenerationConfig(_common.BaseModel): + """The configuration for generating memories.""" + + model: Optional[str] = Field( + default=None, + description="""The model used to generate memories. + + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + ) + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""The configuration for triggering memory generation.""", + ) + + +class ManagedSemanticMemoryConfigGenerationConfigDict(TypedDict, total=False): + """The configuration for generating memories.""" + + model: Optional[str] + """The model used to generate memories. + + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """The configuration for triggering memory generation.""" + + +ManagedSemanticMemoryConfigGenerationConfigOrDict = Union[ + ManagedSemanticMemoryConfigGenerationConfig, + ManagedSemanticMemoryConfigGenerationConfigDict, +] + + +class ManagedSemanticMemoryConfigSimilaritySearchConfig(_common.BaseModel): + """The configuration for similarity search.""" + + embedding_model: Optional[str] = Field( + default=None, + description="""The model used to generate embeddings to look up similar memories. + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + ) + + +class ManagedSemanticMemoryConfigSimilaritySearchConfigDict(TypedDict, total=False): + """The configuration for similarity search.""" + + embedding_model: Optional[str] + """The model used to generate embeddings to look up similar memories. + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + + +ManagedSemanticMemoryConfigSimilaritySearchConfigOrDict = Union[ + ManagedSemanticMemoryConfigSimilaritySearchConfig, + ManagedSemanticMemoryConfigSimilaritySearchConfigDict, +] + + +class ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig(_common.BaseModel): + """The configuration for granular TTL.""" + + create_ttl: Optional[str] = Field( + default=None, + description="""Optional. The TTL duration for memories uploaded via + CreateMemory.""", + ) + generate_created_ttl: Optional[str] = Field( + default=None, + description="""Optional. The TTL duration for memories generated via + GenerateMemories.""", + ) + generate_updated_ttl: Optional[str] = Field( + default=None, + description="""Optional. The TTL duration for memories updated via + GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). + In the case of an UPDATE action, the `expire_time` of the existing memory + will be updated to the new value (now + TTL).""", + ) + + +class ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict(TypedDict, total=False): + """The configuration for granular TTL.""" + + create_ttl: Optional[str] + """Optional. The TTL duration for memories uploaded via + CreateMemory.""" + + generate_created_ttl: Optional[str] + """Optional. The TTL duration for memories generated via + GenerateMemories.""" + + generate_updated_ttl: Optional[str] + """Optional. The TTL duration for memories updated via + GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). + In the case of an UPDATE action, the `expire_time` of the existing memory + will be updated to the new value (now + TTL).""" + + +ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigOrDict = Union[ + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig, + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict, +] + + +class ManagedSemanticMemoryConfigTtlConfig(_common.BaseModel): + """The configuration for automatic TTL ('time-to-live') of the memories.""" + + default_ttl: Optional[str] = Field( + default=None, + description="""The default TTL for 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.""", + ) + granular_ttl_config: Optional[ + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig + ] = Field(default=None, description="""The granular TTL config for memories.""") + memory_revision_default_ttl: Optional[str] = Field( + default=None, + description="""The default TTL for memory revisions 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.""", + ) + + +class ManagedSemanticMemoryConfigTtlConfigDict(TypedDict, total=False): + """The configuration for automatic TTL ('time-to-live') of the memories.""" + + default_ttl: Optional[str] + """The default TTL for 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.""" + + granular_ttl_config: Optional[ + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict + ] + """The granular TTL config for memories.""" + + memory_revision_default_ttl: Optional[str] + """The default TTL for memory revisions 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.""" + + +ManagedSemanticMemoryConfigTtlConfigOrDict = Union[ + ManagedSemanticMemoryConfigTtlConfig, ManagedSemanticMemoryConfigTtlConfigDict +] + + +class ManagedSemanticMemoryConfig(_common.BaseModel): + """The configuration for managed semantic memory.""" + + generation_config: Optional[ManagedSemanticMemoryConfigGenerationConfig] = Field( + default=None, description="""Represents configuration for LLMs calls.""" + ) + ttl_config: Optional[ManagedSemanticMemoryConfigTtlConfig] = Field( + default=None, + description="""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.""", + ) + disable_memory_revisions: Optional[bool] = Field( + default=None, + description="""If true, no memory revisions will be created for any requests to + Memory Bank.""", + ) + similarity_search_config: Optional[ + ManagedSemanticMemoryConfigSimilaritySearchConfig + ] = Field( + default=None, + description="""Configuration for how to perform similarity search on memories.""", + ) + unstructured_memory_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( + default=None, + description="""Configuration for how to customize Memory Bank behavior for a + particular scope for unstructured memories.""", + ) + structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + default=None, + description="""Configuration for organizing structured memories for a particular + scope.""", + ) + + +class ManagedSemanticMemoryConfigDict(TypedDict, total=False): + """The configuration for managed semantic memory.""" + + generation_config: Optional[ManagedSemanticMemoryConfigGenerationConfigDict] + """Represents configuration for LLMs calls.""" + + ttl_config: Optional[ManagedSemanticMemoryConfigTtlConfigDict] + """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.""" + + disable_memory_revisions: Optional[bool] + """If true, no memory revisions will be created for any requests to + Memory Bank.""" + + similarity_search_config: Optional[ + ManagedSemanticMemoryConfigSimilaritySearchConfigDict + ] + """Configuration for how to perform similarity search on memories.""" + + unstructured_memory_configs: Optional[list[MemoryBankCustomizationConfigDict]] + """Configuration for how to customize Memory Bank behavior for a + particular scope for unstructured memories.""" + + structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] + """Configuration for organizing structured memories for a particular + scope.""" + + +ManagedSemanticMemoryConfigOrDict = Union[ + ManagedSemanticMemoryConfig, ManagedSemanticMemoryConfigDict +] + + class MemoryBank(_common.BaseModel): """A memory bank.""" @@ -11211,6 +11459,27 @@ class MemoryBank(_common.BaseModel): 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.""", ) + managed_semantic_memory_config: Optional[ManagedSemanticMemoryConfig] = Field( + default=None, + description="""Represents the configuration for managed memories in Memory Bank. If not set, then the default configuration will be used.""", + ) + display_name: Optional[str] = Field( + default=None, description="""Represents the display name of the Memory Bank.""" + ) + description: Optional[str] = Field( + default=None, description="""Represents the description of the Memory Bank.""" + ) + create_time: Optional[datetime.datetime] = Field( + default=None, description="""Timestamp when this Memory Bank was created.""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Timestamp when this Memory Bank was most recently updated.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for a Memory Bank. If set, this Memory Bank and all sub-resources of this Memory Bank will be secured by this key.""", + ) class MemoryBankDict(TypedDict, total=False): @@ -11219,6 +11488,24 @@ class MemoryBankDict(TypedDict, total=False): name: 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.""" + managed_semantic_memory_config: Optional[ManagedSemanticMemoryConfigDict] + """Represents the configuration for managed memories in Memory Bank. If not set, then the default configuration will be used.""" + + display_name: Optional[str] + """Represents the display name of the Memory Bank.""" + + description: Optional[str] + """Represents the description of the Memory Bank.""" + + create_time: Optional[datetime.datetime] + """Timestamp when this Memory Bank was created.""" + + update_time: Optional[datetime.datetime] + """Timestamp when this Memory Bank was most recently updated.""" + + encryption_spec: Optional[genai_types.EncryptionSpecDict] + """Customer-managed encryption key spec for a Memory Bank. If set, this Memory Bank and all sub-resources of this Memory Bank will be secured by this key.""" + MemoryBankOrDict = Union[MemoryBank, MemoryBankDict] @@ -11360,6 +11647,48 @@ class DeleteMemoryBankOperationDict(TypedDict, total=False): ] +class GetMemoryBankConfig(_common.BaseModel): + """Config for getting a Memory Bank.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetMemoryBankConfigDict(TypedDict, total=False): + """Config for getting a Memory Bank.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetMemoryBankConfigOrDict = Union[GetMemoryBankConfig, GetMemoryBankConfigDict] + + +class _GetMemoryBankRequestParameters(_common.BaseModel): + """Parameters for getting a Memory Bank.""" + + name: Optional[str] = Field( + default=None, description="""Name of the Memory Bank.""" + ) + config: Optional[GetMemoryBankConfig] = Field(default=None, description="""""") + + +class _GetMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Memory Bank.""" + + name: Optional[str] + """Name of the Memory Bank.""" + + config: Optional[GetMemoryBankConfigDict] + """""" + + +_GetMemoryBankRequestParametersOrDict = Union[ + _GetMemoryBankRequestParameters, _GetMemoryBankRequestParametersDict +] + + class IngestionDirectContentsSourceEvent(_common.BaseModel): """The direct contents source event for ingesting events.""" @@ -11626,6 +11955,50 @@ class MemoryBankIngestEventsOperationDict(TypedDict, total=False): ] +class ListMemoryBanksConfig(_common.BaseModel): + """Config for listing Memory Banks.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + + +class ListMemoryBanksConfigDict(TypedDict, total=False): + """Config for listing Memory Banks.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + +ListMemoryBanksConfigOrDict = Union[ListMemoryBanksConfig, ListMemoryBanksConfigDict] + + +class _ListMemoryBanksRequestParameters(_common.BaseModel): + """Parameters for listing Memory Banks.""" + + config: Optional[ListMemoryBanksConfig] = Field(default=None, description="""""") + + +class _ListMemoryBanksRequestParametersDict(TypedDict, total=False): + """Parameters for listing Memory Banks.""" + + config: Optional[ListMemoryBanksConfigDict] + """""" + + +_ListMemoryBanksRequestParametersOrDict = Union[ + _ListMemoryBanksRequestParameters, _ListMemoryBanksRequestParametersDict +] + + class GetMemoryBankOperationConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( @@ -32389,3 +32762,23 @@ class DeployOptionDict(TypedDict, total=False): DeployOptionOrDict = Union[DeployOption, DeployOptionDict] + + +class ListMemoryBanksResponse(_common.BaseModel): + """The response for listing Memory Banks.""" + + memory_banks: Optional[list[MemoryBank]] = Field( + default=None, description="""The list of Memory Banks.""" + ) + + +class ListMemoryBanksResponseDict(TypedDict, total=False): + """The response for listing Memory Banks.""" + + memory_banks: Optional[list[MemoryBankDict]] + """The list of Memory Banks.""" + + +ListMemoryBanksResponseOrDict = Union[ + ListMemoryBanksResponse, ListMemoryBanksResponseDict +] diff --git a/tests/unit/agentplatform/genai/replays/test_memories_private_create.py b/tests/unit/agentplatform/genai/replays/test_memories_private_create.py index ed2904d9e4..91a4638dd3 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_create.py @@ -34,5 +34,5 @@ def test_private_create_memory(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._create", + test_method="memory_banks.memories._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py b/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py index ac9cc30941..3e55d2a683 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py @@ -14,12 +14,11 @@ # # pylint: disable=protected-access,bad-continuation,missing-function-docstring -from tests.unit.agentplatform.genai.replays import pytest_helper from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper def test_generate_and_retrieve_profile(client): - # TODO: Switch to Memory Bank for creation once it supports configs. customization_config = {"disable_natural_language_memories": True} memory_bank_customization_config = types.MemoryBankCustomizationConfig( **customization_config @@ -41,30 +40,23 @@ def test_generate_and_retrieve_profile(client): structured_memory_config_obj = types.StructuredMemoryConfig( **structured_memory_config ) - memory_bank = 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"}, - }, + memory_bank = client.memory_banks.create( + managed_semantic_memory_config={ + "unstructured_memory_configs": [memory_bank_customization_config], + "structured_memory_configs": [structured_memory_config_obj], + } ) try: - memory_bank = client.agent_engines.get(name=memory_bank.api_resource.name) - memory_bank_config = memory_bank.api_resource.context_spec.memory_bank_config - assert memory_bank_config.customization_configs == [ + memory_bank = client.memory_banks.get(name=memory_bank.name) + memory_config = memory_bank.managed_semantic_memory_config + assert memory_config.unstructured_memory_configs == [ memory_bank_customization_config ] - assert memory_bank_config.structured_memory_configs == [ - structured_memory_config_obj - ] + assert memory_config.structured_memory_configs == [structured_memory_config_obj] scope = {"user_id": "123"} client.memory_banks.memories.generate( - name=memory_bank.api_resource.name, + name=memory_bank.name, scope=scope, direct_contents_source={ "events": [{"content": {"parts": [{"text": "My name is Kim."}]}}] @@ -72,7 +64,7 @@ def test_generate_and_retrieve_profile(client): ) memories = list( client.memory_banks.memories.retrieve( - name=memory_bank.api_resource.name, + name=memory_bank.name, scope=scope, config={"memory_types": ["STRUCTURED_PROFILE"]}, ) @@ -81,13 +73,13 @@ def test_generate_and_retrieve_profile(client): assert memories[0].memory.structured_content is not None response = client.memory_banks.memories.retrieve_profiles( - name=memory_bank.api_resource.name, scope=scope + name=memory_bank.name, scope=scope ) assert len(response.profiles) == 1 finally: # Clean up resources. - client.memory_banks.delete(name=memory_bank.api_resource.name, force=True) + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( diff --git a/tests/unit/agentplatform/genai/replays/test_memory_banks_create.py b/tests/unit/agentplatform/genai/replays/test_memory_banks_create.py new file mode 100644 index 0000000000..d7577f6496 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memory_banks_create.py @@ -0,0 +1,164 @@ +# 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" + ), +} +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"} +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_create_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, + ) + ) + + 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.create", +) + +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, + ) + ) + + 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_ingest_events.py b/tests/unit/agentplatform/genai/replays/test_memory_banks_ingest_events.py index 76f45adca0..f3c8b96384 100644 --- a/tests/unit/agentplatform/genai/replays/test_memory_banks_ingest_events.py +++ b/tests/unit/agentplatform/genai/replays/test_memory_banks_ingest_events.py @@ -93,23 +93,21 @@ def test_ingest_events(client): "page_size": 1, }, ).page - # TODO: Re-enable this test once the bug resulting in no metadata being - # applied is fixed. - # # With `wait_for_completion` and `force_flush` set to True, there should be - # # memories immediately after the call. - # assert len(memories) >= 1 - # # The user-provided `metadata` should be applied to the generated memory. - # assert memories[0].memory.metadata["topic"].string_value == "jobs" + # With `wait_for_completion` and `force_flush` set to True, there should be + # memories immediately after the call. + assert len(memories) >= 1 + # The user-provided `metadata` should be applied to the generated memory. + assert memories[0].memory.metadata["topic"].string_value == "jobs" - # # The user-provided `revision_labels` are applied to the generated memory's - # # revision (not the Memory itself), so list the memory's revisions to verify. - # revisions = list( - # client.memory_banks.memories.revisions.list( - # name=memories[0].memory.name, - # ) - # ) - # assert revisions - # assert revisions[0].labels == {"source": "ingest-events-test"} + # The user-provided `revision_labels` are applied to the generated memory's + # revision (not the Memory itself), so list the memory's revisions to verify. + revisions = list( + client.memory_banks.memories.revisions.list( + name=memories[0].memory.name, + ) + ) + assert revisions + assert revisions[0].labels == {"source": "ingest-events-test"} finally: time.sleep(10) diff --git a/tests/unit/agentplatform/genai/replays/test_memory_banks_list.py b/tests/unit/agentplatform/genai/replays/test_memory_banks_list.py new file mode 100644 index 0000000000..59a022a3de --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memory_banks_list.py @@ -0,0 +1,62 @@ +# 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 + + +def test_list_memory_banks(client): + memory_bank_1 = client.memory_banks.create() + memory_bank_2 = client.memory_banks.create() + try: + pager = client.memory_banks.list(config={"page_size": 1}) + assert len(pager.page) == 1 + # Calling list(pager) automatically handles pagination. + assert len(list(pager)) > 1 + assert isinstance(pager.page[0], types.MemoryBank) + + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank_1.name, force=True) + client.memory_banks.delete(name=memory_bank_2.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.list", +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_list_memory_banks_async(client): + memory_bank_1 = await client.aio.memory_banks.create() + memory_bank_2 = await client.aio.memory_banks.create() + try: + pager = await client.aio.memory_banks.list(config={"page_size": 1}) + memory_list = [item async for item in pager] + # Calling list(pager) automatically handles pagination. + assert len(memory_list) > 1 + assert isinstance(memory_list[0], types.MemoryBank) + + finally: + # Clean up resources. + await client.aio.memory_banks.delete(name=memory_bank_1.name, force=True) + await client.aio.memory_banks.delete(name=memory_bank_2.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_memory_banks_private_create.py b/tests/unit/agentplatform/genai/replays/test_memory_banks_private_create.py new file mode 100644 index 0000000000..2a8af4f34d --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memory_banks_private_create.py @@ -0,0 +1,96 @@ +# 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 time + +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper + +DISPLAY_NAME = "My Test Memory Bank" +DESCRIPTION = "My Test Memory Bank Description" +GENERATION_CONFIG = { + "model": ( + "projects/test-project/locations/test-location/publishers/google/models/gemini-3.5-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"} +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_private_create_memory_bank(client): + memory_bank_config = { + "generation_config": GENERATION_CONFIG, + "similarity_search_config": SIMILARITY_SEARCH_CONFIG, + "customization_configs": UNSTRUCTURED_MEMORY_CONFIGS, + "structured_memory_configs": STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttl_config": TTL_CONFIG, + "disable_memory_revisions": True, + } + memory_operation = client.memory_banks._create( + memory_bank_config=memory_bank_config, + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + }, + ) + assert isinstance(memory_operation, types.MemoryBankOperation) + # Give time for the operation to complete. + time.sleep(10) + # Extract the ReasoningEngine name from the operation name. + name = "/".join(memory_operation.name.split("/")[0:-2]) + client.memory_banks.delete(name=name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks._create", +) diff --git a/tests/unit/agentplatform/genai/test_memory_banks.py b/tests/unit/agentplatform/genai/test_memory_banks.py new file mode 100644 index 0000000000..34cc8c4263 --- /dev/null +++ b/tests/unit/agentplatform/genai/test_memory_banks.py @@ -0,0 +1,570 @@ +# 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. +# + +# pylint: disable=protected-access + +import json +from unittest import mock +import google.auth.credentials +from agentplatform import _genai as genai +from agentplatform._genai import client as agentplatform_client +from agentplatform._genai import types as agentplatform_types +from google.genai import types as genai_types +import pytest + + +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": "gemini-3.5-flash", +} +SIMILARITY_SEARCH_CONFIG = { + "embedding_model": ( + "projects/test-project/locations/test-location/" + "publishers/google/models/text-embedding-005" + ) +} +ManagedTopicEnum = agentplatform_types.ManagedTopicEnum +UNSTRUCTURED_MEMORY_CONFIGS = [ + { + "memory_topics": [ + { + "managed_memory_topic": { + "managed_topic_enum": ManagedTopicEnum.USER_PERSONAL_INFO + } + }, + { + "managed_memory_topic": { + "managed_topic_enum": ManagedTopicEnum.USER_PREFERENCES + } + }, + { + "managed_memory_topic": { + "managed_topic_enum": ManagedTopicEnum.KEY_CONVERSATION_DETAILS + } + }, + { + "managed_memory_topic": { + "managed_topic_enum": ManagedTopicEnum.EXPLICIT_INSTRUCTIONS + } + }, + ], + "consolidation_config": {"revisions_per_candidate_count": 1}, + "generate_memories_examples": [], + "enable_third_person_memories": False, + } +] +TTL_CONFIG = {"memory_revision_default_ttl": f"{365 * 24 * 60 * 60}s"} +MEMORY_SCHEMA = { + "properties": { + "name": { + "description": "User's name", + "type": genai_types.Type.STRING, + } + }, + "type": genai_types.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, + } + ], + } +] +API_STRUCTURED_MEMORY_SCHEMA_CONFIGS = [ + { + "scopeKeys": ["user_id"], + "schemaConfigs": [ + { + "id": "user-profile", + "schema": MEMORY_SCHEMA, + } + ], + } +] + + +@pytest.fixture +def memory_banks_client(): + creds = mock.create_autospec(google.auth.credentials.Credentials, instance=True) + creds.token = "test_token" + client = agentplatform_client.Client( + project=PROJECT, location=LOCATION, credentials=creds + ) + return client.memory_banks + + +@pytest.fixture +def async_memory_banks_client(): + creds = mock.create_autospec(google.auth.credentials.Credentials, instance=True) + creds.token = "test_token" + client = agentplatform_client.Client( + project=PROJECT, location=LOCATION, credentials=creds + ) + return client.aio.memory_banks + + +class TestMemoryBanks: + """Tests for the Memory Banks module.""" + + def test_create_memory_bank_no_config(self, memory_banks_client): + """Tests the creation of a Memory Bank with no config.""" + reasoning_engine = { + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + "context_spec": {"memory_bank_config": {}}, + } + with mock.patch.object( + memory_banks_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(PENDING_OP)), + genai_types.HttpResponse(body=json.dumps(FINISHED_OP)), + genai_types.HttpResponse(body=json.dumps({"name": MEMORY_BANK_NAME})), + ] + + memory_bank = memory_banks_client.create( + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + } + ) + + request_mock.assert_has_calls( + [ + mock.call( + "post", + "reasoningEngines", + reasoning_engine, + None, + ), + mock.call( + "get", + OPERATION_NAME, + {"_url": {"operationName": OPERATION_NAME}}, + None, + ), + ] + ) + assert isinstance(memory_bank, genai.types.MemoryBank) + assert memory_bank.name == MEMORY_BANK_NAME + + @pytest.mark.asyncio + async def test_async_create_memory_bank_no_config(self, async_memory_banks_client): + """Tests the creation of a Memory Bank with no config.""" + + reasoning_engine = { + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + "context_spec": {"memory_bank_config": {}}, + } + with mock.patch.object( + async_memory_banks_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(PENDING_OP)), + genai_types.HttpResponse(body=json.dumps(FINISHED_OP)), + genai_types.HttpResponse(body=json.dumps({"name": MEMORY_BANK_NAME})), + ] + + memory_bank = await async_memory_banks_client.create( + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + } + ) + + request_mock.assert_has_calls( + [ + mock.call( + "post", + "reasoningEngines", + reasoning_engine, + None, + ), + mock.call( + "get", + OPERATION_NAME, + {"_url": {"operationName": OPERATION_NAME}}, + None, + ), + ] + ) + assert isinstance(memory_bank, genai.types.MemoryBank) + assert memory_bank.name == MEMORY_BANK_NAME + + def test_create_memory_bank_with_config(self, memory_banks_client): + """Tests the creation of a Memory Bank with config.""" + memory_bank_response = { + "name": MEMORY_BANK_NAME, + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + } + reasoning_engine = { + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + "context_spec": { + "memory_bank_config": { + "generationConfig": GENERATION_CONFIG, + "similaritySearchConfig": SIMILARITY_SEARCH_CONFIG, + "customizationConfigs": UNSTRUCTURED_MEMORY_CONFIGS, + "structuredMemoryConfigs": API_STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttlConfig": TTL_CONFIG, + "disableMemoryRevisions": False, + } + }, + } + with mock.patch.object( + memory_banks_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(PENDING_OP)), + genai_types.HttpResponse(body=json.dumps(FINISHED_OP)), + genai_types.HttpResponse(body=json.dumps(memory_bank_response)), + ] + + memory_bank = memory_banks_client.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": False, + }, + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + }, + ) + + request_mock.assert_has_calls( + [ + mock.call( + "post", + "reasoningEngines", + reasoning_engine, + None, + ), + mock.call( + "get", + OPERATION_NAME, + {"_url": {"operationName": OPERATION_NAME}}, + None, + ), + ] + ) + assert isinstance(memory_bank, genai.types.MemoryBank) + assert memory_bank.name == MEMORY_BANK_NAME + + @pytest.mark.asyncio + async def test_async_create_memory_bank_with_config( + self, async_memory_banks_client + ): + """Tests the creation of a Memory Bank with config.""" + + memory_bank_response = { + "name": MEMORY_BANK_NAME, + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + } + reasoning_engine = { + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + "context_spec": { + "memory_bank_config": { + "generationConfig": GENERATION_CONFIG, + "similaritySearchConfig": SIMILARITY_SEARCH_CONFIG, + "customizationConfigs": UNSTRUCTURED_MEMORY_CONFIGS, + "structuredMemoryConfigs": API_STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttlConfig": TTL_CONFIG, + "disableMemoryRevisions": False, + } + }, + } + with mock.patch.object( + async_memory_banks_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(PENDING_OP)), + genai_types.HttpResponse(body=json.dumps(FINISHED_OP)), + genai_types.HttpResponse(body=json.dumps(memory_bank_response)), + ] + + memory_bank = await async_memory_banks_client.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": False, + }, + config={ + "display_name": DISPLAY_NAME, + "description": DESCRIPTION, + }, + ) + + request_mock.assert_has_calls( + [ + mock.call( + "post", + "reasoningEngines", + reasoning_engine, + None, + ), + mock.call( + "get", + OPERATION_NAME, + {"_url": {"operationName": OPERATION_NAME}}, + None, + ), + ] + ) + assert isinstance(memory_bank, genai.types.MemoryBank) + assert memory_bank.name == MEMORY_BANK_NAME + + def test_get_memory_bank(self, memory_banks_client): + """Tests the retrieval of a Memory Bank.""" + reasoning_engine = { + "name": MEMORY_BANK_NAME, + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + "contextSpec": { + "memoryBankConfig": { + "generationConfig": GENERATION_CONFIG, + "similaritySearchConfig": SIMILARITY_SEARCH_CONFIG, + "customizationConfigs": UNSTRUCTURED_MEMORY_CONFIGS, + "structuredMemoryConfigs": API_STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttlConfig": TTL_CONFIG, + "disableMemoryRevisions": False, + } + }, + } + with mock.patch.object( + memory_banks_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(reasoning_engine)), + ] + memory_bank = memory_banks_client.get(name=MEMORY_BANK_NAME) + request_mock.assert_called_once_with( + "get", + MEMORY_BANK_NAME, + {"_url": {"name": MEMORY_BANK_NAME}}, + None, + ) + assert isinstance(memory_bank, genai.types.MemoryBank) + 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 + == agentplatform_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=False, + ) + ) + + @pytest.mark.asyncio + async def test_async_get_memory_bank(self, async_memory_banks_client): + """Tests the retrieval of a Memory Bank.""" + reasoning_engine = { + "name": MEMORY_BANK_NAME, + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + "contextSpec": { + "memoryBankConfig": { + "generationConfig": GENERATION_CONFIG, + "similaritySearchConfig": SIMILARITY_SEARCH_CONFIG, + "customizationConfigs": UNSTRUCTURED_MEMORY_CONFIGS, + "structuredMemoryConfigs": API_STRUCTURED_MEMORY_SCHEMA_CONFIGS, + "ttlConfig": TTL_CONFIG, + "disableMemoryRevisions": False, + } + }, + } + with mock.patch.object( + async_memory_banks_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(reasoning_engine)), + ] + memory_bank = await async_memory_banks_client.get(name=MEMORY_BANK_NAME) + request_mock.assert_called_once_with( + "get", + MEMORY_BANK_NAME, + {"_url": {"name": MEMORY_BANK_NAME}}, + None, + ) + assert isinstance(memory_bank, genai.types.MemoryBank) + 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 + == agentplatform_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=False, + ) + ) + + def test_list_memory_banks(self, memory_banks_client): + """Tests the listing of Memory Banks.""" + reasoning_engine_1 = { + "name": MEMORY_BANK_NAME, + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + } + reasoning_engine_2 = { + "name": "projects/test-project/locations/test-location/reasoningEngines/456", # pylint: disable=line-too-long + "displayName": "My Second Test Memory Bank", + "description": "My Second Test Memory Bank Description", + } + response = {"reasoningEngines": [reasoning_engine_1, reasoning_engine_2]} + with mock.patch.object( + memory_banks_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(response)), + ] + memory_banks = list(memory_banks_client.list()) + request_mock.assert_called_once_with( + "get", + "reasoningEngines", + {}, + None, + ) + assert len(memory_banks) == 2 + assert isinstance(memory_banks[0], genai.types.MemoryBank) + assert memory_banks[0].name == MEMORY_BANK_NAME + assert memory_banks[0].display_name == DISPLAY_NAME + assert memory_banks[0].description == DESCRIPTION + assert isinstance(memory_banks[1], genai.types.MemoryBank) + assert ( + memory_banks[1].name + == "projects/test-project/locations/test-location/reasoningEngines/456" + ) + assert memory_banks[1].display_name == "My Second Test Memory Bank" + assert ( + memory_banks[1].description == "My Second Test Memory Bank Description" + ) + + @pytest.mark.asyncio + async def test_async_list_memory_banks(self, async_memory_banks_client): + """Tests the listing of Memory Banks.""" + reasoning_engine_1 = { + "name": MEMORY_BANK_NAME, + "displayName": DISPLAY_NAME, + "description": DESCRIPTION, + } + reasoning_engine_2 = { + "name": "projects/test-project/locations/test-location/reasoningEngines/456", # pylint: disable=line-too-long + "displayName": "My Second Test Memory Bank", + "description": "My Second Test Memory Bank Description", + } + response = {"reasoningEngines": [reasoning_engine_1, reasoning_engine_2]} + with mock.patch.object( + async_memory_banks_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(response)), + ] + memory_banks = [] + async for memory_bank in await async_memory_banks_client.list(): + memory_banks.append(memory_bank) + request_mock.assert_called_once_with( + "get", + "reasoningEngines", + {}, + None, + ) + assert len(memory_banks) == 2 + assert isinstance(memory_banks[0], genai.types.MemoryBank) + assert memory_banks[0].name == MEMORY_BANK_NAME + assert memory_banks[0].display_name == DISPLAY_NAME + assert memory_banks[0].description == DESCRIPTION + assert isinstance(memory_banks[1], genai.types.MemoryBank) + assert ( + memory_banks[1].name + == "projects/test-project/locations/test-location/reasoningEngines/456" + ) + assert memory_banks[1].display_name == "My Second Test Memory Bank" + assert ( + memory_banks[1].description == "My Second Test Memory Bank Description" + ) + + def test_delete_memory_bank(self, memory_banks_client): + """Tests the deletion of a Memory Bank.""" + with mock.patch.object( + memory_banks_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=None), + ] + memory_banks_client.delete(name=MEMORY_BANK_NAME, force=True) + request_mock.assert_called_once_with( + "delete", + MEMORY_BANK_NAME, + {"_url": {"name": MEMORY_BANK_NAME}, "force": True}, + None, + ) + + @pytest.mark.asyncio + async def test_async_delete_memory_bank(self, async_memory_banks_client): + """Tests the deletion of a Memory Bank.""" + with mock.patch.object( + async_memory_banks_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=None), + ] + await async_memory_banks_client.delete(name=MEMORY_BANK_NAME, force=True) + request_mock.assert_called_once_with( + "delete", + MEMORY_BANK_NAME, + {"_url": {"name": MEMORY_BANK_NAME}, "force": True}, + None, + )