From d7e752ace644b8340ca8f4b66adace30c6c7173c Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 3 Aug 2026 17:45:53 -0700 Subject: [PATCH] feat: Support configs. PiperOrigin-RevId: 958671891 --- agentplatform/_genai/memory_banks.py | 133 ++++++++++++++++++++++++- agentplatform/_genai/types/__init__.py | 6 ++ agentplatform/_genai/types/common.py | 46 +++++++++ 3 files changed, 181 insertions(+), 4 deletions(-) diff --git a/agentplatform/_genai/memory_banks.py b/agentplatform/_genai/memory_banks.py index c0b0cee277..bfce514a93 100644 --- a/agentplatform/_genai/memory_banks.py +++ b/agentplatform/_genai/memory_banks.py @@ -45,6 +45,15 @@ def _CreateMemoryBankRequestParameters_to_vertex( ) -> dict[str, Any]: to_object: dict[str, Any] = {} + if getv(from_object, ["memory_bank_config"]) is not None: + setv( + parent_object, + ["context_spec", "memoryBankConfig"], + _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + getv(from_object, ["memory_bank_config"]), to_object + ), + ) + return to_object @@ -155,10 +164,98 @@ def _IngestEventsRequestParameters_to_vertex( 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 _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_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 +263,7 @@ def _create( parameter_model = types._CreateMemoryBankRequestParameters( config=config, + memory_bank_config=memory_bank_config, ) request_url_dict: Optional[dict[str, str]] @@ -459,10 +557,31 @@ 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.""" + import json - operation = self._create() + if managed_semantic_memory_config: + memory_bank_config = json.loads( + managed_semantic_memory_config.model_dump_json() + ) + else: + memory_bank_config = {} + if "unstructured_memory_configs" in memory_bank_config: + memory_bank_config["customization_configs"] = memory_bank_config[ + "unstructured_memory_configs" + ] + + operation = self._create( + memory_bank_config=memory_bank_config, + ) operation = _memory_bank_utils._await_operation( operation_name=operation.name, @@ -587,7 +706,12 @@ def ingest_events( 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. @@ -595,6 +719,7 @@ async def _create( parameter_model = types._CreateMemoryBankRequestParameters( config=config, + memory_bank_config=memory_bank_config, ) request_url_dict: Optional[dict[str, str]] diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 9e9278030a..7bd00e288f 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -1038,6 +1038,9 @@ from .common import MachineSpec from .common import MachineSpecDict from .common import MachineSpecOrDict +from .common import ManagedSemanticMemoryConfig +from .common import ManagedSemanticMemoryConfigDict +from .common import ManagedSemanticMemoryConfigOrDict from .common import ManagedTopicEnum from .common import MapInstance from .common import MapInstanceDict @@ -3740,6 +3743,9 @@ "DeployOption", "DeployOptionDict", "DeployOptionOrDict", + "ManagedSemanticMemoryConfig", + "ManagedSemanticMemoryConfigDict", + "ManagedSemanticMemoryConfigOrDict", "A2aTaskState", "State", "Strategy", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index b0343429c1..0db8eb9a89 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -10206,6 +10206,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): @@ -10214,6 +10217,9 @@ class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): config: Optional[CreateMemoryBankConfigDict] """""" + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + """""" + _CreateMemoryBankRequestParametersOrDict = Union[ _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict @@ -27671,3 +27677,43 @@ class DeployOptionDict(TypedDict, total=False): DeployOptionOrDict = Union[DeployOption, DeployOptionDict] + + +class ManagedSemanticMemoryConfig(_common.BaseModel): + """The configuration for managed semantic memory.""" + + generation_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfig + ] = Field(default=None, description="""""") + similarity_search_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig + ] = Field(default=None, description="""""") + ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfig] = Field( + default=None, description="""""" + ) + disable_memory_revisions: Optional[bool] = Field(default=None, description="""""") + + +class ManagedSemanticMemoryConfigDict(TypedDict, total=False): + """The configuration for managed semantic memory.""" + + generation_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict + ] + """""" + + similarity_search_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict + ] + """""" + + ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict] + """""" + + disable_memory_revisions: Optional[bool] + """""" + + +ManagedSemanticMemoryConfigOrDict = Union[ + ManagedSemanticMemoryConfig, ManagedSemanticMemoryConfigDict +]