diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index d6b1820371..6cb42423c4 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -441,11 +441,7 @@ class _RequirementsValidationResult(TypedDict): actions: _RequirementsValidationActions -AgentEngineOperationUnion = Union[ - genai_types.AgentEngineOperation, - genai_types.AgentEngineMemoryOperation, - genai_types.AgentEngineGenerateMemoriesOperation, -] +AgentEngineOperationUnion = Union[genai_types.AgentEngineOperation] class GetOperationFunction(Protocol): diff --git a/agentplatform/_genai/_memory_bank_utils.py b/agentplatform/_genai/_memory_bank_utils.py new file mode 100644 index 0000000000..694fa4e3cf --- /dev/null +++ b/agentplatform/_genai/_memory_bank_utils.py @@ -0,0 +1,142 @@ +# 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. +# +"""Utility functions for memory banks.""" + +import asyncio +import re +import time +from typing import ( + Any, + Protocol, + Union, +) + +from . import types as genai_types + + +MemoryBankOperation = Union[ + genai_types.MemoryBankOperation, + genai_types.MemoryOperation, + genai_types.GenerateMemoriesOperation, +] + + +class GetOperationFunction(Protocol): + def __call__(self, *, operation_name: str, **kwargs: Any) -> MemoryBankOperation: + pass + + +class GetAsyncOperationFunction(Protocol): + async def __call__( + self, *, operation_name: str, **kwargs: Any + ) -> MemoryBankOperation: + pass + + +def _get_memory_bank_id(operation_name: str = "", resource_name: str = "") -> str: + """Returns Memory Bank ID from operation name or resource name.""" + if not resource_name and not operation_name: + raise ValueError("Resource name or operation name cannot be empty.") + + if resource_name: + match = re.match( + r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)$", + resource_name, + ) + if match: + return match.group(1) + match = re.match( + r"^projects/[^/]+/locations/[^/]+/memoryBanks/([^/]+)$", + resource_name, + ) + if match: + return match.group(1) + raise ValueError( + "Failed to parse Memory Bank ID from resource name: " f"`{resource_name}`" + ) + + if not operation_name: + raise ValueError("Operation name cannot be empty.") + + match = re.match( + r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)/operations/[^/]+$", + operation_name, + ) + if match: + return match.group(1) + + match = re.match( + r"^projects/[^/]+/locations/[^/]+/memoryBanks/([^/]+)/operations/[^/]+$", + operation_name, + ) + raise ValueError( + "Failed to parse Memory Bank ID from operation name: " f"`{operation_name}`" + ) + + +def _await_operation( + *, + operation_name: str, + get_operation_fn: GetOperationFunction, + poll_interval_seconds: float = 1, +) -> Any: + """Waits for the operation to complete. + + Args: + operation_name (str): + Required. The name of the operation. + poll_interval_seconds (float): + The number of seconds to wait between each poll. + get_operation_fn (Callable[[str], Any]): + Optional. The function to use for getting the operation. If not + provided, `self._get_memory_bank_operation` will be used. + + Returns: + The operation that has completed (i.e. `operation.done==True`). + """ + operation = get_operation_fn(operation_name=operation_name) + while not operation.done: + time.sleep(poll_interval_seconds) + operation = get_operation_fn(operation_name=operation.name) + + return operation + + +async def _await_async_operation( + *, + operation_name: str, + get_operation_fn: GetAsyncOperationFunction, + poll_interval_seconds: float = 1, +) -> Any: + """Waits for the operation to complete. + + Args: + operation_name (str): + Required. The name of the operation. + poll_interval_seconds (float): + The number of seconds to wait between each poll. + get_operation_fn (Callable[[str], Awaitable[Any]]): + Optional. The async function to use for getting the operation. If not + provided, `self._get_memory_bank_operation` will be used. + + Returns: + The operation that has completed (i.e. `operation.done==True`). + """ + operation = await get_operation_fn(operation_name=operation_name) + while not operation.done: + await asyncio.sleep(poll_interval_seconds) + operation = await get_operation_fn(operation_name=operation.name) + + return operation diff --git a/agentplatform/_genai/agent_engines.py b/agentplatform/_genai/agent_engines.py index ef4d6f2a09..2fd26fda80 100644 --- a/agentplatform/_genai/agent_engines.py +++ b/agentplatform/_genai/agent_engines.py @@ -37,12 +37,10 @@ if typing.TYPE_CHECKING: from . import sessions as sessions_module - from . import memories as memories_module from . import a2a_tasks as a2a_tasks_module from . import runtimes as runtimes_module _ = sessions_module - __ = memories_module ___ = a2a_tasks_module ____ = runtimes_module @@ -1513,7 +1511,6 @@ def _update( return return_value _a2a_tasks = None - _memories = None _sandboxes = None _sessions = None _runtimes = None @@ -1548,21 +1545,6 @@ def a2a_tasks(self) -> "a2a_tasks_module.A2aTasks": ) from e return self._a2a_tasks.A2aTasks(self._api_client) # type: ignore[no-any-return] - @property - def memories(self) -> "memories_module.Memories": - if self._memories is None: - try: - # We need to lazy load the memories module to handle the - # possibility of ImportError when dependencies are not installed. - self._memories = importlib.import_module(".memories", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._memories.Memories(self._api_client) # type: ignore[no-any-return] - @property def sandboxes(self) -> Any: if self._sandboxes is None: @@ -2899,145 +2881,6 @@ async def _async_stream_query( async for response in async_iterator: yield response - def create_memory( - self, - *, - name: str, - fact: str, - scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.create instead.""" - warnings.warn( - ( - "agent_engines.create_memory is deprecated. " - "Use agent_engines.memories.create instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.create( - name=name, - fact=fact, - scope=scope, - config=config, - ) - - def delete_memory( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_memory is deprecated. " - "Use agent_engines.memories.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.delete(name=name, config=config) - - def generate_memories( - self, - *, - name: str, - vertex_session_source: Optional[ - types.GenerateMemoriesRequestVertexSessionSourceOrDict - ] = None, - direct_contents_source: Optional[ - types.GenerateMemoriesRequestDirectContentsSourceOrDict - ] = None, - direct_memories_source: Optional[ - types.GenerateMemoriesRequestDirectMemoriesSourceOrDict - ] = None, - scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Deprecated. Use agent_engines.memories.generate instead.""" - warnings.warn( - ( - "agent_engines.generate_memories is deprecated. " - "Use agent_engines.memories.generate instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.generate( - name=name, - vertex_session_source=vertex_session_source, - direct_contents_source=direct_contents_source, - direct_memories_source=direct_memories_source, - scope=scope, - config=config, - ) - - def get_memory( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, - ) -> types.Memory: - """Deprecated. Use agent_engines.memories.get instead.""" - warnings.warn( - ( - "agent_engines.get_memory is deprecated. " - "Use agent_engines.memories.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.get(name=name, config=config) - - def list_memories( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, - ) -> Iterator[types.Memory]: - """Deprecated. Use agent_engines.memories.list instead.""" - warnings.warn( - ( - "agent_engines.list_memories is deprecated. " - "Use agent_engines.memories.list instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.list(name=name, config=config) - - def retrieve_memories( - self, - *, - name: str, - scope: dict[str, str], - similarity_search_params: Optional[ - types.RetrieveMemoriesRequestSimilaritySearchParamsOrDict - ] = None, - simple_retrieval_params: Optional[ - types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict - ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, - ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: - """Deprecated. Use agent_engines.memories.retrieve instead.""" - warnings.warn( - ( - "agent_engines.retrieve_memories is deprecated. " - "Use agent_engines.memories.retrieve instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.retrieve( - name=name, - scope=scope, - similarity_search_params=similarity_search_params, - simple_retrieval_params=simple_retrieval_params, - config=config, - ) - def create_session( self, *, @@ -3928,7 +3771,6 @@ async def _update( return return_value _a2a_tasks = None - _memories = None _sessions = None _runtimes = None @@ -3990,21 +3832,6 @@ def a2a_tasks(self) -> "a2a_tasks_module.AsyncA2aTasks": ) from e return self._a2a_tasks.AsyncA2aTasks(self._api_client) # type: ignore[no-any-return] - @property - def memories(self) -> "memories_module.AsyncMemories": - if self._memories is None: - try: - # We need to lazy load the memories module to handle the - # possibility of ImportError when dependencies are not installed. - self._memories = importlib.import_module(".memories", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._memories.AsyncMemories(self._api_client) # type: ignore[no-any-return] - @property def sessions(self) -> "sessions_module.AsyncSessions": if self._sessions is None: @@ -4046,23 +3873,6 @@ async def append_session_event( config=config, ) - async def delete_memory( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_memory is deprecated. " - "Use agent_engines.memories.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.memories.delete(name=name, config=config) - async def delete_session( self, *, @@ -4080,23 +3890,6 @@ async def delete_session( ) return await self.sessions.delete(name=name, config=config) - async def get_memory( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, - ) -> types.Memory: - """Deprecated. Use agent_engines.memories.get instead.""" - warnings.warn( - ( - "agent_engines.get_memory is deprecated. " - "Use agent_engines.memories.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.memories.get(name=name, config=config) - async def get_session( self, *, diff --git a/agentplatform/_genai/client.py b/agentplatform/_genai/client.py index 0dcb0d445e..e93e977980 100644 --- a/agentplatform/_genai/client.py +++ b/agentplatform/_genai/client.py @@ -47,6 +47,9 @@ from agentplatform._genai import ( feedback_entries as feedback_entries_module, ) + from agentplatform._genai import ( + memory_banks as memory_banks_module, + ) _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules" @@ -92,6 +95,7 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._rag: Optional[ModuleType] = None self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None + self._memory_banks: Optional[ModuleType] = None @property @_common.experimental_warning( @@ -207,6 +211,12 @@ def model_garden(self) -> "model_garden_module.AsyncModelGarden": ) return self._model_garden.AsyncModelGarden(self._api_client) # type: ignore[no-any-return] + @property + def memory_banks(self) -> "memory_banks_module.AsyncMemoryBanks": + if self._memory_banks is None: + self._memory_banks = importlib.import_module(".memory_banks", __package__) + return self._memory_banks.AsyncMemoryBanks(self._api_client) # type: ignore[no-any-return] + async def aclose(self) -> None: """Closes the async client explicitly. @@ -314,6 +324,7 @@ def __init__( self._rag: Optional[ModuleType] = None self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None + self._memory_banks: Optional[ModuleType] = None @property def evals(self) -> "evals_module.Evals": @@ -453,3 +464,9 @@ def model_garden(self) -> "model_garden_module.ModelGarden": __package__, ) return self._model_garden.ModelGarden(self._api_client) # type: ignore[no-any-return] + + @property + def memory_banks(self) -> "memory_banks_module.MemoryBanks": + if self._memory_banks is None: + self._memory_banks = importlib.import_module(".memory_banks", __package__) + return self._memory_banks.MemoryBanks(self._api_client) # type: ignore[no-any-return] diff --git a/agentplatform/_genai/memories.py b/agentplatform/_genai/memories.py index c83cd240e1..6c660c14d2 100644 --- a/agentplatform/_genai/memories.py +++ b/agentplatform/_genai/memories.py @@ -30,7 +30,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import AsyncPager, Pager -from . import _agent_engines_utils +from . import _memory_bank_utils from . import types if typing.TYPE_CHECKING: @@ -41,63 +41,8 @@ logger = logging.getLogger("agentplatform_genai.memories") -logger.setLevel(logging.INFO) - -def _AgentEngineMemoryConfig_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, ["ttl"]) is not None: - setv(parent_object, ["ttl"], getv(from_object, ["ttl"])) - - if getv(from_object, ["expire_time"]) is not None: - setv(parent_object, ["expireTime"], getv(from_object, ["expire_time"])) - - if getv(from_object, ["revision_expire_time"]) is not None: - setv( - parent_object, - ["revisionExpireTime"], - getv(from_object, ["revision_expire_time"]), - ) - - if getv(from_object, ["revision_ttl"]) is not None: - setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) - - if getv(from_object, ["disable_memory_revisions"]) is not None: - setv( - parent_object, - ["disableMemoryRevisions"], - getv(from_object, ["disable_memory_revisions"]), - ) - - if getv(from_object, ["topics"]) is not None: - setv( - parent_object, ["topics"], [item for item in getv(from_object, ["topics"])] - ) - - if getv(from_object, ["metadata"]) is not None: - setv( - parent_object, - ["metadata"], - {k: v for k, v in getv(from_object, ["metadata"]).items()}, - ) - - if getv(from_object, ["memory_id"]) is not None: - setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"])) - - return to_object - - -def _CreateAgentEngineMemoryRequestParameters_to_vertex( +def _CreateMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -112,12 +57,12 @@ def _CreateAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _AgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + _MemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineMemoryRequestParameters_to_vertex( +def _DeleteMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -128,7 +73,7 @@ def _DeleteAgentEngineMemoryRequestParameters_to_vertex( return to_object -def _GenerateAgentEngineMemoriesConfig_to_vertex( +def _GenerateMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -185,7 +130,7 @@ def _GenerateAgentEngineMemoriesConfig_to_vertex( return to_object -def _GenerateAgentEngineMemoriesRequestParameters_to_vertex( +def _GenerateMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -218,14 +163,12 @@ def _GenerateAgentEngineMemoriesRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _GenerateAgentEngineMemoriesConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _GenerateMemoriesConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( +def _GetGenerateMemoriesOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -238,7 +181,7 @@ def _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( return to_object -def _GetAgentEngineMemoryOperationParameters_to_vertex( +def _GetMemoryOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -251,7 +194,7 @@ def _GetAgentEngineMemoryOperationParameters_to_vertex( return to_object -def _GetAgentEngineMemoryRequestParameters_to_vertex( +def _GetMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -262,53 +205,28 @@ def _GetAgentEngineMemoryRequestParameters_to_vertex( return to_object -def _IngestEventsConfig_to_vertex( +def _ListMemoriesConfig_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, ["force_flush"]) is not None: - setv(parent_object, ["forceFlush"], getv(from_object, ["force_flush"])) - - if getv(from_object, ["revision_labels"]) is not None: - setv(parent_object, ["revisionLabels"], getv(from_object, ["revision_labels"])) - - if getv(from_object, ["revision_expire_time"]) is not None: - setv( - parent_object, - ["revisionExpireTime"], - getv(from_object, ["revision_expire_time"]), - ) - - if getv(from_object, ["revision_ttl"]) is not None: - setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) + if getv(from_object, ["page_size"]) is not None: + setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) - if getv(from_object, ["disable_memory_revisions"]) is not None: - setv( - parent_object, - ["disableMemoryRevisions"], - getv(from_object, ["disable_memory_revisions"]), - ) + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) - if getv(from_object, ["metadata"]) is not None: - setv( - parent_object, - ["metadata"], - {k: v for k, v in getv(from_object, ["metadata"]).items()}, - ) + if getv(from_object, ["filter"]) is not None: + setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) - if getv(from_object, ["metadata_merge_strategy"]) is not None: - setv( - parent_object, - ["metadataMergeStrategy"], - getv(from_object, ["metadata_merge_strategy"]), - ) + if getv(from_object, ["order_by"]) is not None: + setv(parent_object, ["_query", "orderBy"], getv(from_object, ["order_by"])) return to_object -def _IngestEventsRequestParameters_to_vertex( +def _ListMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -316,68 +234,66 @@ def _IngestEventsRequestParameters_to_vertex( if getv(from_object, ["name"]) is not None: setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - if getv(from_object, ["stream_id"]) is not None: - setv(to_object, ["streamId"], getv(from_object, ["stream_id"])) - - if getv(from_object, ["direct_contents_source"]) is not None: - setv( - to_object, - ["directContentsSource"], - getv(from_object, ["direct_contents_source"]), - ) - - if getv(from_object, ["scope"]) is not None: - setv(to_object, ["scope"], getv(from_object, ["scope"])) - - if getv(from_object, ["generation_trigger_config"]) is not None: - setv( - to_object, - ["generationTriggerConfig"], - getv(from_object, ["generation_trigger_config"]), - ) - if getv(from_object, ["config"]) is not None: - _IngestEventsConfig_to_vertex(getv(from_object, ["config"]), to_object) + _ListMemoriesConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _ListAgentEngineMemoryConfig_to_vertex( +def _MemoryConfig_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, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) - if getv(from_object, ["page_token"]) is not None: - setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) - if getv(from_object, ["filter"]) is not None: - setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) + if getv(from_object, ["ttl"]) is not None: + setv(parent_object, ["ttl"], getv(from_object, ["ttl"])) - if getv(from_object, ["order_by"]) is not None: - setv(parent_object, ["_query", "orderBy"], getv(from_object, ["order_by"])) + if getv(from_object, ["expire_time"]) is not None: + setv(parent_object, ["expireTime"], getv(from_object, ["expire_time"])) - return to_object + if getv(from_object, ["revision_expire_time"]) is not None: + setv( + parent_object, + ["revisionExpireTime"], + getv(from_object, ["revision_expire_time"]), + ) + if getv(from_object, ["revision_ttl"]) is not None: + setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) -def _ListAgentEngineMemoryRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + if getv(from_object, ["disable_memory_revisions"]) is not None: + setv( + parent_object, + ["disableMemoryRevisions"], + getv(from_object, ["disable_memory_revisions"]), + ) - if getv(from_object, ["config"]) is not None: - _ListAgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + if getv(from_object, ["topics"]) is not None: + setv( + parent_object, ["topics"], [item for item in getv(from_object, ["topics"])] + ) + + if getv(from_object, ["metadata"]) is not None: + setv( + parent_object, + ["metadata"], + {k: v for k, v in getv(from_object, ["metadata"]).items()}, + ) + + if getv(from_object, ["memory_id"]) is not None: + setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"])) return to_object -def _PurgeAgentEngineMemoriesRequestParameters_to_vertex( +def _PurgeMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -401,7 +317,7 @@ def _PurgeAgentEngineMemoriesRequestParameters_to_vertex( return to_object -def _RetrieveAgentEngineMemoriesConfig_to_vertex( +def _RetrieveMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -423,7 +339,7 @@ def _RetrieveAgentEngineMemoriesConfig_to_vertex( return to_object -def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( +def _RetrieveMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -449,9 +365,7 @@ def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( ) if getv(from_object, ["config"]) is not None: - _RetrieveAgentEngineMemoriesConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _RetrieveMemoriesConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -470,7 +384,7 @@ def _RetrieveMemoryProfilesRequestParameters_to_vertex( return to_object -def _RollbackAgentEngineMemoryRequestParameters_to_vertex( +def _RollbackMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -484,7 +398,7 @@ def _RollbackAgentEngineMemoryRequestParameters_to_vertex( return to_object -def _UpdateAgentEngineMemoryConfig_to_vertex( +def _UpdateMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -542,7 +456,7 @@ def _UpdateAgentEngineMemoryConfig_to_vertex( return to_object -def _UpdateAgentEngineMemoryRequestParameters_to_vertex( +def _UpdateMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -557,9 +471,7 @@ def _UpdateAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _UpdateAgentEngineMemoryConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _UpdateMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -572,13 +484,13 @@ def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.MemoryConfigOrDict] = None, + ) -> types.MemoryOperation: """ - Creates a new memory in the Agent Engine. + Creates a new memory in the Memory Bank. """ - parameter_model = types._CreateAgentEngineMemoryRequestParameters( + parameter_model = types._CreateMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -591,9 +503,7 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _CreateMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -620,7 +530,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.MemoryOperation._from_response( response=response_dict, kwargs=( { @@ -645,24 +555,21 @@ def _create( return return_value def delete( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: + self, *, name: str, config: Optional[types.DeleteMemoryConfigOrDict] = None + ) -> types.DeleteMemoryOperation: """ - Delete an Agent Engine memory. + Delete a memory. Args: name (str): - Required. The name of the Agent Engine memory to be deleted. Format: + Required. The name of the memory to delete. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteAgentEngineMemoryConfig): - Optional. Additional configurations for deleting the Agent Engine. + config (DeleteRuntimeMemoryConfig): + Optional. Additional configurations for deleting the Agent Runtime. """ - parameter_model = types._DeleteAgentEngineMemoryRequestParameters( + parameter_model = types._DeleteMemoryRequestParameters( name=name, config=config, ) @@ -673,9 +580,7 @@ def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _DeleteMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -702,7 +607,7 @@ def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineMemoryOperation._from_response( + return_value = types.DeleteMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -740,13 +645,13 @@ def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: + config: Optional[types.GenerateMemoriesConfigOrDict] = None, + ) -> types.GenerateMemoriesOperation: """ - Generates memories for an Agent Engine. + Generates memories for a Memory Bank. """ - parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( + parameter_model = types._GenerateMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -761,9 +666,7 @@ def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GenerateMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:generate".format_map(request_url_dict) @@ -790,7 +693,7 @@ def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.GenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -815,13 +718,10 @@ def _generate( return return_value def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetMemoryConfigOrDict] = None ) -> types.Memory: """ - Gets an agent engine memory. + Gets a memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -830,7 +730,7 @@ def get( """ - parameter_model = types._GetAgentEngineMemoryRequestParameters( + parameter_model = types._GetMemoryRequestParameters( name=name, config=config, ) @@ -841,9 +741,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -894,101 +792,14 @@ def get( 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._IngestEventsRequestParameters( - name=name, - stream_id=stream_id, - direct_contents_source=direct_contents_source, - scope=scope, - generation_trigger_config=generation_trigger_config, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _IngestEventsRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}/memories:ingestEvents".format_map(request_url_dict) - else: - path = "{name}/memories:ingestEvents" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.MemoryBankIngestEventsOperation._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 _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, - ) -> types.ListReasoningEnginesMemoriesResponse: + self, *, name: str, config: Optional[types.ListMemoriesConfigOrDict] = None + ) -> types.ListMemoriesResponse: """ - Lists Agent Engine memories. + Lists memories. """ - parameter_model = types._ListAgentEngineMemoryRequestParameters( + parameter_model = types._ListMemoriesRequestParameters( name=name, config=config, ) @@ -999,9 +810,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _ListMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -1028,7 +837,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListReasoningEnginesMemoriesResponse._from_response( + return_value = types.ListMemoriesResponse._from_response( response=response_dict, kwargs=( { @@ -1056,9 +865,9 @@ def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - parameter_model = types._GetAgentEngineMemoryOperationParameters( + config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, + ) -> types.MemoryOperation: + parameter_model = types._GetMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -1069,9 +878,7 @@ def _get_memory_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( - parameter_model - ) + request_dict = _GetMemoryOperationParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{operationName}".format_map(request_url_dict) @@ -1098,7 +905,7 @@ def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.MemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1126,9 +933,9 @@ def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( + config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, + ) -> types.GenerateMemoriesOperation: + parameter_model = types._GetGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -1139,7 +946,7 @@ def _get_generate_memories_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1168,7 +975,7 @@ def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.GenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1203,13 +1010,13 @@ def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories for an Agent Engine. + Retrieves memories from a Memory Bank. """ - parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( + parameter_model = types._RetrieveMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -1223,9 +1030,7 @@ def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( - parameter_model - ) + request_dict = _RetrieveMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:retrieve".format_map(request_url_dict) @@ -1284,13 +1089,13 @@ def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for an Agent Engine. + Retrieves memory profiles for a Memory Bank. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.agent_engines.memories.retrieve_profiles( + result = client.memory_banks.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -1383,13 +1188,13 @@ def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackMemoryConfigOrDict] = None, + ) -> types.RollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackAgentEngineMemoryRequestParameters( + parameter_model = types._RollbackMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -1401,9 +1206,7 @@ def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _RollbackMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}:rollback".format_map(request_url_dict) @@ -1430,7 +1233,7 @@ def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineRollbackMemoryOperation._from_response( + return_value = types.RollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1460,13 +1263,13 @@ def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.UpdateMemoryConfigOrDict] = None, + ) -> types.MemoryOperation: """ - Updates an Agent Engine memory. + Updates a memory. """ - parameter_model = types._UpdateAgentEngineMemoryRequestParameters( + parameter_model = types._UpdateMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -1479,9 +1282,7 @@ def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _UpdateMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -1508,7 +1309,7 @@ def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.MemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1541,13 +1342,13 @@ def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: + config: Optional[types.PurgeMemoriesConfigOrDict] = None, + ) -> types.PurgeMemoriesOperation: """ - Purges memories from an Agent Engine. + Purges memories from a Memory Bank. """ - parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( + parameter_model = types._PurgeMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -1561,9 +1362,7 @@ def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( - parameter_model - ) + request_dict = _PurgeMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:purge".format_map(request_url_dict) @@ -1590,7 +1389,7 @@ def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEnginePurgeMemoriesOperation._from_response( + return_value = types.PurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1619,18 +1418,7 @@ def _purge( @property def revisions(self) -> "memory_revisions_module.MemoryRevisions": if self._revisions is None: - try: - # We need to lazy load the revisions module to handle the - # possibility of ImportError when dependencies are not installed. - self._revisions = importlib.import_module( - ".memory_revisions", __package__ - ) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories.revisions' module requires " - "additional packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e + self._revisions = importlib.import_module(".memory_revisions", __package__) return self._revisions.MemoryRevisions(self._api_client) # type: ignore[no-any-return] def create( @@ -1639,9 +1427,9 @@ def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Creates a new memory in the Agent Engine. + config: Optional[types.MemoryConfigOrDict] = None, + ) -> types.MemoryOperation: + """Creates a new memory in Memory Bank. Args: name (str): @@ -1650,16 +1438,16 @@ def create( Required. The fact to be stored in the memory. scope (dict[str, str]): Required. The scope of the memory. For example, {"user_id": "123"}. - config (AgentEngineMemoryConfigOrDict): + config (MemoryConfigOrDict): Optional. The configuration for the memory. Returns: - AgentEngineMemoryOperation: The operation for creating the memory. + MemoryOperation: The operation for creating the memory. """ if config is None: - config = types.AgentEngineMemoryConfig() + config = types.MemoryConfig() elif isinstance(config, dict): - config = types.AgentEngineMemoryConfig.model_validate(config) + config = types.MemoryConfig.model_validate(config) operation = self._create( name=name, fact=fact, @@ -1668,7 +1456,7 @@ def create( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _memory_bank_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1697,13 +1485,13 @@ def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Generates memories for the agent engine. + config: Optional[types.GenerateMemoriesConfigOrDict] = None, + ) -> types.GenerateMemoriesOperation: + """Generates memories with Memory Bank. Args: name (str): - Required. The name of the agent engine to generate memories for. + Required. The name of the memory bank to generate memories with. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -1724,13 +1512,13 @@ def generate( Optional. The configuration for the memories to generate. Returns: - AgentEngineGenerateMemoriesOperation: + GenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateAgentEngineMemoriesConfig() + config = types.GenerateMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) + config = types.GenerateMemoriesConfig.model_validate(config) operation = self._generate( name=name, vertex_session_source=vertex_session_source, @@ -1740,7 +1528,7 @@ def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _memory_bank_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -1753,14 +1541,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + config: Optional[types.ListMemoriesConfigOrDict] = None, ) -> Iterator[types.Memory]: - """Lists Agent Engine memories. + """Lists Memory Bank memories. Args: name (str): - Required. The name of the agent engine to list memories for. - config (ListAgentEngineMemoryConfig): + Required. The name of the memory bank to list memories for. + config (ListMemoriesConfig): Optional. The configuration for the memories to list. Returns: @@ -1785,13 +1573,13 @@ def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveMemoriesConfigOrDict] = None, ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: - """Retrieves memories for the agent. + """Retrieves memories from Memory Bank. Args: name (str): - Required. The name of the agent engine to retrieve memories for. + Required. The name of the memory bank to retrieve memories from. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -1801,7 +1589,7 @@ def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveAgentEngineMemoriesConfig): + config (RetrieveMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -1832,8 +1620,8 @@ def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackMemoryConfigOrDict] = None, + ) -> types.RollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -1841,24 +1629,24 @@ def rollback( Required. The name of the memory to rollback. target_revision_id (str): Required. The revision ID to roll back to - config (RollbackAgentEngineMemoryConfig): + config (RollbackMemoryConfig): Optional. The configuration for the rollback. Returns: - AgentEngineRollbackMemoryOperation: + RollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackAgentEngineMemoryConfig() + config = types.RollbackMemoryConfig() elif isinstance(config, dict): - config = types.RollbackAgentEngineMemoryConfig.model_validate(config) + config = types.RollbackMemoryConfig.model_validate(config) operation = self._rollback( name=name, target_revision_id=target_revision_id, config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _memory_bank_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1874,13 +1662,13 @@ def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: - """Purges memories from an Agent Engine. + config: Optional[types.PurgeMemoriesConfigOrDict] = None, + ) -> types.PurgeMemoriesOperation: + """Purges memories from Memory Bank. Args: name (str): - Required. The name of the Agent Engine to purge memories from. + Required. The name of the Memory Bank to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -1890,17 +1678,17 @@ def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeAgentEngineMemoriesConfig): + config (PurgeMemoriesConfig): Optional. The configuration for the purge operation. Returns: - AgentEnginePurgeMemoriesOperation: + PurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeAgentEngineMemoriesConfig() + config = types.PurgeMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) + config = types.PurgeMemoriesConfig.model_validate(config) operation = self._purge( name=name, filter=filter, @@ -1909,7 +1697,7 @@ def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _memory_bank_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1918,89 +1706,6 @@ def purge( raise RuntimeError(f"Failed to purge memories: {operation.error}") return operation - def ingest_events( - self, - *, - name: str, - scope: dict[str, str], - stream_id: str = "", - direct_contents_source: Optional[ - types.IngestionDirectContentsSourceOrDict - ] = None, - generation_trigger_config: Optional[ - types.MemoryGenerationTriggerConfigOrDict - ] = None, - config: Optional[types.IngestEventsConfigOrDict] = None, - ) -> types.MemoryBankIngestEventsOperation: - """Ingests events into an Agent Engine. - - Example usage: - ``` - client.agent_engines.memories.ingest_events( - name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", - scope={"user_id": "test-user-id"}, - direct_contents_source={ - "events": [ - { - "content": { - "role": "user", - "parts": [ - {"text": "I am a software engineer."} - ], - } - } - ] - }, - generation_trigger_config={ - "generation_rule": { - "idle_duration": "60s" - } - } - ) - ``` - - Args: - name (str): - Required. The name of the Agent Engine to ingest events into. - scope (dict[str, str]): - Required. The scope of the events to ingest. For example, - {"user_id": "123"}. - stream_id (str): - Optional. The ID of the stream to ingest events into. If not - specified, the events will be ingested into the default stream. - direct_contents_source (IngestionDirectContentsSource): - The direct contents source, containing the events to ingest. - generation_trigger_config (MemoryGenerationTriggerConfig): - Optional. The configuration for the generation trigger config. - config (IngestEventsConfig): - Optional. The configuration for the ingest events operation. - - Returns: - AgentEngineIngestEventsOperation: - The operation for ingesting the events. - """ - if config is None: - config = types.IngestEventsConfig() - elif isinstance(config, dict): - config = types.IngestEventsConfig.model_validate(config) - operation = self._ingest_events( - name=name, - scope=scope, - stream_id=stream_id, - generation_trigger_config=generation_trigger_config, - direct_contents_source=direct_contents_source, - config=config, - ) - if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( - operation_name=operation.name, - get_operation_fn=self._get_memory_operation, - poll_interval_seconds=0.5, - ) - if operation.error: - raise RuntimeError(f"Failed to ingest events: {operation.error}") - return operation - class AsyncMemories(_api_module.BaseModule): @@ -2010,13 +1715,13 @@ async def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.MemoryConfigOrDict] = None, + ) -> types.MemoryOperation: """ - Creates a new memory in the Agent Engine. + Creates a new memory in the Memory Bank. """ - parameter_model = types._CreateAgentEngineMemoryRequestParameters( + parameter_model = types._CreateMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -2029,9 +1734,7 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _CreateMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -2060,7 +1763,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.MemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2085,24 +1788,21 @@ async def _create( return return_value async def delete( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: + self, *, name: str, config: Optional[types.DeleteMemoryConfigOrDict] = None + ) -> types.DeleteMemoryOperation: """ - Delete an Agent Engine memory. + Delete a memory. Args: name (str): - Required. The name of the Agent Engine memory to be deleted. Format: + Required. The name of the memory to delete. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteAgentEngineMemoryConfig): - Optional. Additional configurations for deleting the Agent Engine. + config (DeleteRuntimeMemoryConfig): + Optional. Additional configurations for deleting the Agent Runtime. """ - parameter_model = types._DeleteAgentEngineMemoryRequestParameters( + parameter_model = types._DeleteMemoryRequestParameters( name=name, config=config, ) @@ -2113,9 +1813,7 @@ async def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _DeleteMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -2144,7 +1842,7 @@ async def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineMemoryOperation._from_response( + return_value = types.DeleteMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2182,13 +1880,13 @@ async def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: + config: Optional[types.GenerateMemoriesConfigOrDict] = None, + ) -> types.GenerateMemoriesOperation: """ - Generates memories for an Agent Engine. + Generates memories for a Memory Bank. """ - parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( + parameter_model = types._GenerateMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -2203,9 +1901,7 @@ async def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GenerateMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:generate".format_map(request_url_dict) @@ -2234,7 +1930,7 @@ async def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.GenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2259,13 +1955,10 @@ async def _generate( return return_value async def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetMemoryConfigOrDict] = None ) -> types.Memory: """ - Gets an agent engine memory. + Gets a memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -2274,7 +1967,7 @@ async def get( """ - parameter_model = types._GetAgentEngineMemoryRequestParameters( + parameter_model = types._GetMemoryRequestParameters( name=name, config=config, ) @@ -2285,9 +1978,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -2340,103 +2031,14 @@ async def get( self._api_client._verify_response(return_value) return return_value - async 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._IngestEventsRequestParameters( - name=name, - stream_id=stream_id, - direct_contents_source=direct_contents_source, - scope=scope, - generation_trigger_config=generation_trigger_config, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _IngestEventsRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}/memories:ingestEvents".format_map(request_url_dict) - else: - path = "{name}/memories:ingestEvents" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.MemoryBankIngestEventsOperation._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 _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, - ) -> types.ListReasoningEnginesMemoriesResponse: + self, *, name: str, config: Optional[types.ListMemoriesConfigOrDict] = None + ) -> types.ListMemoriesResponse: """ - Lists Agent Engine memories. + Lists memories. """ - parameter_model = types._ListAgentEngineMemoryRequestParameters( + parameter_model = types._ListMemoriesRequestParameters( name=name, config=config, ) @@ -2447,9 +2049,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _ListMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -2478,7 +2078,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListReasoningEnginesMemoriesResponse._from_response( + return_value = types.ListMemoriesResponse._from_response( response=response_dict, kwargs=( { @@ -2506,9 +2106,9 @@ async def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - parameter_model = types._GetAgentEngineMemoryOperationParameters( + config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, + ) -> types.MemoryOperation: + parameter_model = types._GetMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -2519,9 +2119,7 @@ async def _get_memory_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( - parameter_model - ) + request_dict = _GetMemoryOperationParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{operationName}".format_map(request_url_dict) @@ -2550,7 +2148,7 @@ async def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.MemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2578,9 +2176,9 @@ async def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( + config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, + ) -> types.GenerateMemoriesOperation: + parameter_model = types._GetGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -2591,7 +2189,7 @@ async def _get_generate_memories_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2622,7 +2220,7 @@ async def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.GenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2657,13 +2255,13 @@ async def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories for an Agent Engine. + Retrieves memories from a Memory Bank. """ - parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( + parameter_model = types._RetrieveMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -2677,9 +2275,7 @@ async def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( - parameter_model - ) + request_dict = _RetrieveMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:retrieve".format_map(request_url_dict) @@ -2740,13 +2336,13 @@ async def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for an Agent Engine. + Retrieves memory profiles for a Memory Bank. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.agent_engines.memories.retrieve_profiles( + result = client.memory_banks.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -2841,13 +2437,13 @@ async def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackMemoryConfigOrDict] = None, + ) -> types.RollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackAgentEngineMemoryRequestParameters( + parameter_model = types._RollbackMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -2859,9 +2455,7 @@ async def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _RollbackMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}:rollback".format_map(request_url_dict) @@ -2890,7 +2484,7 @@ async def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineRollbackMemoryOperation._from_response( + return_value = types.RollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2920,13 +2514,13 @@ async def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.UpdateMemoryConfigOrDict] = None, + ) -> types.MemoryOperation: """ - Updates an Agent Engine memory. + Updates a memory. """ - parameter_model = types._UpdateAgentEngineMemoryRequestParameters( + parameter_model = types._UpdateMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -2939,9 +2533,7 @@ async def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _UpdateMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -2970,7 +2562,7 @@ async def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.MemoryOperation._from_response( response=response_dict, kwargs=( { @@ -3003,13 +2595,13 @@ async def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: + config: Optional[types.PurgeMemoriesConfigOrDict] = None, + ) -> types.PurgeMemoriesOperation: """ - Purges memories from an Agent Engine. + Purges memories from a Memory Bank. """ - parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( + parameter_model = types._PurgeMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -3023,9 +2615,7 @@ async def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( - parameter_model - ) + request_dict = _PurgeMemoriesRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:purge".format_map(request_url_dict) @@ -3054,7 +2644,7 @@ async def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEnginePurgeMemoriesOperation._from_response( + return_value = types.PurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -3083,18 +2673,7 @@ async def _purge( @property def revisions(self) -> "memory_revisions_module.AsyncMemoryRevisions": if self._revisions is None: - try: - # We need to lazy load the revisions module to handle the - # possibility of ImportError when dependencies are not installed. - self._revisions = importlib.import_module( - ".memory_revisions", __package__ - ) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories.revisions' module requires " - "additional packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e + self._revisions = importlib.import_module(".memory_revisions", __package__) return self._revisions.AsyncMemoryRevisions(self._api_client) # type: ignore[no-any-return] async def create( @@ -3103,9 +2682,9 @@ async def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Creates a new memory in the Agent Engine. + config: Optional[types.MemoryConfigOrDict] = None, + ) -> types.MemoryOperation: + """Creates a new memory in Memory Bank. Args: name (str): @@ -3114,16 +2693,16 @@ async def create( Required. The fact to be stored in the memory. scope (dict[str, str]): Required. The scope of the memory. For example, {"user_id": "123"}. - config (AgentEngineMemoryConfigOrDict): + config (MemoryConfigOrDict): Optional. The configuration for the memory. Returns: - AgentEngineMemoryOperation: The operation for creating the memory. + MemoryOperation: The operation for creating the memory. """ if config is None: - config = types.AgentEngineMemoryConfig() + config = types.MemoryConfig() elif isinstance(config, dict): - config = types.AgentEngineMemoryConfig.model_validate(config) + config = types.MemoryConfig.model_validate(config) operation = await self._create( name=name, fact=fact, @@ -3132,7 +2711,7 @@ async def create( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _memory_bank_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3161,13 +2740,13 @@ async def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Generates memories for the agent engine. + config: Optional[types.GenerateMemoriesConfigOrDict] = None, + ) -> types.GenerateMemoriesOperation: + """Generates memories for Memory Bank. Args: name (str): - Required. The name of the agent engine to generate memories for. + Required. The name of the Memory Bank to generate memories for. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -3188,13 +2767,13 @@ async def generate( Optional. The configuration for the memories to generate. Returns: - AgentEngineGenerateMemoriesOperation: + GenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateAgentEngineMemoriesConfig() + config = types.GenerateMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) + config = types.GenerateMemoriesConfig.model_validate(config) operation = await self._generate( name=name, vertex_session_source=vertex_session_source, @@ -3204,7 +2783,7 @@ async def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _memory_bank_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -3217,14 +2796,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + config: Optional[types.ListMemoriesConfigOrDict] = None, ) -> AsyncPager[types.Memory]: - """Lists Agent Engine memories. + """Lists memories in Memory Bank. Args: name (str): - Required. The name of the agent engine to list memories for. - config (ListAgentEngineMemoryConfig): + Required. The name of the Memory Bank to list memories for. + config (ListMemoriesConfig): Optional. The configuration for the memories to list. Returns: @@ -3249,13 +2828,13 @@ async def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveMemoriesConfigOrDict] = None, ) -> AsyncPager[types.RetrieveMemoriesResponseRetrievedMemory]: - """Retrieves memories for the agent. + """Retrieves memories from Memory Bank. Args: name (str): - Required. The name of the agent engine to retrieve memories for. + Required. The name of the Memory Bank to retrieve memories for. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -3265,7 +2844,7 @@ async def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveAgentEngineMemoriesConfig): + config (RetrieveMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -3296,8 +2875,8 @@ async def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackMemoryConfigOrDict] = None, + ) -> types.RollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -3305,24 +2884,24 @@ async def rollback( Required. The name of the memory to rollback. target_revision_id (str): Required. The revision ID to roll back to - config (RollbackAgentEngineMemoryConfig): + config (RollbackMemoryConfig): Optional. The configuration for the rollback. Returns: - AgentEngineRollbackMemoryOperation: + RollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackAgentEngineMemoryConfig() + config = types.RollbackMemoryConfig() elif isinstance(config, dict): - config = types.RollbackAgentEngineMemoryConfig.model_validate(config) + config = types.RollbackMemoryConfig.model_validate(config) operation = await self._rollback( name=name, target_revision_id=target_revision_id, config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _memory_bank_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3338,13 +2917,13 @@ async def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: - """Purges memories from an Agent Engine. + config: Optional[types.PurgeMemoriesConfigOrDict] = None, + ) -> types.PurgeMemoriesOperation: + """Purges memories from a Memory Bank. Args: name (str): - Required. The name of the Agent Engine to purge memories from. + Required. The name of the Memory Bank to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -3354,17 +2933,17 @@ async def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeAgentEngineMemoriesConfig): + config (PurgeMemoriesConfig): Optional. The configuration for the purge operation. Returns: - AgentEnginePurgeMemoriesOperation: + MemoryBankPurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeAgentEngineMemoriesConfig() + config = types.PurgeMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) + config = types.PurgeMemoriesConfig.model_validate(config) operation = await self._purge( name=name, filter=filter, @@ -3373,7 +2952,7 @@ async def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _memory_bank_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3381,86 +2960,3 @@ async def purge( if operation.error: raise RuntimeError(f"Failed to purge memories: {operation.error}") return operation - - async def ingest_events( - self, - *, - name: str, - scope: dict[str, str], - stream_id: str = "", - direct_contents_source: Optional[ - types.IngestionDirectContentsSourceOrDict - ] = None, - generation_trigger_config: Optional[ - types.MemoryGenerationTriggerConfigOrDict - ] = None, - config: Optional[types.IngestEventsConfigOrDict] = None, - ) -> types.MemoryBankIngestEventsOperation: - """Ingests events into an Agent Engine. - - Example usage: - ``` - await client.aio.agent_engines.memories.ingest_events( - name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", - scope={"user_id": "test-user-id"}, - direct_contents_source={ - "events": [ - { - "content": { - "role": "user", - "parts": [ - {"text": "I am a software engineer."} - ], - } - } - ] - }, - generation_trigger_config={ - "generation_rule": { - "idle_duration": "60s" - } - } - ) - ``` - - Args: - name (str): - Required. The name of the Agent Engine to ingest events into. - scope (dict[str, str]): - Required. The scope of the events to ingest. For example, - {"user_id": "123"}. - stream_id (str): - Optional. The ID of the stream to ingest events into. If not - specified, the events will be ingested into the default stream. - direct_contents_source (IngestionDirectContentsSource): - The direct contents source, containing the events to ingest. - generation_trigger_config (MemoryGenerationTriggerConfig): - Optional. The configuration for the generation trigger config. - config (IngestEventsConfig): - Optional. The configuration for the ingest events operation. - - Returns: - AgentEngineIngestEventsOperation: - The operation for ingesting the events. - """ - if config is None: - config = types.IngestEventsConfig() - elif isinstance(config, dict): - config = types.IngestEventsConfig.model_validate(config) - operation = await self._ingest_events( - name=name, - scope=scope, - stream_id=stream_id, - generation_trigger_config=generation_trigger_config, - direct_contents_source=direct_contents_source, - config=config, - ) - if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( - operation_name=operation.name, - get_operation_fn=self._get_memory_operation, - poll_interval_seconds=0.5, - ) - if operation.error: - raise RuntimeError(f"Failed to ingest events: {operation.error}") - return operation diff --git a/agentplatform/_genai/memory_banks.py b/agentplatform/_genai/memory_banks.py new file mode 100644 index 0000000000..40b896fb86 --- /dev/null +++ b/agentplatform/_genai/memory_banks.py @@ -0,0 +1,1023 @@ +# 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. +# + +# Code generated by the Google Gen AI SDK generator DO NOT EDIT. + +import importlib +import json +import logging +import typing +from typing import Any, 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 . import _memory_bank_utils +from . import types + +if typing.TYPE_CHECKING: + from . import memories as memories_module + + _ = memories_module + + +logger = logging.getLogger("agentplatform_genai.memorybanks") + + +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] = {} + + return to_object + + +def _DeleteMemoryBankRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["force"]) is not None: + setv(to_object, ["force"], getv(from_object, ["force"])) + + return to_object + + +def _GetMemoryBankOperationParameters_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, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + +def _IngestEventsConfig_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, ["force_flush"]) is not None: + setv(parent_object, ["forceFlush"], getv(from_object, ["force_flush"])) + + if getv(from_object, ["revision_labels"]) is not None: + setv(parent_object, ["revisionLabels"], getv(from_object, ["revision_labels"])) + + if getv(from_object, ["revision_expire_time"]) is not None: + setv( + parent_object, + ["revisionExpireTime"], + getv(from_object, ["revision_expire_time"]), + ) + + if getv(from_object, ["revision_ttl"]) is not None: + setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) + + if getv(from_object, ["disable_memory_revisions"]) is not None: + setv( + parent_object, + ["disableMemoryRevisions"], + getv(from_object, ["disable_memory_revisions"]), + ) + + if getv(from_object, ["metadata"]) is not None: + setv( + parent_object, + ["metadata"], + {k: v for k, v in getv(from_object, ["metadata"]).items()}, + ) + + if getv(from_object, ["metadata_merge_strategy"]) is not None: + setv( + parent_object, + ["metadataMergeStrategy"], + getv(from_object, ["metadata_merge_strategy"]), + ) + + return to_object + + +def _IngestEventsRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["stream_id"]) is not None: + setv(to_object, ["streamId"], getv(from_object, ["stream_id"])) + + if getv(from_object, ["direct_contents_source"]) is not None: + setv( + to_object, + ["directContentsSource"], + getv(from_object, ["direct_contents_source"]), + ) + + if getv(from_object, ["scope"]) is not None: + setv(to_object, ["scope"], getv(from_object, ["scope"])) + + if getv(from_object, ["generation_trigger_config"]) is not None: + setv( + to_object, + ["generationTriggerConfig"], + getv(from_object, ["generation_trigger_config"]), + ) + + if getv(from_object, ["config"]) is not None: + _IngestEventsConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +class MemoryBanks(_api_module.BaseModule): + + def _create( + self, *, config: Optional[types.CreateMemoryBankConfigOrDict] = None + ) -> types.MemoryBankOperation: + """ + Creates a new Memory Bank. + """ + + parameter_model = types._CreateMemoryBankRequestParameters( + 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 = _CreateMemoryBankRequestParameters_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 = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.MemoryBankOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteMemoryBankConfigOrDict] = None, + ) -> types.DeleteMemoryBankOperation: + """ + 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 _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._IngestEventsRequestParameters( + name=name, + stream_id=stream_id, + direct_contents_source=direct_contents_source, + scope=scope, + generation_trigger_config=generation_trigger_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _IngestEventsRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/memories:ingestEvents".format_map(request_url_dict) + else: + path = "{name}/memories:ingestEvents" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.MemoryBankIngestEventsOperation._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_memory_bank_operation( + self, + *, + operation_name: str, + config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, + ) -> types.MemoryBankOperation: + parameter_model = types._GetMemoryBankOperationParameters( + operation_name=operation_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 = _GetMemoryBankOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + 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) + + return_value = types.MemoryBankOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + _memories = None + + @property + def memories(self) -> "memories_module.Memories": + if self._memories is None: + self._memories = importlib.import_module(".memories", __package__) + return self._memories.Memories(self._api_client) # type: ignore[no-any-return] + + def create(self): + """Creates a new Memory Bank.""" + + operation = self._create() + + operation = _memory_bank_utils._await_operation( + operation_name=operation.name, + get_operation_fn=self._get_memory_bank_operation, + ) + if operation.error: + raise RuntimeError(operation.error) + memory_bank = types.MemoryBank(name=operation.response.name) + return memory_bank + + def delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteMemoryBankConfigOrDict] = None, + ) -> types.DeleteMemoryBankOperation: + """ + Delete a Memory Bank resource. + + Args: + name (str): + Required. The name of the Memory Bank to be deleted. Format: + `projects/{project}/locations/{location}/memoryBanks/{resource_id}` + or `memoryBanks/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Memory Bank has undeleted child resources. Defaults to False. + config (DeleteMemoryBankConfig): + Optional. Additional configurations for deleting the Memory Bank. + """ + name = name.replace("memoryBanks", "reasoningEngines") + logger.info(f"Deleting MemoryBank resource: {name}") + operation = self._delete(name=name, force=force, config=config) + logger.info(f"Started MemoryBank delete operation: {operation.name}") + return operation.name + + def ingest_events( + self, + *, + name: str, + scope: dict[str, str], + stream_id: str = "", + direct_contents_source: Optional[ + types.IngestionDirectContentsSourceOrDict + ] = None, + generation_trigger_config: Optional[ + types.MemoryGenerationTriggerConfigOrDict + ] = None, + config: Optional[types.IngestEventsConfigOrDict] = None, + ) -> types.MemoryBankIngestEventsOperation: + """Ingests events into Memory Bank. + + Example usage: + ``` + client.memory_banks.ingest_events( + name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", + scope={"user_id": "test-user-id"}, + direct_contents_source={ + "events": [ + { + "content": { + "role": "user", + "parts": [ + {"text": "I am a software engineer."} + ], + } + } + ] + }, + generation_trigger_config={ + "generation_rule": { + "idle_duration": "60s" + } + } + ) + ``` + + Args: + name (str): + Required. The name of the Memory Bank to ingest events into. + scope (dict[str, str]): + Required. The scope of the events to ingest. For example, + {"user_id": "123"}. + stream_id (str): + Optional. The ID of the stream to ingest events into. If not + specified, the events will be ingested into the default stream. + direct_contents_source (IngestionDirectContentsSource): + The direct contents source, containing the events to ingest. + generation_trigger_config (MemoryGenerationTriggerConfig): + Optional. The configuration for the generation trigger config. + config (IngestEventsConfig): + Optional. The configuration for the ingest events operation. + + Returns: + IngestEventsOperation: + The operation for ingesting the events. + """ + if config is None: + config = types.IngestEventsConfig() + elif isinstance(config, dict): + config = types.IngestEventsConfig.model_validate(config) + operation = self._ingest_events( + name=name, + scope=scope, + stream_id=stream_id, + generation_trigger_config=generation_trigger_config, + direct_contents_source=direct_contents_source, + config=config, + ) + if config.wait_for_completion and not operation.done: + operation = _memory_bank_utils._await_operation( + operation_name=operation.name, + get_operation_fn=self._get_memory_bank_operation, + poll_interval_seconds=0.5, + ) + if operation.error: + raise RuntimeError(f"Failed to ingest events: {operation.error}") + return operation + + +class AsyncMemoryBanks(_api_module.BaseModule): + + async def _create( + self, *, config: Optional[types.CreateMemoryBankConfigOrDict] = None + ) -> types.MemoryBankOperation: + """ + Creates a new Memory Bank. + """ + + parameter_model = types._CreateMemoryBankRequestParameters( + 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 = _CreateMemoryBankRequestParameters_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( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.MemoryBankOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteMemoryBankConfigOrDict] = None, + ) -> types.DeleteMemoryBankOperation: + """ + 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 = await self._api_client.async_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 + + async 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._IngestEventsRequestParameters( + name=name, + stream_id=stream_id, + direct_contents_source=direct_contents_source, + scope=scope, + generation_trigger_config=generation_trigger_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _IngestEventsRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/memories:ingestEvents".format_map(request_url_dict) + else: + path = "{name}/memories:ingestEvents" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.MemoryBankIngestEventsOperation._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, + *, + operation_name: str, + config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, + ) -> types.MemoryBankOperation: + parameter_model = types._GetMemoryBankOperationParameters( + operation_name=operation_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 = _GetMemoryBankOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + 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) + + return_value = types.MemoryBankOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + _memories = None + + @property + def memories(self) -> "memories_module.AsyncMemories": + if self._memories is None: + self._memories = importlib.import_module(".memories", __package__) + return self._memories.AsyncMemories(self._api_client) # type: ignore[no-any-return] + + async def create(self): + """Creates a new Memory Bank.""" + + operation = await self._create() + + operation = await _memory_bank_utils._await_async_operation( + operation_name=operation.name, + get_operation_fn=self._get_memory_bank_operation, + ) + if operation.error: + raise RuntimeError(operation.error) + memory_bank = types.MemoryBank(name=operation.response.name) + return memory_bank + + async def delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteMemoryBankConfigOrDict] = None, + ) -> types.DeleteMemoryBankOperation: + """ + Delete a Memory Bank resource. + + Args: + name (str): + Required. The name of the Memory Bank to be deleted. Format: + `projects/{project}/locations/{location}/memoryBanks/{resource_id}` + or `memoryBanks/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Memory Bank has undeleted child resources. Defaults to False. + config (DeleteMemoryBankConfig): + Optional. Additional configurations for deleting the Memory Bank. + """ + name = name.replace("memoryBanks", "reasoningEngines") + logger.info(f"Deleting MemoryBank resource: {name}") + operation = await self._delete(name=name, force=force, config=config) + logger.info(f"Started MemoryBank delete operation: {operation.name}") + return operation.name + + async def ingest_events( + self, + *, + name: str, + scope: dict[str, str], + stream_id: str = "", + direct_contents_source: Optional[ + types.IngestionDirectContentsSourceOrDict + ] = None, + generation_trigger_config: Optional[ + types.MemoryGenerationTriggerConfigOrDict + ] = None, + config: Optional[types.IngestEventsConfigOrDict] = None, + ) -> types.MemoryBankIngestEventsOperation: + """Ingests events into Memory Bank. + + Example usage: + ``` + await client.aio.memory_banks.ingest_events( + name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", + scope={"user_id": "test-user-id"}, + direct_contents_source={ + "events": [ + { + "content": { + "role": "user", + "parts": [ + {"text": "I am a software engineer."} + ], + } + } + ] + }, + generation_trigger_config={ + "generation_rule": { + "idle_duration": "60s" + } + } + ) + ``` + + Args: + name (str): + Required. The name of the Memory Bank to ingest events into. + scope (dict[str, str]): + Required. The scope of the events to ingest. For example, + {"user_id": "123"}. + stream_id (str): + Optional. The ID of the stream to ingest events into. If not + specified, the events will be ingested into the default stream. + direct_contents_source (IngestionDirectContentsSource): + The direct contents source, containing the events to ingest. + generation_trigger_config (MemoryGenerationTriggerConfig): + Optional. The configuration for the generation trigger config. + config (IngestEventsConfig): + Optional. The configuration for the ingest events operation. + + Returns: + IngestEventsOperation: + The operation for ingesting the events. + """ + if config is None: + config = types.IngestEventsConfig() + elif isinstance(config, dict): + config = types.IngestEventsConfig.model_validate(config) + operation = await self._ingest_events( + name=name, + scope=scope, + stream_id=stream_id, + generation_trigger_config=generation_trigger_config, + direct_contents_source=direct_contents_source, + config=config, + ) + if config.wait_for_completion and not operation.done: + operation = await _memory_bank_utils._await_async_operation( + operation_name=operation.name, + get_operation_fn=self._get_memory_bank_operation, + poll_interval_seconds=0.5, + ) + if operation.error: + raise RuntimeError(f"Failed to ingest events: {operation.error}") + return operation diff --git a/agentplatform/_genai/memory_revisions.py b/agentplatform/_genai/memory_revisions.py index c1a1e307b1..73917a973b 100644 --- a/agentplatform/_genai/memory_revisions.py +++ b/agentplatform/_genai/memory_revisions.py @@ -31,10 +31,8 @@ logger = logging.getLogger("agentplatform_genai.memoryrevisions") -logger.setLevel(logging.INFO) - -def _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( +def _GetMemoryRevisionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -45,7 +43,7 @@ def _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( return to_object -def _ListAgentEngineMemoryRevisionsConfig_to_vertex( +def _ListMemoryRevisionsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -63,7 +61,7 @@ def _ListAgentEngineMemoryRevisionsConfig_to_vertex( return to_object -def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( +def _ListMemoryRevisionsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -72,9 +70,7 @@ def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineMemoryRevisionsConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _ListMemoryRevisionsConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -82,26 +78,23 @@ def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( class MemoryRevisions(_api_module.BaseModule): def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetMemoryRevisionConfigOrDict] = None ) -> types.MemoryRevision: """ - Gets an agent engine memory revision. + Gets a memory revision. Args: - name (str): Required. The name of the Agent Engine memory revision to get. Format: + name (str): Required. The name of the memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetAgentEngineMemoryRevisionConfig): - Optional. Additional configurations for getting the Agent Engine memory revision. + config (GetMemoryRevisionConfig): + Optional. Additional configurations for getting the memory revision. Returns: - AgentEngineMemoryRevision: The requested Agent Engine memory revision. + MemoryRevision: The requested memory revision. """ - parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( + parameter_model = types._GetMemoryRevisionRequestParameters( name=name, config=config, ) @@ -112,7 +105,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -169,23 +162,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, - ) -> types.ListAgentEngineMemoryRevisionsResponse: + config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, + ) -> types.ListMemoryRevisionsResponse: """ - Lists Agent Engine memory revisions. + Lists memory revisions. Args: - name (str): Required. The name of the Agent Engine memory to list revisions for. Format: + name (str): Required. The name of the memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListAgentEngineMemoryRevisionsConfig): - Optional. Additional configurations for listing the Agent Engine memory revisions. + config (ListMemoryRevisionsConfig): + Optional. Additional configurations for listing the memory revisions. Returns: - ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. + LisemoryRevisionsResponse: The requested memory revisions. """ - parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( + parameter_model = types._ListMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -196,7 +189,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -225,7 +218,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( + return_value = types.ListMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -253,14 +246,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, ) -> Iterator[types.MemoryRevision]: - """Lists Agent Engine memory revisions. + """Lists memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListAgentEngineMemoryRevisionsConfigOrDict): + config (ListMemoryRevisionsConfigOrDict): Optional. The configuration for the memories to list revisions. Returns: @@ -278,26 +271,23 @@ def list( class AsyncMemoryRevisions(_api_module.BaseModule): async def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetMemoryRevisionConfigOrDict] = None ) -> types.MemoryRevision: """ - Gets an agent engine memory revision. + Gets a memory revision. Args: - name (str): Required. The name of the Agent Engine memory revision to get. Format: + name (str): Required. The name of the memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetAgentEngineMemoryRevisionConfig): - Optional. Additional configurations for getting the Agent Engine memory revision. + config (GetMemoryRevisionConfig): + Optional. Additional configurations for getting the memory revision. Returns: - AgentEngineMemoryRevision: The requested Agent Engine memory revision. + MemoryRevision: The requested memory revision. """ - parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( + parameter_model = types._GetMemoryRevisionRequestParameters( name=name, config=config, ) @@ -308,7 +298,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -367,23 +357,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, - ) -> types.ListAgentEngineMemoryRevisionsResponse: + config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, + ) -> types.ListMemoryRevisionsResponse: """ - Lists Agent Engine memory revisions. + Lists memory revisions. Args: - name (str): Required. The name of the Agent Engine memory to list revisions for. Format: + name (str): Required. The name of the memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListAgentEngineMemoryRevisionsConfig): - Optional. Additional configurations for listing the Agent Engine memory revisions. + config (ListMemoryRevisionsConfig): + Optional. Additional configurations for listing the memory revisions. Returns: - ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. + LisemoryRevisionsResponse: The requested memory revisions. """ - parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( + parameter_model = types._ListMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -394,7 +384,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -425,7 +415,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( + return_value = types.ListMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -453,14 +443,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, ) -> AsyncPager[types.MemoryRevision]: - """Lists Agent Engine memory revisions. + """Lists memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListAgentEngineMemoryRevisionsConfigOrDict): + config (ListMemoryRevisionsConfigOrDict): Optional. The configuration for the memories to list revisions. Returns: diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 6d719ddd84..9e9278030a 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -29,7 +29,6 @@ from .common import _AssessDatasetParameters from .common import _CancelQueryJobAgentEngineRequestParameters from .common import _CheckQueryJobAgentEngineRequestParameters -from .common import _CreateAgentEngineMemoryRequestParameters from .common import _CreateAgentEngineRequestParameters from .common import _CreateAgentEngineSandboxRequestParameters from .common import _CreateAgentEngineSessionRequestParameters @@ -41,6 +40,8 @@ from .common import _CreateEvaluationMetricParameters from .common import _CreateEvaluationRunParameters from .common import _CreateEvaluationSetParameters +from .common import _CreateMemoryBankRequestParameters +from .common import _CreateMemoryRequestParameters from .common import _CreateMultimodalDatasetParameters from .common import _CreateRagCorpusRequestParameters from .common import _CreateRuntimeFeedbackEntryRequestParameters @@ -49,7 +50,6 @@ from .common import _CreateSkillRequestParameters from .common import _CustomJobParameters from .common import _CustomJobParameters -from .common import _DeleteAgentEngineMemoryRequestParameters from .common import _DeleteAgentEngineRequestParameters from .common import _DeleteAgentEngineRuntimeRevisionRequestParameters from .common import _DeleteAgentEngineSandboxRequestParameters @@ -58,6 +58,8 @@ from .common import _DeleteDatasetRequestParameters from .common import _DeleteEvaluationExperimentParameters from .common import _DeleteEvaluationMetricParameters +from .common import _DeleteMemoryBankRequestParameters +from .common import _DeleteMemoryRequestParameters from .common import _DeleteMultimodalDatasetRequestParameters from .common import _DeletePromptVersionRequestParameters from .common import _DeleteRagCorpusRequestParameters @@ -69,14 +71,10 @@ from .common import _EvaluateInstancesRequestParameters from .common import _ExecuteCodeAgentEngineSandboxRequestParameters from .common import _ExportPublisherModelRequestParameters -from .common import _GenerateAgentEngineMemoriesRequestParameters from .common import _GenerateInstanceRubricsRequest from .common import _GenerateLossClustersParameters +from .common import _GenerateMemoriesRequestParameters from .common import _GenerateUserScenariosParameters -from .common import _GetAgentEngineGenerateMemoriesOperationParameters -from .common import _GetAgentEngineMemoryOperationParameters -from .common import _GetAgentEngineMemoryRequestParameters -from .common import _GetAgentEngineMemoryRevisionRequestParameters from .common import _GetAgentEngineOperationParameters from .common import _GetAgentEngineRequestParameters from .common import _GetAgentEngineRuntimeRevisionRequestParameters @@ -99,7 +97,12 @@ from .common import _GetEvaluationRunParameters from .common import _GetEvaluationSetParameters from .common import _GetExportPublisherModelOperationParameters +from .common import _GetGenerateMemoriesOperationParameters from .common import _GetImportFilesOperationParameters +from .common import _GetMemoryBankOperationParameters +from .common import _GetMemoryOperationParameters +from .common import _GetMemoryRequestParameters +from .common import _GetMemoryRevisionRequestParameters from .common import _GetMultimodalDatasetOperationParameters from .common import _GetMultimodalDatasetParameters from .common import _GetPublisherModelRequestParameters @@ -120,8 +123,6 @@ from .common import _GetSkillRevisionRequestParameters from .common import _ImportRagFilesRequestParameters from .common import _IngestEventsRequestParameters -from .common import _ListAgentEngineMemoryRequestParameters -from .common import _ListAgentEngineMemoryRevisionsRequestParameters from .common import _ListAgentEngineRequestParameters from .common import _ListAgentEngineRuntimeRevisionsRequestParameters from .common import _ListAgentEngineSandboxesRequestParameters @@ -133,6 +134,8 @@ from .common import _ListDatasetVersionsRequestParameters from .common import _ListEvaluationExperimentsParameters from .common import _ListEvaluationMetricsParameters +from .common import _ListMemoriesRequestParameters +from .common import _ListMemoryRevisionsRequestParameters from .common import _ListMultimodalDatasetsRequestParameters from .common import _ListPublisherModelsRequestParameters from .common import _ListRagCorporaRequestParameters @@ -144,25 +147,25 @@ from .common import _ListSkillsRequestParameters from .common import _OptimizeRequestParameters from .common import _OptimizeRequestParameters -from .common import _PurgeAgentEngineMemoriesRequestParameters +from .common import _PurgeMemoriesRequestParameters from .common import _QueryAgentEngineRequestParameters from .common import _QueryAgentEngineRuntimeRevisionRequestParameters from .common import _RecommendSpecRequestParameters from .common import _RestoreVersionRequestParameters -from .common import _RetrieveAgentEngineMemoriesRequestParameters +from .common import _RetrieveMemoriesRequestParameters from .common import _RetrieveMemoryProfilesRequestParameters from .common import _RetrieveRagContextsRequestParameters from .common import _RetrieveSkillsRequestParameters -from .common import _RollbackAgentEngineMemoryRequestParameters +from .common import _RollbackMemoryRequestParameters from .common import _RunQueryJobAgentEngineConfig from .common import _RunQueryJobAgentEngineConfigDict from .common import _RunQueryJobAgentEngineConfigOrDict from .common import _RunQueryJobAgentEngineRequestParameters -from .common import _UpdateAgentEngineMemoryRequestParameters from .common import _UpdateAgentEngineRequestParameters from .common import _UpdateAgentEngineSessionRequestParameters from .common import _UpdateDatasetParameters from .common import _UpdateEvaluationExperimentParameters +from .common import _UpdateMemoryRequestParameters from .common import _UpdateMultimodalDatasetParameters from .common import _UpdateRagConfigRequestParameters from .common import _UpdateRagCorpusRequestParameters @@ -180,25 +183,10 @@ from .common import AgentEngineConfigDict from .common import AgentEngineConfigOrDict from .common import AgentEngineDict -from .common import AgentEngineGenerateMemoriesOperation -from .common import AgentEngineGenerateMemoriesOperationDict -from .common import AgentEngineGenerateMemoriesOperationOrDict -from .common import AgentEngineMemoryConfig -from .common import AgentEngineMemoryConfigDict -from .common import AgentEngineMemoryConfigOrDict -from .common import AgentEngineMemoryOperation -from .common import AgentEngineMemoryOperationDict -from .common import AgentEngineMemoryOperationOrDict from .common import AgentEngineOperation from .common import AgentEngineOperationDict from .common import AgentEngineOperationOrDict from .common import AgentEngineOrDict -from .common import AgentEnginePurgeMemoriesOperation -from .common import AgentEnginePurgeMemoriesOperationDict -from .common import AgentEnginePurgeMemoriesOperationOrDict -from .common import AgentEngineRollbackMemoryOperation -from .common import AgentEngineRollbackMemoryOperationDict -from .common import AgentEngineRollbackMemoryOperationOrDict from .common import AgentEngineRuntimeRevision from .common import AgentEngineRuntimeRevisionDict from .common import AgentEngineRuntimeRevisionOrDict @@ -362,6 +350,9 @@ from .common import CreateEvaluationSetConfig from .common import CreateEvaluationSetConfigDict from .common import CreateEvaluationSetConfigOrDict +from .common import CreateMemoryBankConfig +from .common import CreateMemoryBankConfigDict +from .common import CreateMemoryBankConfigOrDict from .common import CreateMultimodalDatasetConfig from .common import CreateMultimodalDatasetConfigDict from .common import CreateMultimodalDatasetConfigOrDict @@ -414,12 +405,6 @@ from .common import DeleteAgentEngineConfig from .common import DeleteAgentEngineConfigDict from .common import DeleteAgentEngineConfigOrDict -from .common import DeleteAgentEngineMemoryConfig -from .common import DeleteAgentEngineMemoryConfigDict -from .common import DeleteAgentEngineMemoryConfigOrDict -from .common import DeleteAgentEngineMemoryOperation -from .common import DeleteAgentEngineMemoryOperationDict -from .common import DeleteAgentEngineMemoryOperationOrDict from .common import DeleteAgentEngineOperation from .common import DeleteAgentEngineOperationDict from .common import DeleteAgentEngineOperationOrDict @@ -456,6 +441,18 @@ from .common import DeleteEvaluationMetricOperation from .common import DeleteEvaluationMetricOperationDict from .common import DeleteEvaluationMetricOperationOrDict +from .common import DeleteMemoryBankConfig +from .common import DeleteMemoryBankConfigDict +from .common import DeleteMemoryBankConfigOrDict +from .common import DeleteMemoryBankOperation +from .common import DeleteMemoryBankOperationDict +from .common import DeleteMemoryBankOperationOrDict +from .common import DeleteMemoryConfig +from .common import DeleteMemoryConfigDict +from .common import DeleteMemoryConfigOrDict +from .common import DeleteMemoryOperation +from .common import DeleteMemoryOperationDict +from .common import DeleteMemoryOperationOrDict from .common import DeletePromptConfig from .common import DeletePromptConfigDict from .common import DeletePromptConfigOrDict @@ -670,9 +667,6 @@ from .common import GeminiTemplateConfig from .common import GeminiTemplateConfigDict from .common import GeminiTemplateConfigOrDict -from .common import GenerateAgentEngineMemoriesConfig -from .common import GenerateAgentEngineMemoriesConfigDict -from .common import GenerateAgentEngineMemoriesConfigOrDict from .common import GenerateInstanceRubricsResponse from .common import GenerateInstanceRubricsResponseDict from .common import GenerateInstanceRubricsResponseOrDict @@ -685,6 +679,12 @@ from .common import GenerateLossClustersResponse from .common import GenerateLossClustersResponseDict from .common import GenerateLossClustersResponseOrDict +from .common import GenerateMemoriesConfig +from .common import GenerateMemoriesConfigDict +from .common import GenerateMemoriesConfigOrDict +from .common import GenerateMemoriesOperation +from .common import GenerateMemoriesOperationDict +from .common import GenerateMemoriesOperationOrDict from .common import GenerateMemoriesRequestDirectContentsSource from .common import GenerateMemoriesRequestDirectContentsSourceDict from .common import GenerateMemoriesRequestDirectContentsSourceEvent @@ -716,12 +716,6 @@ from .common import GetAgentEngineConfig from .common import GetAgentEngineConfigDict from .common import GetAgentEngineConfigOrDict -from .common import GetAgentEngineMemoryConfig -from .common import GetAgentEngineMemoryConfigDict -from .common import GetAgentEngineMemoryConfigOrDict -from .common import GetAgentEngineMemoryRevisionConfig -from .common import GetAgentEngineMemoryRevisionConfigDict -from .common import GetAgentEngineMemoryRevisionConfigOrDict from .common import GetAgentEngineOperationConfig from .common import GetAgentEngineOperationConfigDict from .common import GetAgentEngineOperationConfigOrDict @@ -767,6 +761,15 @@ from .common import GetImportFilesOperationConfig from .common import GetImportFilesOperationConfigDict from .common import GetImportFilesOperationConfigOrDict +from .common import GetMemoryBankOperationConfig +from .common import GetMemoryBankOperationConfigDict +from .common import GetMemoryBankOperationConfigOrDict +from .common import GetMemoryConfig +from .common import GetMemoryConfigDict +from .common import GetMemoryConfigOrDict +from .common import GetMemoryRevisionConfig +from .common import GetMemoryRevisionConfigDict +from .common import GetMemoryRevisionConfigOrDict from .common import GetMultimodalDatasetOperationConfig from .common import GetMultimodalDatasetOperationConfigDict from .common import GetMultimodalDatasetOperationConfigOrDict @@ -874,15 +877,6 @@ from .common import ListAgentEngineConfig from .common import ListAgentEngineConfigDict from .common import ListAgentEngineConfigOrDict -from .common import ListAgentEngineMemoryConfig -from .common import ListAgentEngineMemoryConfigDict -from .common import ListAgentEngineMemoryConfigOrDict -from .common import ListAgentEngineMemoryRevisionsConfig -from .common import ListAgentEngineMemoryRevisionsConfigDict -from .common import ListAgentEngineMemoryRevisionsConfigOrDict -from .common import ListAgentEngineMemoryRevisionsResponse -from .common import ListAgentEngineMemoryRevisionsResponseDict -from .common import ListAgentEngineMemoryRevisionsResponseOrDict from .common import ListAgentEngineRuntimeRevisionsConfig from .common import ListAgentEngineRuntimeRevisionsConfigDict from .common import ListAgentEngineRuntimeRevisionsConfigOrDict @@ -937,6 +931,18 @@ from .common import ListEvaluationMetricsResponse from .common import ListEvaluationMetricsResponseDict from .common import ListEvaluationMetricsResponseOrDict +from .common import ListMemoriesConfig +from .common import ListMemoriesConfigDict +from .common import ListMemoriesConfigOrDict +from .common import ListMemoriesResponse +from .common import ListMemoriesResponseDict +from .common import ListMemoriesResponseOrDict +from .common import ListMemoryRevisionsConfig +from .common import ListMemoryRevisionsConfigDict +from .common import ListMemoryRevisionsConfigOrDict +from .common import ListMemoryRevisionsResponse +from .common import ListMemoryRevisionsResponseDict +from .common import ListMemoryRevisionsResponseOrDict from .common import ListModelGardenModelsConfig from .common import ListModelGardenModelsConfigDict from .common import ListModelGardenModelsConfigOrDict @@ -970,9 +976,6 @@ from .common import ListRagFilesResponse from .common import ListRagFilesResponseDict from .common import ListRagFilesResponseOrDict -from .common import ListReasoningEnginesMemoriesResponse -from .common import ListReasoningEnginesMemoriesResponseDict -from .common import ListReasoningEnginesMemoriesResponseOrDict from .common import ListReasoningEnginesResponse from .common import ListReasoningEnginesResponseDict from .common import ListReasoningEnginesResponseOrDict @@ -1040,6 +1043,7 @@ from .common import MapInstanceDict from .common import MapInstanceOrDict from .common import Memory +from .common import MemoryBank from .common import MemoryBankCustomizationConfig from .common import MemoryBankCustomizationConfigConsolidationConfig from .common import MemoryBankCustomizationConfigConsolidationConfigDict @@ -1083,9 +1087,17 @@ from .common import MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict from .common import MemoryBankCustomizationConfigMemoryTopicOrDict from .common import MemoryBankCustomizationConfigOrDict +from .common import MemoryBankDict from .common import MemoryBankIngestEventsOperation from .common import MemoryBankIngestEventsOperationDict from .common import MemoryBankIngestEventsOperationOrDict +from .common import MemoryBankOperation +from .common import MemoryBankOperationDict +from .common import MemoryBankOperationOrDict +from .common import MemoryBankOrDict +from .common import MemoryConfig +from .common import MemoryConfigDict +from .common import MemoryConfigOrDict from .common import MemoryConjunctionFilter from .common import MemoryConjunctionFilterDict from .common import MemoryConjunctionFilterOrDict @@ -1103,6 +1115,9 @@ from .common import MemoryMetadataValue from .common import MemoryMetadataValueDict from .common import MemoryMetadataValueOrDict +from .common import MemoryOperation +from .common import MemoryOperationDict +from .common import MemoryOperationOrDict from .common import MemoryOrDict from .common import MemoryProfile from .common import MemoryProfileDict @@ -1269,9 +1284,12 @@ from .common import PublisherModelResourceReference from .common import PublisherModelResourceReferenceDict from .common import PublisherModelResourceReferenceOrDict -from .common import PurgeAgentEngineMemoriesConfig -from .common import PurgeAgentEngineMemoriesConfigDict -from .common import PurgeAgentEngineMemoriesConfigOrDict +from .common import PurgeMemoriesConfig +from .common import PurgeMemoriesConfigDict +from .common import PurgeMemoriesConfigOrDict +from .common import PurgeMemoriesOperation +from .common import PurgeMemoriesOperationDict +from .common import PurgeMemoriesOperationOrDict from .common import PurgeMemoriesResponse from .common import PurgeMemoriesResponseDict from .common import PurgeMemoriesResponseOrDict @@ -1546,9 +1564,6 @@ from .common import RestoreVersionOperation from .common import RestoreVersionOperationDict from .common import RestoreVersionOperationOrDict -from .common import RetrieveAgentEngineMemoriesConfig -from .common import RetrieveAgentEngineMemoriesConfigDict -from .common import RetrieveAgentEngineMemoriesConfigOrDict from .common import RetrieveContextsConfig from .common import RetrieveContextsConfigDict from .common import RetrieveContextsConfigOrDict @@ -1558,6 +1573,9 @@ from .common import RetrievedSkill from .common import RetrievedSkillDict from .common import RetrievedSkillOrDict +from .common import RetrieveMemoriesConfig +from .common import RetrieveMemoriesConfigDict +from .common import RetrieveMemoriesConfigOrDict from .common import RetrieveMemoriesRequestSimilaritySearchParams from .common import RetrieveMemoriesRequestSimilaritySearchParamsDict from .common import RetrieveMemoriesRequestSimilaritySearchParamsOrDict @@ -1582,9 +1600,12 @@ from .common import RetrieveSkillsResponse from .common import RetrieveSkillsResponseDict from .common import RetrieveSkillsResponseOrDict -from .common import RollbackAgentEngineMemoryConfig -from .common import RollbackAgentEngineMemoryConfigDict -from .common import RollbackAgentEngineMemoryConfigOrDict +from .common import RollbackMemoryConfig +from .common import RollbackMemoryConfigDict +from .common import RollbackMemoryConfigOrDict +from .common import RollbackMemoryOperation +from .common import RollbackMemoryOperationDict +from .common import RollbackMemoryOperationOrDict from .common import RougeInput from .common import RougeInputDict from .common import RougeInputOrDict @@ -1922,15 +1943,15 @@ from .common import UpdateAgentEngineConfig from .common import UpdateAgentEngineConfigDict from .common import UpdateAgentEngineConfigOrDict -from .common import UpdateAgentEngineMemoryConfig -from .common import UpdateAgentEngineMemoryConfigDict -from .common import UpdateAgentEngineMemoryConfigOrDict from .common import UpdateAgentEngineSessionConfig from .common import UpdateAgentEngineSessionConfigDict from .common import UpdateAgentEngineSessionConfigOrDict from .common import UpdateEvaluationExperimentConfig from .common import UpdateEvaluationExperimentConfigDict from .common import UpdateEvaluationExperimentConfigOrDict +from .common import UpdateMemoryConfig +from .common import UpdateMemoryConfigDict +from .common import UpdateMemoryConfigOrDict from .common import UpdatePromptConfig from .common import UpdatePromptConfigDict from .common import UpdatePromptConfigOrDict @@ -2639,27 +2660,57 @@ "UpdateAgentEngineConfig", "UpdateAgentEngineConfigDict", "UpdateAgentEngineConfigOrDict", + "CreateMemoryBankConfig", + "CreateMemoryBankConfigDict", + "CreateMemoryBankConfigOrDict", + "MemoryBank", + "MemoryBankDict", + "MemoryBankOrDict", + "MemoryBankOperation", + "MemoryBankOperationDict", + "MemoryBankOperationOrDict", + "DeleteMemoryBankConfig", + "DeleteMemoryBankConfigDict", + "DeleteMemoryBankConfigOrDict", + "DeleteMemoryBankOperation", + "DeleteMemoryBankOperationDict", + "DeleteMemoryBankOperationOrDict", + "IngestionDirectContentsSourceEvent", + "IngestionDirectContentsSourceEventDict", + "IngestionDirectContentsSourceEventOrDict", + "IngestionDirectContentsSource", + "IngestionDirectContentsSourceDict", + "IngestionDirectContentsSourceOrDict", "MemoryMetadataValue", "MemoryMetadataValueDict", "MemoryMetadataValueOrDict", - "AgentEngineMemoryConfig", - "AgentEngineMemoryConfigDict", - "AgentEngineMemoryConfigOrDict", + "IngestEventsConfig", + "IngestEventsConfigDict", + "IngestEventsConfigOrDict", + "MemoryBankIngestEventsOperation", + "MemoryBankIngestEventsOperationDict", + "MemoryBankIngestEventsOperationOrDict", + "GetMemoryBankOperationConfig", + "GetMemoryBankOperationConfigDict", + "GetMemoryBankOperationConfigOrDict", + "MemoryConfig", + "MemoryConfigDict", + "MemoryConfigOrDict", "MemoryStructuredContent", "MemoryStructuredContentDict", "MemoryStructuredContentOrDict", "Memory", "MemoryDict", "MemoryOrDict", - "AgentEngineMemoryOperation", - "AgentEngineMemoryOperationDict", - "AgentEngineMemoryOperationOrDict", - "DeleteAgentEngineMemoryConfig", - "DeleteAgentEngineMemoryConfigDict", - "DeleteAgentEngineMemoryConfigOrDict", - "DeleteAgentEngineMemoryOperation", - "DeleteAgentEngineMemoryOperationDict", - "DeleteAgentEngineMemoryOperationOrDict", + "MemoryOperation", + "MemoryOperationDict", + "MemoryOperationOrDict", + "DeleteMemoryConfig", + "DeleteMemoryConfigDict", + "DeleteMemoryConfigOrDict", + "DeleteMemoryOperation", + "DeleteMemoryOperationDict", + "DeleteMemoryOperationOrDict", "GenerateMemoriesRequestVertexSessionSource", "GenerateMemoriesRequestVertexSessionSourceDict", "GenerateMemoriesRequestVertexSessionSourceOrDict", @@ -2675,39 +2726,27 @@ "GenerateMemoriesRequestDirectMemoriesSource", "GenerateMemoriesRequestDirectMemoriesSourceDict", "GenerateMemoriesRequestDirectMemoriesSourceOrDict", - "GenerateAgentEngineMemoriesConfig", - "GenerateAgentEngineMemoriesConfigDict", - "GenerateAgentEngineMemoriesConfigOrDict", + "GenerateMemoriesConfig", + "GenerateMemoriesConfigDict", + "GenerateMemoriesConfigOrDict", "GenerateMemoriesResponseGeneratedMemory", "GenerateMemoriesResponseGeneratedMemoryDict", "GenerateMemoriesResponseGeneratedMemoryOrDict", "GenerateMemoriesResponse", "GenerateMemoriesResponseDict", "GenerateMemoriesResponseOrDict", - "AgentEngineGenerateMemoriesOperation", - "AgentEngineGenerateMemoriesOperationDict", - "AgentEngineGenerateMemoriesOperationOrDict", - "GetAgentEngineMemoryConfig", - "GetAgentEngineMemoryConfigDict", - "GetAgentEngineMemoryConfigOrDict", - "IngestionDirectContentsSourceEvent", - "IngestionDirectContentsSourceEventDict", - "IngestionDirectContentsSourceEventOrDict", - "IngestionDirectContentsSource", - "IngestionDirectContentsSourceDict", - "IngestionDirectContentsSourceOrDict", - "IngestEventsConfig", - "IngestEventsConfigDict", - "IngestEventsConfigOrDict", - "MemoryBankIngestEventsOperation", - "MemoryBankIngestEventsOperationDict", - "MemoryBankIngestEventsOperationOrDict", - "ListAgentEngineMemoryConfig", - "ListAgentEngineMemoryConfigDict", - "ListAgentEngineMemoryConfigOrDict", - "ListReasoningEnginesMemoriesResponse", - "ListReasoningEnginesMemoriesResponseDict", - "ListReasoningEnginesMemoriesResponseOrDict", + "GenerateMemoriesOperation", + "GenerateMemoriesOperationDict", + "GenerateMemoriesOperationOrDict", + "GetMemoryConfig", + "GetMemoryConfigDict", + "GetMemoryConfigOrDict", + "ListMemoriesConfig", + "ListMemoriesConfigDict", + "ListMemoriesConfigOrDict", + "ListMemoriesResponse", + "ListMemoriesResponseDict", + "ListMemoriesResponseOrDict", "RetrieveMemoriesRequestSimilaritySearchParams", "RetrieveMemoriesRequestSimilaritySearchParamsDict", "RetrieveMemoriesRequestSimilaritySearchParamsOrDict", @@ -2720,9 +2759,9 @@ "MemoryConjunctionFilter", "MemoryConjunctionFilterDict", "MemoryConjunctionFilterOrDict", - "RetrieveAgentEngineMemoriesConfig", - "RetrieveAgentEngineMemoriesConfigDict", - "RetrieveAgentEngineMemoriesConfigOrDict", + "RetrieveMemoriesConfig", + "RetrieveMemoriesConfigDict", + "RetrieveMemoriesConfigOrDict", "RetrieveMemoriesResponseRetrievedMemory", "RetrieveMemoriesResponseRetrievedMemoryDict", "RetrieveMemoriesResponseRetrievedMemoryOrDict", @@ -2738,39 +2777,39 @@ "RetrieveProfilesResponse", "RetrieveProfilesResponseDict", "RetrieveProfilesResponseOrDict", - "RollbackAgentEngineMemoryConfig", - "RollbackAgentEngineMemoryConfigDict", - "RollbackAgentEngineMemoryConfigOrDict", - "AgentEngineRollbackMemoryOperation", - "AgentEngineRollbackMemoryOperationDict", - "AgentEngineRollbackMemoryOperationOrDict", - "UpdateAgentEngineMemoryConfig", - "UpdateAgentEngineMemoryConfigDict", - "UpdateAgentEngineMemoryConfigOrDict", - "PurgeAgentEngineMemoriesConfig", - "PurgeAgentEngineMemoriesConfigDict", - "PurgeAgentEngineMemoriesConfigOrDict", + "RollbackMemoryConfig", + "RollbackMemoryConfigDict", + "RollbackMemoryConfigOrDict", + "RollbackMemoryOperation", + "RollbackMemoryOperationDict", + "RollbackMemoryOperationOrDict", + "UpdateMemoryConfig", + "UpdateMemoryConfigDict", + "UpdateMemoryConfigOrDict", + "PurgeMemoriesConfig", + "PurgeMemoriesConfigDict", + "PurgeMemoriesConfigOrDict", "PurgeMemoriesResponse", "PurgeMemoriesResponseDict", "PurgeMemoriesResponseOrDict", - "AgentEnginePurgeMemoriesOperation", - "AgentEnginePurgeMemoriesOperationDict", - "AgentEnginePurgeMemoriesOperationOrDict", - "GetAgentEngineMemoryRevisionConfig", - "GetAgentEngineMemoryRevisionConfigDict", - "GetAgentEngineMemoryRevisionConfigOrDict", + "PurgeMemoriesOperation", + "PurgeMemoriesOperationDict", + "PurgeMemoriesOperationOrDict", + "GetMemoryRevisionConfig", + "GetMemoryRevisionConfigDict", + "GetMemoryRevisionConfigOrDict", "IntermediateExtractedMemory", "IntermediateExtractedMemoryDict", "IntermediateExtractedMemoryOrDict", "MemoryRevision", "MemoryRevisionDict", "MemoryRevisionOrDict", - "ListAgentEngineMemoryRevisionsConfig", - "ListAgentEngineMemoryRevisionsConfigDict", - "ListAgentEngineMemoryRevisionsConfigOrDict", - "ListAgentEngineMemoryRevisionsResponse", - "ListAgentEngineMemoryRevisionsResponseDict", - "ListAgentEngineMemoryRevisionsResponseOrDict", + "ListMemoryRevisionsConfig", + "ListMemoryRevisionsConfigDict", + "ListMemoryRevisionsConfigOrDict", + "ListMemoryRevisionsResponse", + "ListMemoryRevisionsResponseDict", + "ListMemoryRevisionsResponseOrDict", "AskContextsConfig", "AskContextsConfigDict", "AskContextsConfigOrDict", @@ -3799,21 +3838,24 @@ "_GetAgentEngineOperationParameters", "_QueryAgentEngineRequestParameters", "_UpdateAgentEngineRequestParameters", - "_CreateAgentEngineMemoryRequestParameters", - "_DeleteAgentEngineMemoryRequestParameters", - "_GenerateAgentEngineMemoriesRequestParameters", - "_GetAgentEngineMemoryRequestParameters", + "_CreateMemoryBankRequestParameters", + "_DeleteMemoryBankRequestParameters", "_IngestEventsRequestParameters", - "_ListAgentEngineMemoryRequestParameters", - "_GetAgentEngineMemoryOperationParameters", - "_GetAgentEngineGenerateMemoriesOperationParameters", - "_RetrieveAgentEngineMemoriesRequestParameters", + "_GetMemoryBankOperationParameters", + "_CreateMemoryRequestParameters", + "_DeleteMemoryRequestParameters", + "_GenerateMemoriesRequestParameters", + "_GetMemoryRequestParameters", + "_ListMemoriesRequestParameters", + "_GetMemoryOperationParameters", + "_GetGenerateMemoriesOperationParameters", + "_RetrieveMemoriesRequestParameters", "_RetrieveMemoryProfilesRequestParameters", - "_RollbackAgentEngineMemoryRequestParameters", - "_UpdateAgentEngineMemoryRequestParameters", - "_PurgeAgentEngineMemoriesRequestParameters", - "_GetAgentEngineMemoryRevisionRequestParameters", - "_ListAgentEngineMemoryRevisionsRequestParameters", + "_RollbackMemoryRequestParameters", + "_UpdateMemoryRequestParameters", + "_PurgeMemoriesRequestParameters", + "_GetMemoryRevisionRequestParameters", + "_ListMemoryRevisionsRequestParameters", "_AskContextsRequestParameters", "_CreateRagCorpusRequestParameters", "_GetCorpusOperationParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 24265d7dab..b0343429c1 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -10184,353 +10184,424 @@ class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False): ] -class MemoryMetadataValue(_common.BaseModel): - """The metadata values for memories.""" +class CreateMemoryBankConfig(_common.BaseModel): + """Config for create memory bank.""" - bool_value: Optional[bool] = Field( - default=None, description="""Represents a boolean value.""" - ) - double_value: Optional[float] = Field( - default=None, description="""Represents a double value.""" - ) - string_value: Optional[str] = Field( - default=None, description="""Represents a string value.""" - ) - timestamp_value: Optional[datetime.datetime] = Field( - default=None, - description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class MemoryMetadataValueDict(TypedDict, total=False): - """The metadata values for memories.""" +class CreateMemoryBankConfigDict(TypedDict, total=False): + """Config for create memory bank.""" - bool_value: Optional[bool] - """Represents a boolean value.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - double_value: Optional[float] - """Represents a double value.""" - string_value: Optional[str] - """Represents a string value.""" +CreateMemoryBankConfigOrDict = Union[CreateMemoryBankConfig, CreateMemoryBankConfigDict] - timestamp_value: Optional[datetime.datetime] - """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" +class _CreateMemoryBankRequestParameters(_common.BaseModel): + """Parameters for creating memory banks.""" -MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] + config: Optional[CreateMemoryBankConfig] = Field(default=None, description="""""") -class AgentEngineMemoryConfig(_common.BaseModel): - """Config for creating a Memory.""" +class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for creating memory banks.""" - 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 display name of the memory.""" - ) - description: Optional[str] = Field( - default=None, description="""The description of the memory.""" - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) - ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. The TTL for this resource. + config: Optional[CreateMemoryBankConfigDict] + """""" - The expiration time is computed: now + TTL.""", - ) - expire_time: Optional[datetime.datetime] = Field( + +_CreateMemoryBankRequestParametersOrDict = Union[ + _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict +] + + +class MemoryBank(_common.BaseModel): + """A memory bank.""" + + name: Optional[str] = Field( default=None, - description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", + 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.""", ) - revision_expire_time: Optional[datetime.datetime] = Field( + + +class MemoryBankDict(TypedDict, total=False): + """A memory bank.""" + + 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.""" + + +MemoryBankOrDict = Union[MemoryBank, MemoryBankDict] + + +class MemoryBankOperation(_common.BaseModel): + """Operation that has an memory bank as a response.""" + + name: Optional[str] = Field( default=None, - description="""Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - revision_ttl: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - disable_memory_revisions: Optional[bool] = Field( + done: Optional[bool] = Field( default=None, - description="""Optional. Input only. If true, no revision will be created for this request.""", - ) - topics: Optional[list[MemoryTopicId]] = Field( - default=None, description="""Optional. The topics of the memory.""" + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - memory_id: Optional[str] = Field( - default=None, - description="""Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", + response: Optional[MemoryBank] = Field( + default=None, description="""The created Memory Bank.""" ) -class AgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for creating a Memory.""" +class MemoryBankOperationDict(TypedDict, total=False): + """Operation that has an memory bank as a response.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - display_name: Optional[str] - """The display name of the memory.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - description: Optional[str] - """The description of the memory.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. + response: Optional[MemoryBankDict] + """The created Memory Bank.""" - The expiration time is computed: now + TTL.""" - expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" +MemoryBankOperationOrDict = Union[MemoryBankOperation, MemoryBankOperationDict] - revision_expire_time: Optional[datetime.datetime] - """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" +class DeleteMemoryBankConfig(_common.BaseModel): + """Config for delete memory bank.""" - disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revision will be created for this request.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics of the memory.""" - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" +class DeleteMemoryBankConfigDict(TypedDict, total=False): + """Config for delete memory bank.""" - memory_id: Optional[str] - """Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +DeleteMemoryBankConfigOrDict = Union[DeleteMemoryBankConfig, DeleteMemoryBankConfigDict] + + +class _DeleteMemoryBankRequestParameters(_common.BaseModel): + """Parameters for deleting a memory bank.""" + + name: Optional[str] = Field( + default=None, description="""Name of the memory bank.""" + ) + force: Optional[bool] = Field( + default=False, + description="""If set to true, any child resources will also be deleted.""", + ) + config: Optional[DeleteMemoryBankConfig] = Field(default=None, description="""""") + + +class _DeleteMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a memory bank.""" + + name: Optional[str] + """Name of the memory bank.""" + force: Optional[bool] + """If set to true, any child resources will also be deleted.""" + + config: Optional[DeleteMemoryBankConfigDict] + """""" -AgentEngineMemoryConfigOrDict = Union[ - AgentEngineMemoryConfig, AgentEngineMemoryConfigDict + +_DeleteMemoryBankRequestParametersOrDict = Union[ + _DeleteMemoryBankRequestParameters, _DeleteMemoryBankRequestParametersDict ] -class _CreateAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Memories.""" +class DeleteMemoryBankOperation(_common.BaseModel): + """Operation for deleting a memory bank.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to create the memory under.""", + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - fact: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The fact of the memory. - - This is the semantic knowledge extracted from the source content).""", + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - scope: Optional[dict[str, str]] = Field( + done: Optional[bool] = Field( default=None, - description="""The scope of the memory. - - Memories are isolated within their scope. The scope is defined when - creating or generating memories. Up to 5 key-value pairs are accepted, - and scope values cannot contain the wildcard character '*'.""", + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) - config: Optional[AgentEngineMemoryConfig] = Field(default=None, description="""""") -class _CreateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Memories.""" +class DeleteMemoryBankOperationDict(TypedDict, total=False): + """Operation for deleting a memory bank.""" name: Optional[str] - """Name of the agent engine to create the memory under.""" - - fact: Optional[str] - """The fact of the memory. - - This is the semantic knowledge extracted from the source content).""" + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - scope: Optional[dict[str, str]] - """The scope of the memory. + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - Memories are isolated within their scope. The scope is defined when - creating or generating memories. Up to 5 key-value pairs are accepted, - and scope values cannot contain the wildcard character '*'.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - config: Optional[AgentEngineMemoryConfigDict] - """""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -_CreateAgentEngineMemoryRequestParametersOrDict = Union[ - _CreateAgentEngineMemoryRequestParameters, - _CreateAgentEngineMemoryRequestParametersDict, +DeleteMemoryBankOperationOrDict = Union[ + DeleteMemoryBankOperation, DeleteMemoryBankOperationDict ] -class MemoryStructuredContent(_common.BaseModel): - """Represents the structured value of the memory.""" +class IngestionDirectContentsSourceEvent(_common.BaseModel): + """The direct contents source event for ingesting events.""" - data: Optional[dict[str, Any]] = Field( + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. The content of the event.""" + ) + event_id: Optional[str] = Field( default=None, - description="""Required. Represents the structured value of the memory.""", + description="""Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""", ) - schema_id: Optional[str] = Field( + event_time: Optional[datetime.datetime] = Field( default=None, - description="""Required. Represents the schema ID for which this structured memory belongs to.""", + description="""Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""", ) -class MemoryStructuredContentDict(TypedDict, total=False): - """Represents the structured value of the memory.""" +class IngestionDirectContentsSourceEventDict(TypedDict, total=False): + """The direct contents source event for ingesting events.""" - data: Optional[dict[str, Any]] - """Required. Represents the structured value of the memory.""" + content: Optional[genai_types.Content] + """Required. The content of the event.""" - schema_id: Optional[str] - """Required. Represents the schema ID for which this structured memory belongs to.""" + event_id: Optional[str] + """Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""" + event_time: Optional[datetime.datetime] + """Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""" -MemoryStructuredContentOrDict = Union[ - MemoryStructuredContent, MemoryStructuredContentDict + +IngestionDirectContentsSourceEventOrDict = Union[ + IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict ] -class Memory(_common.BaseModel): - """A memory.""" +class IngestionDirectContentsSource(_common.BaseModel): + """The direct contents source for ingesting events.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp when this Memory was created.""", + events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( + default=None, description="""Required. The events to ingest.""" ) - description: Optional[str] = Field( - default=None, - description="""Optional. Represents the description of the Memory.""", + + +class IngestionDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for ingesting events.""" + + events: Optional[list[IngestionDirectContentsSourceEventDict]] + """Required. The events to ingest.""" + + +IngestionDirectContentsSourceOrDict = Union[ + IngestionDirectContentsSource, IngestionDirectContentsSourceDict +] + + +class MemoryMetadataValue(_common.BaseModel): + """The metadata values for memories.""" + + bool_value: Optional[bool] = Field( + default=None, description="""Represents a boolean value.""" ) - disable_memory_revisions: Optional[bool] = Field( - default=None, - description="""Optional. Input only. Indicates whether no revision will be created for this request.""", + double_value: Optional[float] = Field( + default=None, description="""Represents a double value.""" ) - display_name: Optional[str] = Field( - default=None, - description="""Optional. Represents the display name of the Memory.""", - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""", + string_value: Optional[str] = Field( + default=None, description="""Represents a string value.""" ) - fact: Optional[str] = Field( + timestamp_value: Optional[datetime.datetime] = Field( default=None, - description="""Optional. Represents semantic knowledge extracted from the source content.""", + description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( - default=None, - description="""Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", + + +class MemoryMetadataValueDict(TypedDict, total=False): + """The metadata values for memories.""" + + bool_value: Optional[bool] + """Represents a boolean value.""" + + double_value: Optional[float] + """Represents a double value.""" + + string_value: Optional[str] + """Represents a string value.""" + + timestamp_value: Optional[datetime.datetime] + """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" + + +MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] + + +class IngestEventsConfig(_common.BaseModel): + """Config for ingesting events.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - name: Optional[str] = Field( - default=None, - description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", + wait_for_completion: Optional[bool] = Field( + default=False, + description="""Waits for the underlying memory generation operation to complete + before returning. Defaults to false.""", ) - revision_expire_time: Optional[datetime.datetime] = Field( + force_flush: Optional[bool] = Field( default=None, - description="""Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", + description="""Optional. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""", ) revision_labels: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""", - ) - revision_ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""", + description="""Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""", ) - scope: Optional[dict[str, str]] = Field( + revision_expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""", - ) - topics: Optional[list[MemoryTopicId]] = Field( - default=None, description="""Optional. Represents the Topics of the Memory.""" + description="""Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", ) - ttl: Optional[str] = Field( + revision_ttl: Optional[str] = Field( default=None, - description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", + description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", ) - update_time: Optional[datetime.datetime] = Field( + disable_memory_revisions: Optional[bool] = Field( default=None, - description="""Output only. Represents the timestamp when this Memory was most recently updated.""", + description="""Optional. Input only. If true, no revisions will be created for this request.""", ) - memory_type: Optional[MemoryType] = Field( + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( default=None, - description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", + description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", ) - structured_content: Optional[MemoryStructuredContent] = Field( + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( default=None, - description="""Optional. Represents the structured content of the memory.""", + description="""Optional. The strategy to use when applying metadata to existing memories.""", ) -class MemoryDict(TypedDict, total=False): - """A memory.""" +class IngestEventsConfigDict(TypedDict, total=False): + """Config for ingesting events.""" - create_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was created.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - description: Optional[str] - """Optional. Represents the description of the Memory.""" + wait_for_completion: Optional[bool] + """Waits for the underlying memory generation operation to complete + before returning. Defaults to false.""" - disable_memory_revisions: Optional[bool] - """Optional. Input only. Indicates whether no revision will be created for this request.""" + force_flush: Optional[bool] + """Optional. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""" - display_name: Optional[str] - """Optional. Represents the display name of the Memory.""" + revision_labels: Optional[dict[str, str]] + """Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""" - expire_time: Optional[datetime.datetime] - """Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""" + revision_expire_time: Optional[datetime.datetime] + """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" - fact: Optional[str] - """Optional. Represents semantic knowledge extracted from the source content.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + + disable_memory_revisions: Optional[bool] + """Optional. Input only. If true, no revisions will be created for this request.""" metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" + """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" - name: Optional[str] - """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""" + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] + """Optional. The strategy to use when applying metadata to existing memories.""" - revision_expire_time: Optional[datetime.datetime] - """Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" - revision_labels: Optional[dict[str, str]] - """Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""" +IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] - revision_ttl: Optional[str] - """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" - scope: Optional[dict[str, str]] - """Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""" +class _IngestEventsRequestParameters(_common.BaseModel): + """Parameters for ingesting events to Memory Bank.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. Represents the Topics of the Memory.""" + name: Optional[str] = Field( + default=None, description="""Name of the Memory Bank to ingest events into.""" + ) + stream_id: Optional[str] = Field( + default=None, description="""The ID of the stream to ingest events into.""" + ) + direct_contents_source: Optional[IngestionDirectContentsSource] = Field( + default=None, + description="""The direct memories source of the events that should be ingested.""", + ) + scope: Optional[dict[str, str]] = Field( + default=None, + description="""The scope of the memories that should be generated from the stream. - ttl: Optional[str] - """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" + Memories will be consolidated across memories with the same scope. Scope + values cannot contain the wildcard character '*'.""", + ) + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""The configuration for the memory generation trigger.""", + ) + config: Optional[IngestEventsConfig] = Field(default=None, description="""""") - update_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was most recently updated.""" - memory_type: Optional[MemoryType] - """Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""" +class _IngestEventsRequestParametersDict(TypedDict, total=False): + """Parameters for ingesting events to Memory Bank.""" - structured_content: Optional[MemoryStructuredContentDict] - """Optional. Represents the structured content of the memory.""" + name: Optional[str] + """Name of the Memory Bank to ingest events into.""" + + stream_id: Optional[str] + """The ID of the stream to ingest events into.""" + direct_contents_source: Optional[IngestionDirectContentsSourceDict] + """The direct memories source of the events that should be ingested.""" -MemoryOrDict = Union[Memory, MemoryDict] + scope: Optional[dict[str, str]] + """The scope of the memories that should be generated from the stream. + + Memories will be consolidated across memories with the same scope. Scope + values cannot contain the wildcard character '*'.""" + + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """The configuration for the memory generation trigger.""" + + config: Optional[IngestEventsConfigDict] + """""" + + +_IngestEventsRequestParametersOrDict = Union[ + _IngestEventsRequestParameters, _IngestEventsRequestParametersDict +] -class AgentEngineMemoryOperation(_common.BaseModel): - """Operation that has an agent engine memory as a response.""" +class MemoryBankIngestEventsOperation(_common.BaseModel): + """Operation that ingests events into a memory bank.""" name: Optional[str] = Field( default=None, @@ -10548,13 +10619,10 @@ class AgentEngineMemoryOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) - response: Optional[Memory] = Field( - default=None, description="""The Agent Engine Memory.""" - ) -class AgentEngineMemoryOperationDict(TypedDict, total=False): - """Operation that has an agent engine memory as a response.""" +class MemoryBankIngestEventsOperationDict(TypedDict, total=False): + """Operation that ingests events into a memory bank.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -10568,465 +10636,449 @@ class AgentEngineMemoryOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[MemoryDict] - """The Agent Engine Memory.""" - -AgentEngineMemoryOperationOrDict = Union[ - AgentEngineMemoryOperation, AgentEngineMemoryOperationDict +MemoryBankIngestEventsOperationOrDict = Union[ + MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict ] -class DeleteAgentEngineMemoryConfig(_common.BaseModel): - """Config for deleting an Agent Engine Memory.""" +class GetMemoryBankOperationConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Memory.""" +class GetMemoryBankOperationConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineMemoryConfigOrDict = Union[ - DeleteAgentEngineMemoryConfig, DeleteAgentEngineMemoryConfigDict +GetMemoryBankOperationConfigOrDict = Union[ + GetMemoryBankOperationConfig, GetMemoryBankOperationConfigDict ] -class _DeleteAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class _GetMemoryBankOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a memory bank as a response.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to delete.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[DeleteAgentEngineMemoryConfig] = Field( - default=None, description="""""" + config: Optional[GetMemoryBankOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class _DeleteAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class _GetMemoryBankOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a memory bank as a response.""" - name: Optional[str] - """Name of the agent engine memory to delete.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[DeleteAgentEngineMemoryConfigDict] - """""" + config: Optional[GetMemoryBankOperationConfigDict] + """Used to override the default configuration.""" -_DeleteAgentEngineMemoryRequestParametersOrDict = Union[ - _DeleteAgentEngineMemoryRequestParameters, - _DeleteAgentEngineMemoryRequestParametersDict, +_GetMemoryBankOperationParametersOrDict = Union[ + _GetMemoryBankOperationParameters, _GetMemoryBankOperationParametersDict ] -class DeleteAgentEngineMemoryOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class MemoryConfig(_common.BaseModel): + """Config for creating a Memory.""" - name: Optional[str] = Field( + 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 display name of the memory.""" + ) + description: Optional[str] = Field( + default=None, description="""The description of the memory.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) + ttl: Optional[str] = Field( default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + description="""Optional. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""", ) - metadata: Optional[dict[str, Any]] = Field( + expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", ) - done: Optional[bool] = Field( + revision_expire_time: Optional[datetime.datetime] = Field( default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + description="""Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", ) - error: Optional[dict[str, Any]] = Field( + revision_ttl: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + ) + disable_memory_revisions: Optional[bool] = Field( + default=None, + description="""Optional. Input only. If true, no revision will be created for this request.""", + ) + topics: Optional[list[MemoryTopicId]] = Field( + default=None, description="""Optional. The topics of the memory.""" + ) + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + default=None, + description="""Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", + ) + memory_id: Optional[str] = Field( + default=None, + description="""Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", ) -class DeleteAgentEngineMemoryOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" +class MemoryConfigDict(TypedDict, total=False): + """Config for creating a Memory.""" - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + display_name: Optional[str] + """The display name of the memory.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + description: Optional[str] + """The description of the memory.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -DeleteAgentEngineMemoryOperationOrDict = Union[ - DeleteAgentEngineMemoryOperation, DeleteAgentEngineMemoryOperationDict -] + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. + The expiration time is computed: now + TTL.""" -class GenerateMemoriesRequestVertexSessionSource(_common.BaseModel): - """The vertex session source for generating memories.""" + expire_time: Optional[datetime.datetime] + """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" - end_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""", - ) - session: Optional[str] = Field( - default=None, - description="""Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""", - ) - start_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""", - ) + revision_expire_time: Optional[datetime.datetime] + """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" -class GenerateMemoriesRequestVertexSessionSourceDict(TypedDict, total=False): - """The vertex session source for generating memories.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. If true, no revision will be created for this request.""" - end_time: Optional[datetime.datetime] - """Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics of the memory.""" - session: Optional[str] - """Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""" + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" - start_time: Optional[datetime.datetime] - """Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""" + memory_id: Optional[str] + """Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" -GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ - GenerateMemoriesRequestVertexSessionSource, - GenerateMemoriesRequestVertexSessionSourceDict, -] +MemoryConfigOrDict = Union[MemoryConfig, MemoryConfigDict] -class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): +class _CreateMemoryRequestParameters(_common.BaseModel): + """Parameters for creating Memories.""" - content: Optional[genai_types.Content] = Field( + name: Optional[str] = Field( default=None, - description="""Required. A single piece of content from which to generate memories.""", + description="""Name of the memory bank to create the memory under.""", ) + fact: Optional[str] = Field( + default=None, + description="""The fact of the memory. + This is the semantic knowledge extracted from the source content).""", + ) + scope: Optional[dict[str, str]] = Field( + default=None, + description="""The scope of the memory. -class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): - - content: Optional[genai_types.Content] - """Required. A single piece of content from which to generate memories.""" + Memories are isolated within their scope. The scope is defined when + creating or generating memories. Up to 5 key-value pairs are accepted, + and scope values cannot contain the wildcard character '*'.""", + ) + config: Optional[MemoryConfig] = Field(default=None, description="""""") -GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ - GenerateMemoriesRequestDirectContentsSourceEvent, - GenerateMemoriesRequestDirectContentsSourceEventDict, -] +class _CreateMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for creating Memories.""" + name: Optional[str] + """Name of the memory bank to create the memory under.""" -class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): - """The direct contents source for generating memories.""" + fact: Optional[str] + """The fact of the memory. - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( - default=None, - description="""Required. The source content (i.e. chat history) to generate memories from.""", - ) + This is the semantic knowledge extracted from the source content).""" + scope: Optional[dict[str, str]] + """The scope of the memory. -class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for generating memories.""" + Memories are isolated within their scope. The scope is defined when + creating or generating memories. Up to 5 key-value pairs are accepted, + and scope values cannot contain the wildcard character '*'.""" - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] - """Required. The source content (i.e. chat history) to generate memories from.""" + config: Optional[MemoryConfigDict] + """""" -GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ - GenerateMemoriesRequestDirectContentsSource, - GenerateMemoriesRequestDirectContentsSourceDict, +_CreateMemoryRequestParametersOrDict = Union[ + _CreateMemoryRequestParameters, _CreateMemoryRequestParametersDict ] -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): - """A direct memory to upload to Memory Bank.""" +class MemoryStructuredContent(_common.BaseModel): + """Represents the structured value of the memory.""" - fact: Optional[str] = Field( + data: Optional[dict[str, Any]] = Field( default=None, - description="""Required. The fact to consolidate with existing memories.""", + description="""Required. Represents the structured value of the memory.""", ) - topics: Optional[list[MemoryTopicId]] = Field( + schema_id: Optional[str] = Field( default=None, - description="""Optional. The topics that the consolidated memories should be associated with.""", + description="""Required. Represents the schema ID for which this structured memory belongs to.""", ) -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( - TypedDict, total=False -): - """A direct memory to upload to Memory Bank.""" +class MemoryStructuredContentDict(TypedDict, total=False): + """Represents the structured value of the memory.""" - fact: Optional[str] - """Required. The fact to consolidate with existing memories.""" + data: Optional[dict[str, Any]] + """Required. Represents the structured value of the memory.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics that the consolidated memories should be associated with.""" + schema_id: Optional[str] + """Required. Represents the schema ID for which this structured memory belongs to.""" -GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, - GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, +MemoryStructuredContentOrDict = Union[ + MemoryStructuredContent, MemoryStructuredContentDict ] -class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): - """The direct memories source for generating memories.""" +class Memory(_common.BaseModel): + """A memory.""" - direct_memories: Optional[ - list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemory] - ] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""", + description="""Output only. Represents the timestamp when this Memory was created.""", ) - - -class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): - """The direct memories source for generating memories.""" - - direct_memories: Optional[ - list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] - ] - """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" - - -GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSource, - GenerateMemoriesRequestDirectMemoriesSourceDict, -] - - -class GenerateAgentEngineMemoriesConfig(_common.BaseModel): - """Config for generating memories.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + description: Optional[str] = Field( + default=None, + description="""Optional. Represents the description of the Memory.""", ) - disable_consolidation: Optional[bool] = Field( + disable_memory_revisions: Optional[bool] = Field( default=None, - description="""Whether to disable consolidation of memories. - - If true, generated memories will not be consolidated with existing - memories; all generated memories will be added as new memories regardless - of whether they are duplicates of or contradictory to existing memories. - By default, memory consolidation is enabled.""", + description="""Optional. Input only. Indicates whether no revision will be created for this request.""", ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + display_name: Optional[str] = Field( + default=None, + description="""Optional. Represents the display name of the Memory.""", ) - revision_labels: Optional[dict[str, str]] = Field( + expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""", + description="""Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""", + ) + fact: Optional[str] = Field( + default=None, + description="""Optional. Represents semantic knowledge extracted from the source content.""", + ) + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + default=None, + description="""Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", ) revision_expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", + description="""Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", + ) + revision_labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""", ) revision_ttl: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + description="""Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""", ) - disable_memory_revisions: Optional[bool] = Field( + scope: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Input only. If true, no revisions will be created for this request.""", + description="""Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + topics: Optional[list[MemoryTopicId]] = Field( + default=None, description="""Optional. Represents the Topics of the Memory.""" + ) + ttl: Optional[str] = Field( default=None, - description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", + description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", ) - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( + update_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. The strategy to use when applying metadata to existing memories.""", + description="""Output only. Represents the timestamp when this Memory was most recently updated.""", ) - allowed_topics: Optional[list[MemoryTopicId]] = Field( + memory_type: Optional[MemoryType] = Field( default=None, - description="""Optional. Restricts memory generation to a subset of memory topics.""", + description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", + ) + structured_content: Optional[MemoryStructuredContent] = Field( + default=None, + description="""Optional. Represents the structured content of the memory.""", ) -class GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False): - """Config for generating memories.""" +class MemoryDict(TypedDict, total=False): + """A memory.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + create_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was created.""" - disable_consolidation: Optional[bool] - """Whether to disable consolidation of memories. + description: Optional[str] + """Optional. Represents the description of the Memory.""" - If true, generated memories will not be consolidated with existing - memories; all generated memories will be added as new memories regardless - of whether they are duplicates of or contradictory to existing memories. - By default, memory consolidation is enabled.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. Indicates whether no revision will be created for this request.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + display_name: Optional[str] + """Optional. Represents the display name of the Memory.""" + + expire_time: Optional[datetime.datetime] + """Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""" + + fact: Optional[str] + """Optional. Represents semantic knowledge extracted from the source content.""" + + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" - revision_labels: Optional[dict[str, str]] - """Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""" + name: Optional[str] + """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""" revision_expire_time: Optional[datetime.datetime] - """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" + """Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" + + revision_labels: Optional[dict[str, str]] + """Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""" revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" - disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revisions will be created for this request.""" + scope: Optional[dict[str, str]] + """Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""" - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. Represents the Topics of the Memory.""" - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] - """Optional. The strategy to use when applying metadata to existing memories.""" + ttl: Optional[str] + """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" - allowed_topics: Optional[list[MemoryTopicIdDict]] - """Optional. Restricts memory generation to a subset of memory topics.""" + update_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was most recently updated.""" + memory_type: Optional[MemoryType] + """Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""" -GenerateAgentEngineMemoriesConfigOrDict = Union[ - GenerateAgentEngineMemoriesConfig, GenerateAgentEngineMemoriesConfigDict -] + structured_content: Optional[MemoryStructuredContentDict] + """Optional. Represents the structured content of the memory.""" + + +MemoryOrDict = Union[Memory, MemoryDict] -class _GenerateAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for generating agent engine memories.""" +class MemoryOperation(_common.BaseModel): + """Operation that has a memory as a response.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to generate memories for.""", + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The vertex session source of the memories that should be generated.""", - ) - direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( - Field( - default=None, - description="""The direct contents source of the memories that should be generated.""", - ) - ) - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( - Field( - default=None, - description="""The direct memories source of the memories that should be generated.""", - ) + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - scope: Optional[dict[str, str]] = Field( + done: Optional[bool] = Field( default=None, - description="""The scope of the memories that should be generated. - - Memories will be consolidated across memories with the same scope. Must be - provided unless the scope is defined in the source content. If `scope` is - provided, it will override the scope defined in the source content. Scope - values cannot contain the wildcard character '*'.""", + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", ) - config: Optional[GenerateAgentEngineMemoriesConfig] = Field( - default=None, description="""""" + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[Memory] = Field(default=None, description="""The Memory.""") -class _GenerateAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for generating agent engine memories.""" +class MemoryOperationDict(TypedDict, total=False): + """Operation that has a memory as a response.""" name: Optional[str] - """Name of the agent engine to generate memories for.""" - - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] - """The vertex session source of the memories that should be generated.""" - - direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] - """The direct contents source of the memories that should be generated.""" + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] - """The direct memories source of the memories that should be generated.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - scope: Optional[dict[str, str]] - """The scope of the memories that should be generated. + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - Memories will be consolidated across memories with the same scope. Must be - provided unless the scope is defined in the source content. If `scope` is - provided, it will override the scope defined in the source content. Scope - values cannot contain the wildcard character '*'.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - config: Optional[GenerateAgentEngineMemoriesConfigDict] - """""" + response: Optional[MemoryDict] + """The Memory.""" -_GenerateAgentEngineMemoriesRequestParametersOrDict = Union[ - _GenerateAgentEngineMemoriesRequestParameters, - _GenerateAgentEngineMemoriesRequestParametersDict, -] +MemoryOperationOrDict = Union[MemoryOperation, MemoryOperationDict] -class GenerateMemoriesResponseGeneratedMemory(_common.BaseModel): - """A memory that was generated.""" +class DeleteMemoryConfig(_common.BaseModel): + """Config for deleting a Memory.""" - memory: Optional[Memory] = Field( - default=None, description="""The generated memory.""" - ) - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] = Field( - default=None, description="""The action to take.""" - ) - previous_revision: Optional[str] = Field( - default=None, - description="""The previous revision of the Memory before the action was performed. This - field is only set if the action is `UPDATED` or `DELETED`. You can use - this to rollback the Memory to the previous revision, undoing the action. - Format: - `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): - """A memory that was generated.""" - - memory: Optional[MemoryDict] - """The generated memory.""" - - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] - """The action to take.""" +class DeleteMemoryConfigDict(TypedDict, total=False): + """Config for deleting a Memory.""" - previous_revision: Optional[str] - """The previous revision of the Memory before the action was performed. This - field is only set if the action is `UPDATED` or `DELETED`. You can use - this to rollback the Memory to the previous revision, undoing the action. - Format: - `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ - GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict -] +DeleteMemoryConfigOrDict = Union[DeleteMemoryConfig, DeleteMemoryConfigDict] -class GenerateMemoriesResponse(_common.BaseModel): - """The response for generating memories.""" +class _DeleteMemoryRequestParameters(_common.BaseModel): + """Parameters for deleting memories.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( - default=None, description="""The generated memories.""" + name: Optional[str] = Field( + default=None, description="""Name of the memory to delete.""" ) + config: Optional[DeleteMemoryConfig] = Field(default=None, description="""""") -class GenerateMemoriesResponseDict(TypedDict, total=False): - """The response for generating memories.""" +class _DeleteMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for deleting memories.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] - """The generated memories.""" + name: Optional[str] + """Name of the memory to delete.""" + config: Optional[DeleteMemoryConfigDict] + """""" -GenerateMemoriesResponseOrDict = Union[ - GenerateMemoriesResponse, GenerateMemoriesResponseDict + +_DeleteMemoryRequestParametersOrDict = Union[ + _DeleteMemoryRequestParameters, _DeleteMemoryRequestParametersDict ] -class AgentEngineGenerateMemoriesOperation(_common.BaseModel): - """Operation that generates memories for an agent engine.""" +class DeleteMemoryOperation(_common.BaseModel): + """Operation for deleting memories.""" name: Optional[str] = Field( default=None, @@ -11044,13 +11096,10 @@ class AgentEngineGenerateMemoriesOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) - response: Optional[GenerateMemoriesResponse] = Field( - default=None, description="""The response for generating memories.""" - ) -class AgentEngineGenerateMemoriesOperationDict(TypedDict, total=False): - """Operation that generates memories for an agent engine.""" +class DeleteMemoryOperationDict(TypedDict, total=False): + """Operation for deleting memories.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -11064,129 +11113,163 @@ class AgentEngineGenerateMemoriesOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[GenerateMemoriesResponseDict] - """The response for generating memories.""" + +DeleteMemoryOperationOrDict = Union[DeleteMemoryOperation, DeleteMemoryOperationDict] + + +class GenerateMemoriesRequestVertexSessionSource(_common.BaseModel): + """The vertex session source for generating memories.""" + + end_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""", + ) + session: Optional[str] = Field( + default=None, + description="""Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""", + ) + start_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""", + ) + + +class GenerateMemoriesRequestVertexSessionSourceDict(TypedDict, total=False): + """The vertex session source for generating memories.""" + + end_time: Optional[datetime.datetime] + """Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""" + + session: Optional[str] + """Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""" + + start_time: Optional[datetime.datetime] + """Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""" -AgentEngineGenerateMemoriesOperationOrDict = Union[ - AgentEngineGenerateMemoriesOperation, AgentEngineGenerateMemoriesOperationDict +GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ + GenerateMemoriesRequestVertexSessionSource, + GenerateMemoriesRequestVertexSessionSourceDict, ] -class GetAgentEngineMemoryConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory.""" +class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + content: Optional[genai_types.Content] = Field( + default=None, + description="""Required. A single piece of content from which to generate memories.""", ) -class GetAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory.""" +class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + content: Optional[genai_types.Content] + """Required. A single piece of content from which to generate memories.""" -GetAgentEngineMemoryConfigOrDict = Union[ - GetAgentEngineMemoryConfig, GetAgentEngineMemoryConfigDict +GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ + GenerateMemoriesRequestDirectContentsSourceEvent, + GenerateMemoriesRequestDirectContentsSourceEventDict, ] -class _GetAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine.""" +class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): + """The direct contents source for generating memories.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[GetAgentEngineMemoryConfig] = Field( - default=None, description="""""" + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( + default=None, + description="""Required. The source content (i.e. chat history) to generate memories from.""", ) -class _GetAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine.""" - - name: Optional[str] - """Name of the agent engine.""" +class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for generating memories.""" - config: Optional[GetAgentEngineMemoryConfigDict] - """""" + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] + """Required. The source content (i.e. chat history) to generate memories from.""" -_GetAgentEngineMemoryRequestParametersOrDict = Union[ - _GetAgentEngineMemoryRequestParameters, _GetAgentEngineMemoryRequestParametersDict +GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ + GenerateMemoriesRequestDirectContentsSource, + GenerateMemoriesRequestDirectContentsSourceDict, ] -class IngestionDirectContentsSourceEvent(_common.BaseModel): - """The direct contents source event for ingesting events.""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): + """A direct memory to upload to Memory Bank.""" - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. The content of the event.""" - ) - event_id: Optional[str] = Field( + fact: Optional[str] = Field( default=None, - description="""Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""", + description="""Required. The fact to consolidate with existing memories.""", ) - event_time: Optional[datetime.datetime] = Field( + topics: Optional[list[MemoryTopicId]] = Field( default=None, - description="""Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""", + description="""Optional. The topics that the consolidated memories should be associated with.""", ) -class IngestionDirectContentsSourceEventDict(TypedDict, total=False): - """The direct contents source event for ingesting events.""" - - content: Optional[genai_types.Content] - """Required. The content of the event.""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( + TypedDict, total=False +): + """A direct memory to upload to Memory Bank.""" - event_id: Optional[str] - """Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""" + fact: Optional[str] + """Required. The fact to consolidate with existing memories.""" - event_time: Optional[datetime.datetime] - """Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics that the consolidated memories should be associated with.""" -IngestionDirectContentsSourceEventOrDict = Union[ - IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict +GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, + GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, ] -class IngestionDirectContentsSource(_common.BaseModel): - """The direct contents source for ingesting events.""" - - events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( - default=None, description="""Required. The events to ingest.""" +class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): + """The direct memories source for generating memories.""" + + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemory] + ] = Field( + default=None, + description="""Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""", ) -class IngestionDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for ingesting events.""" +class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): + """The direct memories source for generating memories.""" - events: Optional[list[IngestionDirectContentsSourceEventDict]] - """Required. The events to ingest.""" + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] + ] + """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" -IngestionDirectContentsSourceOrDict = Union[ - IngestionDirectContentsSource, IngestionDirectContentsSourceDict +GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSource, + GenerateMemoriesRequestDirectMemoriesSourceDict, ] -class IngestEventsConfig(_common.BaseModel): - """Config for ingesting events.""" +class GenerateMemoriesConfig(_common.BaseModel): + """Config for generating memories.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - wait_for_completion: Optional[bool] = Field( - default=False, - description="""Waits for the underlying memory generation operation to complete - before returning. Defaults to false.""", - ) - force_flush: Optional[bool] = Field( + disable_consolidation: Optional[bool] = Field( default=None, - description="""Optional. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""", + description="""Whether to disable consolidation of memories. + + If true, generated memories will not be consolidated with existing + memories; all generated memories will be added as new memories regardless + of whether they are duplicates of or contradictory to existing memories. + By default, memory consolidation is enabled.""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) revision_labels: Optional[dict[str, str]] = Field( default=None, @@ -11212,20 +11295,28 @@ class IngestEventsConfig(_common.BaseModel): default=None, description="""Optional. The strategy to use when applying metadata to existing memories.""", ) + allowed_topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. Restricts memory generation to a subset of memory topics.""", + ) -class IngestEventsConfigDict(TypedDict, total=False): - """Config for ingesting events.""" +class GenerateMemoriesConfigDict(TypedDict, total=False): + """Config for generating memories.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - wait_for_completion: Optional[bool] - """Waits for the underlying memory generation operation to complete - before returning. Defaults to false.""" + disable_consolidation: Optional[bool] + """Whether to disable consolidation of memories. - force_flush: Optional[bool] - """Optional. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""" + If true, generated memories will not be consolidated with existing + memories; all generated memories will be added as new memories regardless + of whether they are duplicates of or contradictory to existing memories. + By default, memory consolidation is enabled.""" + + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" revision_labels: Optional[dict[str, str]] """Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""" @@ -11245,69 +11336,143 @@ class IngestEventsConfigDict(TypedDict, total=False): metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] """Optional. The strategy to use when applying metadata to existing memories.""" + allowed_topics: Optional[list[MemoryTopicIdDict]] + """Optional. Restricts memory generation to a subset of memory topics.""" + -IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] +GenerateMemoriesConfigOrDict = Union[GenerateMemoriesConfig, GenerateMemoriesConfigDict] -class _IngestEventsRequestParameters(_common.BaseModel): - """Parameters for purging agent engine memories.""" +class _GenerateMemoriesRequestParameters(_common.BaseModel): + """Parameters for generating memories.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine to ingest events into.""" - ) - stream_id: Optional[str] = Field( - default=None, description="""The ID of the stream to ingest events into.""" + default=None, + description="""Name of the Memory Bank to generate memories with.""", ) - direct_contents_source: Optional[IngestionDirectContentsSource] = Field( + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( default=None, - description="""The direct memories source of the events that should be ingested.""", + description="""The vertex session source of the memories that should be generated.""", + ) + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( + Field( + default=None, + description="""The direct contents source of the memories that should be generated.""", + ) + ) + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( + Field( + default=None, + description="""The direct memories source of the memories that should be generated.""", + ) ) scope: Optional[dict[str, str]] = Field( default=None, - description="""The scope of the memories that should be generated from the stream. + description="""The scope of the memories that should be generated. - Memories will be consolidated across memories with the same scope. Scope + Memories will be consolidated across memories with the same scope. Must be + provided unless the scope is defined in the source content. If `scope` is + provided, it will override the scope defined in the source content. Scope values cannot contain the wildcard character '*'.""", ) - generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( - default=None, - description="""The configuration for the memory generation trigger.""", - ) - config: Optional[IngestEventsConfig] = Field(default=None, description="""""") + config: Optional[GenerateMemoriesConfig] = Field(default=None, description="""""") -class _IngestEventsRequestParametersDict(TypedDict, total=False): - """Parameters for purging agent engine memories.""" +class _GenerateMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for generating memories.""" name: Optional[str] - """Name of the Agent Engine to ingest events into.""" + """Name of the Memory Bank to generate memories with.""" - stream_id: Optional[str] - """The ID of the stream to ingest events into.""" + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] + """The vertex session source of the memories that should be generated.""" - direct_contents_source: Optional[IngestionDirectContentsSourceDict] - """The direct memories source of the events that should be ingested.""" + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] + """The direct contents source of the memories that should be generated.""" + + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] + """The direct memories source of the memories that should be generated.""" scope: Optional[dict[str, str]] - """The scope of the memories that should be generated from the stream. + """The scope of the memories that should be generated. - Memories will be consolidated across memories with the same scope. Scope + Memories will be consolidated across memories with the same scope. Must be + provided unless the scope is defined in the source content. If `scope` is + provided, it will override the scope defined in the source content. Scope values cannot contain the wildcard character '*'.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """The configuration for the memory generation trigger.""" - - config: Optional[IngestEventsConfigDict] + config: Optional[GenerateMemoriesConfigDict] """""" -_IngestEventsRequestParametersOrDict = Union[ - _IngestEventsRequestParameters, _IngestEventsRequestParametersDict +_GenerateMemoriesRequestParametersOrDict = Union[ + _GenerateMemoriesRequestParameters, _GenerateMemoriesRequestParametersDict ] -class MemoryBankIngestEventsOperation(_common.BaseModel): - """Operation that ingests events into a memory bank.""" +class GenerateMemoriesResponseGeneratedMemory(_common.BaseModel): + """A memory that was generated.""" + + memory: Optional[Memory] = Field( + default=None, description="""The generated memory.""" + ) + action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] = Field( + default=None, description="""The action to take.""" + ) + previous_revision: Optional[str] = Field( + default=None, + description="""The previous revision of the Memory before the action was performed. This + field is only set if the action is `UPDATED` or `DELETED`. You can use + this to rollback the Memory to the previous revision, undoing the action. + Format: + `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""", + ) + + +class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): + """A memory that was generated.""" + + memory: Optional[MemoryDict] + """The generated memory.""" + + action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] + """The action to take.""" + + previous_revision: Optional[str] + """The previous revision of the Memory before the action was performed. This + field is only set if the action is `UPDATED` or `DELETED`. You can use + this to rollback the Memory to the previous revision, undoing the action. + Format: + `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""" + + +GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ + GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict +] + + +class GenerateMemoriesResponse(_common.BaseModel): + """The response for generating memories.""" + + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( + default=None, description="""The generated memories.""" + ) + + +class GenerateMemoriesResponseDict(TypedDict, total=False): + """The response for generating memories.""" + + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] + """The generated memories.""" + + +GenerateMemoriesResponseOrDict = Union[ + GenerateMemoriesResponse, GenerateMemoriesResponseDict +] + + +class GenerateMemoriesOperation(_common.BaseModel): + """Operation that generates memories with a Memory Bank.""" name: Optional[str] = Field( default=None, @@ -11325,10 +11490,13 @@ class MemoryBankIngestEventsOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[GenerateMemoriesResponse] = Field( + default=None, description="""The response for generating memories.""" + ) -class MemoryBankIngestEventsOperationDict(TypedDict, total=False): - """Operation that ingests events into a memory bank.""" +class GenerateMemoriesOperationDict(TypedDict, total=False): + """Operation that generates memories with a Memory Bank.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -11342,14 +11510,57 @@ class MemoryBankIngestEventsOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" + response: Optional[GenerateMemoriesResponseDict] + """The response for generating memories.""" + -MemoryBankIngestEventsOperationOrDict = Union[ - MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict +GenerateMemoriesOperationOrDict = Union[ + GenerateMemoriesOperation, GenerateMemoriesOperationDict +] + + +class GetMemoryConfig(_common.BaseModel): + """Config for getting a Memory.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetMemoryConfigDict(TypedDict, total=False): + """Config for getting a Memory.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetMemoryConfigOrDict = Union[GetMemoryConfig, GetMemoryConfigDict] + + +class _GetMemoryRequestParameters(_common.BaseModel): + """Parameters for getting a Memory.""" + + name: Optional[str] = Field(default=None, description="""Name of the memory.""") + config: Optional[GetMemoryConfig] = Field(default=None, description="""""") + + +class _GetMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Memory.""" + + name: Optional[str] + """Name of the memory.""" + + config: Optional[GetMemoryConfigDict] + """""" + + +_GetMemoryRequestParametersOrDict = Union[ + _GetMemoryRequestParameters, _GetMemoryRequestParametersDict ] -class ListAgentEngineMemoryConfig(_common.BaseModel): - """Config for listing agent engine memories.""" +class ListMemoriesConfig(_common.BaseModel): + """Config for listing memories.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -11374,8 +11585,8 @@ class ListAgentEngineMemoryConfig(_common.BaseModel): ) -class ListAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for listing agent engine memories.""" +class ListMemoriesConfigDict(TypedDict, total=False): + """Config for listing memories.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -11401,51 +11612,47 @@ class ListAgentEngineMemoryConfigDict(TypedDict, total=False): * `update_time`""" -ListAgentEngineMemoryConfigOrDict = Union[ - ListAgentEngineMemoryConfig, ListAgentEngineMemoryConfigDict -] +ListMemoriesConfigOrDict = Union[ListMemoriesConfig, ListMemoriesConfigDict] -class _ListAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _ListMemoriesRequestParameters(_common.BaseModel): + """Parameters for listing memories.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[ListAgentEngineMemoryConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the Memory Bank.""" ) + config: Optional[ListMemoriesConfig] = Field(default=None, description="""""") -class _ListAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _ListMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for listing memories.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the Memory Bank.""" - config: Optional[ListAgentEngineMemoryConfigDict] + config: Optional[ListMemoriesConfigDict] """""" -_ListAgentEngineMemoryRequestParametersOrDict = Union[ - _ListAgentEngineMemoryRequestParameters, _ListAgentEngineMemoryRequestParametersDict +_ListMemoriesRequestParametersOrDict = Union[ + _ListMemoriesRequestParameters, _ListMemoriesRequestParametersDict ] -class ListReasoningEnginesMemoriesResponse(_common.BaseModel): - """Response for listing agent engine memories.""" +class ListMemoriesResponse(_common.BaseModel): + """Response for listing memories.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" ) next_page_token: Optional[str] = Field(default=None, description="""""") memories: Optional[list[Memory]] = Field( - default=None, description="""List of agent engine memories.""" + default=None, description="""List of memories.""" ) -class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): - """Response for listing agent engine memories.""" +class ListMemoriesResponseDict(TypedDict, total=False): + """Response for listing memories.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -11454,65 +11661,61 @@ class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): """""" memories: Optional[list[MemoryDict]] - """List of agent engine memories.""" + """List of memories.""" -ListReasoningEnginesMemoriesResponseOrDict = Union[ - ListReasoningEnginesMemoriesResponse, ListReasoningEnginesMemoriesResponseDict -] +ListMemoriesResponseOrDict = Union[ListMemoriesResponse, ListMemoriesResponseDict] -class _GetAgentEngineMemoryOperationParameters(_common.BaseModel): +class _GetMemoryOperationParameters(_common.BaseModel): """Parameters for getting an operation with a memory as a response.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetMemoryBankOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetAgentEngineMemoryOperationParametersDict(TypedDict, total=False): +class _GetMemoryOperationParametersDict(TypedDict, total=False): """Parameters for getting an operation with a memory as a response.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetMemoryBankOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineMemoryOperationParametersOrDict = Union[ - _GetAgentEngineMemoryOperationParameters, - _GetAgentEngineMemoryOperationParametersDict, +_GetMemoryOperationParametersOrDict = Union[ + _GetMemoryOperationParameters, _GetMemoryOperationParametersDict ] -class _GetAgentEngineGenerateMemoriesOperationParameters(_common.BaseModel): +class _GetGenerateMemoriesOperationParameters(_common.BaseModel): """Parameters for getting an operation with generated memories as a response.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetMemoryBankOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetAgentEngineGenerateMemoriesOperationParametersDict(TypedDict, total=False): +class _GetGenerateMemoriesOperationParametersDict(TypedDict, total=False): """Parameters for getting an operation with generated memories as a response.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetMemoryBankOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineGenerateMemoriesOperationParametersOrDict = Union[ - _GetAgentEngineGenerateMemoriesOperationParameters, - _GetAgentEngineGenerateMemoriesOperationParametersDict, +_GetGenerateMemoriesOperationParametersOrDict = Union[ + _GetGenerateMemoriesOperationParameters, _GetGenerateMemoriesOperationParametersDict ] @@ -11633,7 +11836,7 @@ class MemoryConjunctionFilterDict(TypedDict, total=False): ] -class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): +class RetrieveMemoriesConfig(_common.BaseModel): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11676,7 +11879,7 @@ class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): ) -class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): +class RetrieveMemoriesConfigDict(TypedDict, total=False): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] @@ -11715,17 +11918,15 @@ class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): retrieve memories matching any of the specified `MemoryType` values.""" -RetrieveAgentEngineMemoriesConfigOrDict = Union[ - RetrieveAgentEngineMemoriesConfig, RetrieveAgentEngineMemoriesConfigDict -] +RetrieveMemoriesConfigOrDict = Union[RetrieveMemoriesConfig, RetrieveMemoriesConfigDict] -class _RetrieveAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for retrieving agent engine memories.""" +class _RetrieveMemoriesRequestParameters(_common.BaseModel): + """Parameters for retrieving memories.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to retrieve memories from.""", + description="""Name of the Memory Bank to retrieve memories from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -11747,16 +11948,14 @@ class _RetrieveAgentEngineMemoriesRequestParameters(_common.BaseModel): description="""Parameters for simple (non-similarity search) retrieval.""", ) ) - config: Optional[RetrieveAgentEngineMemoriesConfig] = Field( - default=None, description="""""" - ) + config: Optional[RetrieveMemoriesConfig] = Field(default=None, description="""""") -class _RetrieveAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving agent engine memories.""" +class _RetrieveMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving memories.""" name: Optional[str] - """Name of the agent engine to retrieve memories from.""" + """Name of the Memory Bank to retrieve memories from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -11773,13 +11972,12 @@ class _RetrieveAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] """Parameters for simple (non-similarity search) retrieval.""" - config: Optional[RetrieveAgentEngineMemoriesConfigDict] + config: Optional[RetrieveMemoriesConfigDict] """""" -_RetrieveAgentEngineMemoriesRequestParametersOrDict = Union[ - _RetrieveAgentEngineMemoriesRequestParameters, - _RetrieveAgentEngineMemoriesRequestParametersDict, +_RetrieveMemoriesRequestParametersOrDict = Union[ + _RetrieveMemoriesRequestParameters, _RetrieveMemoriesRequestParametersDict ] @@ -11858,11 +12056,11 @@ class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): - """Parameters for retrieving agent engine memory profiles.""" + """Parameters for retrieving memory profiles.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to retrieve memory profiles from.""", + description="""Name of the Memory Bank to retrieve memory profiles from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -11878,10 +12076,10 @@ class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving agent engine memory profiles.""" + """Parameters for retrieving memory profiles.""" name: Optional[str] - """Name of the agent engine to retrieve memory profiles from.""" + """Name of the Memory Bank to retrieve memory profiles from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -11952,7 +12150,7 @@ class RetrieveProfilesResponseDict(TypedDict, total=False): ] -class RollbackAgentEngineMemoryConfig(_common.BaseModel): +class RollbackMemoryConfig(_common.BaseModel): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11964,7 +12162,7 @@ class RollbackAgentEngineMemoryConfig(_common.BaseModel): ) -class RollbackAgentEngineMemoryConfigDict(TypedDict, total=False): +class RollbackMemoryConfigDict(TypedDict, total=False): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] @@ -11974,45 +12172,40 @@ class RollbackAgentEngineMemoryConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -RollbackAgentEngineMemoryConfigOrDict = Union[ - RollbackAgentEngineMemoryConfig, RollbackAgentEngineMemoryConfigDict -] +RollbackMemoryConfigOrDict = Union[RollbackMemoryConfig, RollbackMemoryConfigDict] -class _RollbackAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for generating agent engine memories.""" +class _RollbackMemoryRequestParameters(_common.BaseModel): + """Parameters for generating memories.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to rollback.""" + default=None, description="""Name of the memory to rollback.""" ) target_revision_id: Optional[str] = Field( default=None, description="""The ID of the revision to rollback to.""" ) - config: Optional[RollbackAgentEngineMemoryConfig] = Field( - default=None, description="""""" - ) + config: Optional[RollbackMemoryConfig] = Field(default=None, description="""""") -class _RollbackAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for generating agent engine memories.""" +class _RollbackMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for generating memories.""" name: Optional[str] - """Name of the agent engine memory to rollback.""" + """Name of the memory to rollback.""" target_revision_id: Optional[str] """The ID of the revision to rollback to.""" - config: Optional[RollbackAgentEngineMemoryConfigDict] + config: Optional[RollbackMemoryConfigDict] """""" -_RollbackAgentEngineMemoryRequestParametersOrDict = Union[ - _RollbackAgentEngineMemoryRequestParameters, - _RollbackAgentEngineMemoryRequestParametersDict, +_RollbackMemoryRequestParametersOrDict = Union[ + _RollbackMemoryRequestParameters, _RollbackMemoryRequestParametersDict ] -class AgentEngineRollbackMemoryOperation(_common.BaseModel): +class RollbackMemoryOperation(_common.BaseModel): """Operation that rolls back a memory.""" name: Optional[str] = Field( @@ -12033,7 +12226,7 @@ class AgentEngineRollbackMemoryOperation(_common.BaseModel): ) -class AgentEngineRollbackMemoryOperationDict(TypedDict, total=False): +class RollbackMemoryOperationDict(TypedDict, total=False): """Operation that rolls back a memory.""" name: Optional[str] @@ -12049,13 +12242,13 @@ class AgentEngineRollbackMemoryOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -AgentEngineRollbackMemoryOperationOrDict = Union[ - AgentEngineRollbackMemoryOperation, AgentEngineRollbackMemoryOperationDict +RollbackMemoryOperationOrDict = Union[ + RollbackMemoryOperation, RollbackMemoryOperationDict ] -class UpdateAgentEngineMemoryConfig(_common.BaseModel): - """Config for updating agent engine memory.""" +class UpdateMemoryConfig(_common.BaseModel): + """Config for updating a memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -12110,8 +12303,8 @@ class UpdateAgentEngineMemoryConfig(_common.BaseModel): ) -class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for updating agent engine memory.""" +class UpdateMemoryConfigDict(TypedDict, total=False): + """Config for updating a memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -12156,16 +12349,14 @@ class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" -UpdateAgentEngineMemoryConfigOrDict = Union[ - UpdateAgentEngineMemoryConfig, UpdateAgentEngineMemoryConfigDict -] +UpdateMemoryConfigOrDict = Union[UpdateMemoryConfig, UpdateMemoryConfigDict] -class _UpdateAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for updating agent engine memories.""" +class _UpdateMemoryRequestParameters(_common.BaseModel): + """Parameters for updating memories.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to update.""" + default=None, description="""Name of the memory to update.""" ) fact: Optional[str] = Field( default=None, @@ -12181,16 +12372,14 @@ class _UpdateAgentEngineMemoryRequestParameters(_common.BaseModel): creating or generating memories. Up to 5 key-value pairs are accepted, and scope values cannot contain the wildcard character '*'.""", ) - config: Optional[UpdateAgentEngineMemoryConfig] = Field( - default=None, description="""""" - ) + config: Optional[UpdateMemoryConfig] = Field(default=None, description="""""") -class _UpdateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engine memories.""" +class _UpdateMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for updating memories.""" name: Optional[str] - """Name of the agent engine memory to update.""" + """Name of the memory to update.""" fact: Optional[str] """The updated fact of the memory. @@ -12204,17 +12393,16 @@ class _UpdateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): creating or generating memories. Up to 5 key-value pairs are accepted, and scope values cannot contain the wildcard character '*'.""" - config: Optional[UpdateAgentEngineMemoryConfigDict] + config: Optional[UpdateMemoryConfigDict] """""" -_UpdateAgentEngineMemoryRequestParametersOrDict = Union[ - _UpdateAgentEngineMemoryRequestParameters, - _UpdateAgentEngineMemoryRequestParametersDict, +_UpdateMemoryRequestParametersOrDict = Union[ + _UpdateMemoryRequestParameters, _UpdateMemoryRequestParametersDict ] -class PurgeAgentEngineMemoriesConfig(_common.BaseModel): +class PurgeMemoriesConfig(_common.BaseModel): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -12226,7 +12414,7 @@ class PurgeAgentEngineMemoriesConfig(_common.BaseModel): ) -class PurgeAgentEngineMemoriesConfigDict(TypedDict, total=False): +class PurgeMemoriesConfigDict(TypedDict, total=False): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] @@ -12236,16 +12424,14 @@ class PurgeAgentEngineMemoriesConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -PurgeAgentEngineMemoriesConfigOrDict = Union[ - PurgeAgentEngineMemoriesConfig, PurgeAgentEngineMemoriesConfigDict -] +PurgeMemoriesConfigOrDict = Union[PurgeMemoriesConfig, PurgeMemoriesConfigDict] -class _PurgeAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for purging agent engine memories.""" +class _PurgeMemoriesRequestParameters(_common.BaseModel): + """Parameters for purging memories.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine to purge memories from.""" + default=None, description="""Name of the Memory Bank to purge memories from.""" ) filter: Optional[str] = Field( default=None, @@ -12273,16 +12459,14 @@ class _PurgeAgentEngineMemoriesRequestParameters(_common.BaseModel): default=None, description="""If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""", ) - config: Optional[PurgeAgentEngineMemoriesConfig] = Field( - default=None, description="""""" - ) + config: Optional[PurgeMemoriesConfig] = Field(default=None, description="""""") -class _PurgeAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for purging agent engine memories.""" +class _PurgeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for purging memories.""" name: Optional[str] - """Name of the Agent Engine to purge memories from.""" + """Name of the Memory Bank to purge memories from.""" filter: Optional[str] """The standard list filter to determine which memories to purge. @@ -12307,13 +12491,12 @@ class _PurgeAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): force: Optional[bool] """If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""" - config: Optional[PurgeAgentEngineMemoriesConfigDict] + config: Optional[PurgeMemoriesConfigDict] """""" -_PurgeAgentEngineMemoriesRequestParametersOrDict = Union[ - _PurgeAgentEngineMemoriesRequestParameters, - _PurgeAgentEngineMemoriesRequestParametersDict, +_PurgeMemoriesRequestParametersOrDict = Union[ + _PurgeMemoriesRequestParameters, _PurgeMemoriesRequestParametersDict ] @@ -12335,8 +12518,8 @@ class PurgeMemoriesResponseDict(TypedDict, total=False): PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] -class AgentEnginePurgeMemoriesOperation(_common.BaseModel): - """Operation that purges memories from an agent engine.""" +class PurgeMemoriesOperation(_common.BaseModel): + """Operation that purges memories from a Memory Bank.""" name: Optional[str] = Field( default=None, @@ -12359,8 +12542,8 @@ class AgentEnginePurgeMemoriesOperation(_common.BaseModel): ) -class AgentEnginePurgeMemoriesOperationDict(TypedDict, total=False): - """Operation that purges memories from an agent engine.""" +class PurgeMemoriesOperationDict(TypedDict, total=False): + """Operation that purges memories from a Memory Bank.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -12378,55 +12561,50 @@ class AgentEnginePurgeMemoriesOperationDict(TypedDict, total=False): """The response for purging memories.""" -AgentEnginePurgeMemoriesOperationOrDict = Union[ - AgentEnginePurgeMemoriesOperation, AgentEnginePurgeMemoriesOperationDict -] +PurgeMemoriesOperationOrDict = Union[PurgeMemoriesOperation, PurgeMemoriesOperationDict] -class GetAgentEngineMemoryRevisionConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory Revision.""" +class GetMemoryRevisionConfig(_common.BaseModel): + """Config for getting a Memory Revision.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineMemoryRevisionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory Revision.""" +class GetMemoryRevisionConfigDict(TypedDict, total=False): + """Config for getting a Memory Revision.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineMemoryRevisionConfigOrDict = Union[ - GetAgentEngineMemoryRevisionConfig, GetAgentEngineMemoryRevisionConfigDict +GetMemoryRevisionConfigOrDict = Union[ + GetMemoryRevisionConfig, GetMemoryRevisionConfigDict ] -class _GetAgentEngineMemoryRevisionRequestParameters(_common.BaseModel): - """Parameters for getting an Agent Engine memory revision.""" +class _GetMemoryRevisionRequestParameters(_common.BaseModel): + """Parameters for getting a memory revision.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[GetAgentEngineMemoryRevisionConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the Memory Revision.""" ) + config: Optional[GetMemoryRevisionConfig] = Field(default=None, description="""""") -class _GetAgentEngineMemoryRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an Agent Engine memory revision.""" +class _GetMemoryRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting a memory revision.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the Memory Revision.""" - config: Optional[GetAgentEngineMemoryRevisionConfigDict] + config: Optional[GetMemoryRevisionConfigDict] """""" -_GetAgentEngineMemoryRevisionRequestParametersOrDict = Union[ - _GetAgentEngineMemoryRevisionRequestParameters, - _GetAgentEngineMemoryRevisionRequestParametersDict, +_GetMemoryRevisionRequestParametersOrDict = Union[ + _GetMemoryRevisionRequestParameters, _GetMemoryRevisionRequestParametersDict ] @@ -12526,8 +12704,8 @@ class MemoryRevisionDict(TypedDict, total=False): MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] -class ListAgentEngineMemoryRevisionsConfig(_common.BaseModel): - """Config for listing Agent Engine memory revisions.""" +class ListMemoryRevisionsConfig(_common.BaseModel): + """Config for listing memory revisions.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -12541,8 +12719,8 @@ class ListAgentEngineMemoryRevisionsConfig(_common.BaseModel): ) -class ListAgentEngineMemoryRevisionsConfigDict(TypedDict, total=False): - """Config for listing Agent Engine memory revisions.""" +class ListMemoryRevisionsConfigDict(TypedDict, total=False): + """Config for listing memory revisions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -12558,40 +12736,37 @@ class ListAgentEngineMemoryRevisionsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineMemoryRevisionsConfigOrDict = Union[ - ListAgentEngineMemoryRevisionsConfig, ListAgentEngineMemoryRevisionsConfigDict +ListMemoryRevisionsConfigOrDict = Union[ + ListMemoryRevisionsConfig, ListMemoryRevisionsConfigDict ] -class _ListAgentEngineMemoryRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing Agent Engine memory revisions.""" +class _ListMemoryRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing memory revisions.""" - name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine memory""" - ) - config: Optional[ListAgentEngineMemoryRevisionsConfig] = Field( + name: Optional[str] = Field(default=None, description="""Name of the memory""") + config: Optional[ListMemoryRevisionsConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineMemoryRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing Agent Engine memory revisions.""" +class _ListMemoryRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing memory revisions.""" name: Optional[str] - """Name of the Agent Engine memory""" + """Name of the memory""" - config: Optional[ListAgentEngineMemoryRevisionsConfigDict] + config: Optional[ListMemoryRevisionsConfigDict] """""" -_ListAgentEngineMemoryRevisionsRequestParametersOrDict = Union[ - _ListAgentEngineMemoryRevisionsRequestParameters, - _ListAgentEngineMemoryRevisionsRequestParametersDict, +_ListMemoryRevisionsRequestParametersOrDict = Union[ + _ListMemoryRevisionsRequestParameters, _ListMemoryRevisionsRequestParametersDict ] -class ListAgentEngineMemoryRevisionsResponse(_common.BaseModel): - """Response for listing agent engine memory revisions.""" +class ListMemoryRevisionsResponse(_common.BaseModel): + """Response for listing memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" @@ -12602,8 +12777,8 @@ class ListAgentEngineMemoryRevisionsResponse(_common.BaseModel): ) -class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): - """Response for listing agent engine memory revisions.""" +class ListMemoryRevisionsResponseDict(TypedDict, total=False): + """Response for listing memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -12615,8 +12790,8 @@ class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): """List of memory revisions.""" -ListAgentEngineMemoryRevisionsResponseOrDict = Union[ - ListAgentEngineMemoryRevisionsResponse, ListAgentEngineMemoryRevisionsResponseDict +ListMemoryRevisionsResponseOrDict = Union[ + ListMemoryRevisionsResponse, ListMemoryRevisionsResponseDict ] diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py deleted file mode 100644 index e4b33ad2d5..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_delete_memory(client): - ae_memory_operation = client.agent_engines.memories.delete( - name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/5605466683931099136", - ) - assert isinstance(ae_memory_operation, types.DeleteAgentEngineMemoryOperation) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories.delete", -) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py deleted file mode 100644 index bca6fca8c9..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_get_memory(client): - memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - ae_memory = client.agent_engines.memories.get(name=memory_name) - assert isinstance(ae_memory, types.Memory) - assert ae_memory.name == memory_name - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories.get", -) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py deleted file mode 100644 index 7b294a01cf..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py +++ /dev/null @@ -1,35 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_private_generate_memory(client): - ae_memory_operation = client.agent_engines.memories._generate( - name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", - vertex_session_source=types.GenerateMemoriesRequestVertexSessionSource( - session="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/sessions/6922431337672474624" - ), - ) - assert isinstance(ae_memory_operation, types.AgentEngineGenerateMemoriesOperation) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories._generate", -) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py deleted file mode 100644 index bdcd2f14a4..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_private_get_generate_memories_operation(client): - memory_operation = client.agent_engines.memories._get_generate_memories_operation( - operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/operations/5669315676343369728" - ) - assert isinstance(memory_operation, types.AgentEngineGenerateMemoriesOperation) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories._get_generate_memories_operation", -) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py deleted file mode 100644 index 85329da47e..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_private_rollback(client): - rollback_operation = client.agent_engines.memories._rollback( - name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688", - target_revision_id="3001207491565453312", - ) - assert isinstance(rollback_operation, types.AgentEngineRollbackMemoryOperation) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories._rollback", -) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py deleted file mode 100644 index f321a921c0..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py +++ /dev/null @@ -1,118 +0,0 @@ -# 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 - -import datetime - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_create_memory_with_ttl(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - - metadata = { - "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), - "my_double_key": types.MemoryMetadataValue(double_value=123.456), - "my_boolean_key": types.MemoryMetadataValue(bool_value=True), - "my_timestamp_key": types.MemoryMetadataValue( - timestamp_value=datetime.datetime( - 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc - ) - ), - } - - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( - display_name="my_memory_fact", - ttl="120s", - metadata=metadata, - ), - ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - assert operation.response.fact == "memory_fact" - assert operation.response.scope == {"user_id": "123"} - assert operation.response.name.startswith(agent_engine.api_resource.name) - # Expire time is calculated by the server, so we only check that it is - # within a reasonable range to avoid flakiness. - assert ( - operation.response.create_time + datetime.timedelta(seconds=119.5) - <= operation.response.expire_time - <= operation.response.create_time + datetime.timedelta(seconds=120.5) - ) - assert operation.response.metadata == metadata - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -def test_create_memory_with_expire_time(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - expire_time = datetime.datetime( - 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc - ) - - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( - display_name="my_memory_fact", expire_time=expire_time - ), - ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - assert operation.response.fact == "memory_fact" - assert operation.response.scope == {"user_id": "123"} - assert operation.response.name.startswith(agent_engine.api_resource.name) - assert operation.response.expire_time == expire_time - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -def test_create_memory_with_custom_memory_id(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( - display_name="my_memory_fact", memory_id="my-memory-id" - ), - ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - assert operation.response.fact == "memory_fact" - assert operation.response.scope == {"user_id": "123"} - assert ( - operation.response.name - == f"{agent_engine.api_resource.name}/memories/my-memory-id" - ) - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.create_memory", -) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py deleted file mode 100644 index 9d3dea4c72..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py +++ /dev/null @@ -1,62 +0,0 @@ -# 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 - -import pytest - - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_delete_memory(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - ) - memory = operation.response - operation = client.agent_engines.memories.delete(name=memory.name) - assert isinstance(operation, types.DeleteAgentEngineMemoryOperation) - assert operation.name.startswith(memory.name + "/operations/") - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.delete_memory", -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_delete_memory_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - ) - memory = operation.response - operation = await client.aio.agent_engines.memories.delete(name=memory.name) - assert isinstance(operation, types.DeleteAgentEngineMemoryOperation) - assert operation.name.startswith(memory.name + "/operations/") - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True - ) diff --git a/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py deleted file mode 100644 index 4942e3ac9e..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py +++ /dev/null @@ -1,381 +0,0 @@ -# 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 - -import datetime -import pytest - - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -from google.genai import types as genai_types - - -def test_generate_and_rollback_memories(client): - agent_engine = client.agent_engines.create() - assert not list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - ) - ) - # Generate memories using source content. This result is non-deterministic, - # because an LLM is used to generate the memories. - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope={"user_id": "test-user-id"}, - direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( - events=[ - types.GenerateMemoriesRequestDirectContentsSourceEvent( - content=genai_types.Content( - role="model", - parts=[ - genai_types.Part( - text="I am a software engineer focusing in security" - ) - ], - ) - ) - ] - ), - config=types.GenerateAgentEngineMemoriesConfig( - revision_labels={"key": "value"} - ), - ) - memories = list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - ) - ) - assert len(memories) >= 1 - - # Every action that modifies a memory creates a new revision. - memory_revisions = list( - client.agent_engines.memories.revisions.list( - name=memories[0].name, - ) - ) - assert len(memory_revisions) >= 1 - # The revision's labels depend on the generation request's revision labels. - assert memory_revisions[0].labels == {"key": "value"} - revision_name = memory_revisions[0].name - - # Update the memory. - client.agent_engines.memories._update( - name=memories[0].name, - fact="This is temporary", - scope={"user_id": "test-user-id"}, - ) - memory = client.agent_engines.memories.get(name=memories[0].name) - assert memory.fact == "This is temporary" - - # Rollback to the revision with the original fact that was created by the - # generation request. - client.agent_engines.memories.rollback( - name=memories[0].name, - target_revision_id=revision_name.split("/")[-1], - ) - memory = client.agent_engines.memories.get(name=memories[0].name) - assert memory.fact == memory_revisions[0].fact - - # Update the memory again using generation. We use the original source - # content to ensure that the original memory is updated. The response should - # refer to the previous revision. - pre_extracted_fact = "I am a software engineer focusing in security" - response = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope={"user_id": "test-user-id"}, - direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( - direct_memories=[ - types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( - fact=pre_extracted_fact - ) - ] - ), - ) - # The memory was updated, so the previous revision is set. - assert response.response.generated_memories[0].previous_revision is not None - memory_revisions = list( - client.agent_engines.memories.revisions.list(name=memories[0].name) - ) - # Memory Revisions are returned in descending order by revision create time. - # We can't make an assertion on the actual value, since it's - # generated and thus non-deterministic. - assert memory_revisions[0].extracted_memories[0].fact is not None - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -def test_generate_memories_direct_memories_source(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope={"user_id": "test-user-id"}, - direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( - direct_memories=[ - types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( - fact="I am a software engineer." - ), - types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( - fact="I like to write replay tests." - ), - ] - ), - config=types.GenerateAgentEngineMemoriesConfig(wait_for_completion=True), - ) - assert ( - len( - list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - ) - ) - ) - >= 1 - ) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -def test_generate_memories_with_metadata(client): - agent_engine = client.agent_engines.create() - metadata = { - "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), - "my_double_key": types.MemoryMetadataValue(double_value=123.456), - "my_boolean_key": types.MemoryMetadataValue(bool_value=True), - "my_timestamp_key": types.MemoryMetadataValue( - timestamp_value=datetime.datetime( - 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc - ) - ), - } - # Reuse the same content and scope for all generation requests to ensure - # that the same memory is updated. - direct_memories_source = types.GenerateMemoriesRequestDirectMemoriesSource( - direct_memories=[ - types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( - fact="I am a software engineer." - ), - ] - ) - scope = {"user_id": "test-user-id"} - - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope=scope, - direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig(metadata=metadata), - ) - assert len(operation.response.generated_memories) >= 1 - memory = client.agent_engines.memories.get( - name=operation.response.generated_memories[0].memory.name - ) - assert memory.metadata == metadata - - # Overwrite the metadata. - overwrite_metadata = { - "my_string_key": types.MemoryMetadataValue(string_value="new_value"), - } - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope=scope, - direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( - metadata=overwrite_metadata, - metadata_merge_strategy=types.MemoryMetadataMergeStrategy.OVERWRITE, - ), - ) - assert len(operation.response.generated_memories) >= 1 - assert ( - operation.response.generated_memories[0].action - == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED - ) - memory = client.agent_engines.memories.get( - name=operation.response.generated_memories[0].memory.name - ) - assert memory.metadata == overwrite_metadata - - # Merge the metadata. - new_metadata = { - "my_double_key": types.MemoryMetadataValue(double_value=123.456), - } - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope=scope, - direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( - metadata=new_metadata, - metadata_merge_strategy=types.MemoryMetadataMergeStrategy.MERGE, - ), - ) - assert len(operation.response.generated_memories) >= 1 - assert ( - operation.response.generated_memories[0].action - == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED - ) - memory = client.agent_engines.memories.get( - name=operation.response.generated_memories[0].memory.name - ) - assert memory.metadata == {**overwrite_metadata, **new_metadata} - - # Restrict consolidation based on metadata values. For the first request, - # there's no existing memories that match the metadata, so a new memory is - # created. - restricted_metadata = { - "my_string_key": types.MemoryMetadataValue(string_value="new_value2"), - } - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope=scope, - direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( - metadata=restricted_metadata, - metadata_merge_strategy="REQUIRE_EXACT_MATCH", - ), - ) - assert len(operation.response.generated_memories) == 1 - # Metadata doesn't match existing memory, so a new memory is created. - assert ( - operation.response.generated_memories[0].action - == types.GenerateMemoriesResponseGeneratedMemoryAction.CREATED - ) - memory = client.agent_engines.memories.get( - name=operation.response.generated_memories[0].memory.name - ) - assert memory.metadata == restricted_metadata - - # Send a second request where the metadata matches only one of the existing - # memories. - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope=scope, - direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( - metadata=restricted_metadata, - metadata_merge_strategy="REQUIRE_EXACT_MATCH", - ), - ) - assert len(operation.response.generated_memories) == 1 - assert ( - operation.response.generated_memories[0].action - == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED - ) - assert operation.response.generated_memories[0].memory.name == memory.name - - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.generate_memories", -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_generate_and_rollback_memories_async(client): - agent_engine = client.agent_engines.create() - await client.aio.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope={"user_id": "test-user-id"}, - direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( - direct_memories=[ - types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( - fact="I am a software engineer." - ), - types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( - fact="I like to write replay tests." - ), - ] - ), - config=types.GenerateAgentEngineMemoriesConfig(wait_for_completion=True), - ) - memories_pager = await client.aio.agent_engines.memories.list( - name=agent_engine.api_resource.name - ) - memory_list = [item async for item in memories_pager] - assert len(memory_list) >= 1 - - revisions_pager = await client.aio.agent_engines.memories.revisions.list( - name=memory_list[0].name - ) - memory_revisions = [item async for item in revisions_pager] - assert len(memory_revisions) >= 1 - revision_name = memory_revisions[0].name - - # Update the memory. - client.agent_engines.memories._update( - name=memory_list[0].name, - fact="This is temporary", - scope={"user_id": "test-user-id"}, - ) - memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name) - assert memory.fact == "This is temporary" - - # Rollback to the revision with the original fact that was created by the - # generation request. - await client.aio.agent_engines.memories.rollback( - name=memory_list[0].name, - target_revision_id=revision_name.split("/")[-1], - ) - memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name) - assert memory.fact == memory_revisions[0].fact - - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True - ) - - -def test_generate_memories_with_allowed_topics(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, - scope={"user_id": "test-user-id"}, - direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( - events=[ - types.GenerateMemoriesRequestDirectContentsSourceEvent( - content=genai_types.Content( - role="user", - parts=[genai_types.Part(text="I am a software engineer.")], - ) - ), - types.GenerateMemoriesRequestDirectContentsSourceEvent( - content=genai_types.Content( - role="user", - parts=[genai_types.Part(text="I like to write replay tests.")], - ) - ), - ] - ), - config=types.GenerateAgentEngineMemoriesConfig( - allowed_topics=[ - types.MemoryTopicId( - managed_memory_topic=types.ManagedTopicEnum.USER_PREFERENCES - ), - ], - wait_for_completion=True, - ), - ) - assert ( - len( - list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - ) - ) - ) - == 1 - ) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py b/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py deleted file mode 100644 index ff7ad46583..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py +++ /dev/null @@ -1,118 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper - - -def test_ingest_events(client): - agent_engine = client.agent_engines.create() - assert not list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - ) - ) - scope = {"user_id": "test-user-id"} - # Generate memories using source content. This result is non-deterministic, - # because an LLM is used to generate the memories. - client.agent_engines.memories.ingest_events( - name=agent_engine.api_resource.name, - scope=scope, - direct_contents_source={ - "events": [ - { - "content": { - "role": "user", - "parts": [{"text": "I like dogs."}], - } - } - ] - }, - # `overlap_event_count` re-includes trailing events from one generation - # window in the next so context is preserved across GenerateMemories - # calls. - generation_trigger_config={ - "generation_rule": {"idle_duration": "60s", "overlap_event_count": 1} - }, - ) - memories = list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope=scope, - ) - ) - # Ingest events should be asynchronous by default, so there should be no - # memories immediately after the call. Processing will only start after 60s - # of inactivity. - assert len(memories) == 0 - - client.agent_engines.memories.ingest_events( - name=agent_engine.api_resource.name, - scope=scope, - direct_contents_source={ - "events": [ - { - "content": { - "role": "user", - "parts": [{"text": "I'm a software engineer."}], - } - } - ] - }, - # `revision_labels`, `metadata`, and `metadata_merge_strategy` are - # applied to the memories generated from the ingested events. Because - # `force_flush` makes generation synchronous, the user-provided metadata - # is observable on the retrieved memory below. - config={ - "wait_for_completion": True, - "force_flush": True, - "revision_labels": {"source": "ingest-events-test"}, - "metadata": {"topic": {"string_value": "jobs"}}, - "metadata_merge_strategy": "OVERWRITE", - }, - ) - memories = list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope=scope, - simple_retrieval_params={ - "page_size": 1, - }, - ) - ) - # 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.agent_engines.memories.revisions.list( - name=memories[0].memory.name, - ) - ) - assert revisions - assert revisions[0].labels == {"source": "ingest-events-test"} - - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories.ingest_events", -) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py deleted file mode 100644 index 4373b4bf25..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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 - -import pytest - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_list_memories(client): - agent_engine = client.agent_engines.create() - assert not list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - ) - ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - config={ - "wait_for_completion": True, - }, - ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_2", - scope={"user_id": "456"}, - config={ - "wait_for_completion": True, - }, - ) - memory_list = client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - config=types.ListAgentEngineMemoryConfig( - page_size=1, - order_by="create_time asc", - ), - ) - assert len(memory_list) == 1 - assert isinstance(memory_list[0], types.Memory) - assert memory_list[0].fact == "memory_fact" - assert memory_list[0].scope["user_id"] == "123" - # Clean up resources. - agent_engine.delete(force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories.list", -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_async_list_memories(client): - agent_engine = client.agent_engines.create() - pager = await client.aio.agent_engines.memories.list( - name=agent_engine.api_resource.name - ) - assert not [item async for item in pager] - - await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_2", - scope={"user_id": "456"}, - config={ - "wait_for_completion": True, - }, - ) - pager = await client.aio.agent_engines.memories.list( - name=agent_engine.api_resource.name - ) - memory_list = [item async for item in pager] - assert len(memory_list) == 1 - assert isinstance(memory_list[0], types.Memory) - - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True - ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_create.py b/tests/unit/agentplatform/genai/replays/test_memories_create.py new file mode 100644 index 0000000000..a820a2f012 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_create.py @@ -0,0 +1,121 @@ +# 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 + +import datetime + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types + + +def test_create_memory_with_ttl(client): + memory_bank = client.memory_banks.create() + try: + assert isinstance(memory_bank, types.MemoryBank) + + metadata = { + "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), + "my_double_key": types.MemoryMetadataValue(double_value=123.456), + "my_boolean_key": types.MemoryMetadataValue(bool_value=True), + "my_timestamp_key": types.MemoryMetadataValue( + timestamp_value=datetime.datetime( + 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc + ) + ), + } + + operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + config=types.MemoryConfig( + display_name="my_memory_fact", + ttl="120s", + metadata=metadata, + ), + ) + assert isinstance(operation, types.MemoryOperation) + assert operation.response.fact == "memory_fact" + assert operation.response.scope == {"user_id": "123"} + assert operation.response.name.startswith(memory_bank.name) + # Expire time is calculated by the server, so we only check that it is + # within a reasonable range to avoid flakiness. + assert ( + operation.response.create_time + datetime.timedelta(seconds=119.5) + <= operation.response.expire_time + <= operation.response.create_time + datetime.timedelta(seconds=120.5) + ) + assert operation.response.metadata == metadata + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank.name, force=True) + + +def test_create_memory_with_expire_time(client): + memory_bank = client.memory_banks.create() + try: + assert isinstance(memory_bank, types.MemoryBank) + expire_time = datetime.datetime( + 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc + ) + + operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + config=types.MemoryConfig( + display_name="my_memory_fact", expire_time=expire_time + ), + ) + assert isinstance(operation, types.MemoryOperation) + assert operation.response.fact == "memory_fact" + assert operation.response.scope == {"user_id": "123"} + assert operation.response.name.startswith(memory_bank.name) + assert operation.response.expire_time == expire_time + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank.name, force=True) + + +def test_create_memory_with_custom_memory_id(client): + memory_bank = client.memory_banks.create() + try: + assert isinstance(memory_bank, types.MemoryBank) + + operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + config=types.MemoryConfig( + display_name="my_memory_fact", memory_id="my-memory-id" + ), + ) + assert isinstance(operation, types.MemoryOperation) + assert operation.response.fact == "memory_fact" + assert operation.response.scope == {"user_id": "123"} + assert ( + operation.response.name + == f"{memory_bank.name}/memories/my-memory-id" + ) + 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.memories.create", +) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_delete.py b/tests/unit/agentplatform/genai/replays/test_memories_delete.py new file mode 100644 index 0000000000..18f31d4991 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_delete.py @@ -0,0 +1,67 @@ +# 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 + +import pytest + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types + + +def test_delete_memory(client): + memory_bank = client.memory_banks.create() + try: + memory_operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + assert isinstance(memory_operation, types.MemoryOperation) + delete_memory_operation = client.memory_banks.memories.delete( + name=memory_operation.response.name, + ) + assert isinstance( + delete_memory_operation, types.DeleteMemoryOperation + ) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.memories.delete", +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_delete_memory_async(client): + memory_bank = client.memory_banks.create() + try: + operation = await client.aio.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + memory = operation.response + operation = await client.aio.memory_banks.memories.delete(name=memory.name) + assert isinstance(operation, types.DeleteMemoryOperation) + assert operation.name.startswith(memory.name + "/operations/") + finally: + await client.aio.memory_banks.delete( + name=memory_bank.name, force=True + ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_generate.py b/tests/unit/agentplatform/genai/replays/test_memories_generate.py new file mode 100644 index 0000000000..5646cedcf5 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_generate.py @@ -0,0 +1,394 @@ +# 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 + +import datetime +import pytest + + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types +from google.genai import types as genai_types + + +def test_generate_and_rollback_memories(client): + memory_bank = client.memory_banks.create() + try: + assert not list( + client.memory_banks.memories.list( + name=memory_bank.name, + ) + ) + # Generate memories using source content. This result is non-deterministic, + # because an LLM is used to generate the memories. + client.memory_banks.memories.generate( + name=memory_bank.name, + scope={"user_id": "test-user-id"}, + direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( + events=[ + types.GenerateMemoriesRequestDirectContentsSourceEvent( + content=genai_types.Content( + role="model", + parts=[ + genai_types.Part( + text="I am a software engineer focusing in security" + ) + ], + ) + ) + ] + ), + config=types.GenerateMemoriesConfig( + revision_labels={"key": "value"} + ), + ) + memories = list( + client.memory_banks.memories.list( + name=memory_bank.name, + ) + ) + assert len(memories) >= 1 + + # Every action that modifies a memory creates a new revision. + memory_revisions = list( + client.memory_banks.memories.revisions.list( + name=memories[0].name, + ) + ) + assert len(memory_revisions) >= 1 + # The revision's labels depend on the generation request's revision labels. + assert memory_revisions[0].labels == {"key": "value"} + revision_name = memory_revisions[0].name + + # Update the memory. + client.memory_banks.memories._update( + name=memories[0].name, + fact="This is temporary", + scope={"user_id": "test-user-id"}, + ) + memory = client.memory_banks.memories.get(name=memories[0].name) + assert memory.fact == "This is temporary" + + # Rollback to the revision with the original fact that was created by the + # generation request. + client.memory_banks.memories.rollback( + name=memories[0].name, + target_revision_id=revision_name.split("/")[-1], + ) + memory = client.memory_banks.memories.get(name=memories[0].name) + assert memory.fact == memory_revisions[0].fact + + # Update the memory again using generation. We use the original source + # content to ensure that the original memory is updated. The response should + # refer to the previous revision. + pre_extracted_fact = "I am a software engineer focusing in security" + response = client.memory_banks.memories.generate( + name=memory_bank.name, + scope={"user_id": "test-user-id"}, + direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( + direct_memories=[ + types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( + fact=pre_extracted_fact + ) + ] + ), + ) + # The memory was updated, so the previous revision is set. + assert response.response.generated_memories[0].previous_revision is not None + memory_revisions = list( + client.memory_banks.memories.revisions.list(name=memories[0].name) + ) + # Memory Revisions are returned in descending order by revision create time. + # We can't make an assertion on the actual value, since it's + # generated and thus non-deterministic. + assert memory_revisions[0].extracted_memories[0].fact is not None + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) + + +def test_direct_memories_source(client): + memory_bank = client.memory_banks.create() + try: + client.memory_banks.memories.generate( + name=memory_bank.name, + scope={"user_id": "test-user-id"}, + direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( + direct_memories=[ + types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( + fact="I am a software engineer." + ), + types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( + fact="I like to write replay tests." + ), + ] + ), + config=types.GenerateMemoriesConfig(wait_for_completion=True), + ) + assert ( + len( + list( + client.memory_banks.memories.list( + name=memory_bank.name, + ) + ) + ) + >= 1 + ) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) + + +def test_generate_with_metadata(client): + memory_bank = client.memory_banks.create() + try: + metadata = { + "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), + "my_double_key": types.MemoryMetadataValue(double_value=123.456), + "my_boolean_key": types.MemoryMetadataValue(bool_value=True), + "my_timestamp_key": types.MemoryMetadataValue( + timestamp_value=datetime.datetime( + 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc + ) + ), + } + # Reuse the same content and scope for all generation requests to ensure + # that the same memory is updated. + direct_memories_source = types.GenerateMemoriesRequestDirectMemoriesSource( + direct_memories=[ + types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( + fact="I am a software engineer." + ), + ] + ) + scope = {"user_id": "test-user-id"} + + operation = client.memory_banks.memories.generate( + name=memory_bank.name, + scope=scope, + direct_memories_source=direct_memories_source, + config=types.GenerateMemoriesConfig(metadata=metadata), + ) + assert len(operation.response.generated_memories) >= 1 + memory = client.memory_banks.memories.get( + name=operation.response.generated_memories[0].memory.name + ) + assert memory.metadata == metadata + + # Overwrite the metadata. + overwrite_metadata = { + "my_string_key": types.MemoryMetadataValue(string_value="new_value"), + } + operation = client.memory_banks.memories.generate( + name=memory_bank.name, + scope=scope, + direct_memories_source=direct_memories_source, + config=types.GenerateMemoriesConfig( + metadata=overwrite_metadata, + metadata_merge_strategy=types.MemoryMetadataMergeStrategy.OVERWRITE, + ), + ) + assert len(operation.response.generated_memories) >= 1 + assert ( + operation.response.generated_memories[0].action + == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED + ) + memory = client.memory_banks.memories.get( + name=operation.response.generated_memories[0].memory.name + ) + assert memory.metadata == overwrite_metadata + + # Merge the metadata. + new_metadata = { + "my_double_key": types.MemoryMetadataValue(double_value=123.456), + } + operation = client.memory_banks.memories.generate( + name=memory_bank.name, + scope=scope, + direct_memories_source=direct_memories_source, + config=types.GenerateMemoriesConfig( + metadata=new_metadata, + metadata_merge_strategy=types.MemoryMetadataMergeStrategy.MERGE, + ), + ) + assert len(operation.response.generated_memories) >= 1 + assert ( + operation.response.generated_memories[0].action + == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED + ) + memory = client.memory_banks.memories.get( + name=operation.response.generated_memories[0].memory.name + ) + assert memory.metadata == {**overwrite_metadata, **new_metadata} + + # Restrict consolidation based on metadata values. For the first request, + # there's no existing memories that match the metadata, so a new memory is + # created. + restricted_metadata = { + "my_string_key": types.MemoryMetadataValue(string_value="new_value2"), + } + operation = client.memory_banks.memories.generate( + name=memory_bank.name, + scope=scope, + direct_memories_source=direct_memories_source, + config=types.GenerateMemoriesConfig( + metadata=restricted_metadata, + metadata_merge_strategy="REQUIRE_EXACT_MATCH", + ), + ) + assert len(operation.response.generated_memories) == 1 + # Metadata doesn't match existing memory, so a new memory is created. + assert ( + operation.response.generated_memories[0].action + == types.GenerateMemoriesResponseGeneratedMemoryAction.CREATED + ) + memory = client.memory_banks.memories.get( + name=operation.response.generated_memories[0].memory.name + ) + assert memory.metadata == restricted_metadata + + # Send a second request where the metadata matches only one of the existing + # memories. + operation = client.memory_banks.memories.generate( + name=memory_bank.name, + scope=scope, + direct_memories_source=direct_memories_source, + config=types.GenerateMemoriesConfig( + metadata=restricted_metadata, + metadata_merge_strategy="REQUIRE_EXACT_MATCH", + ), + ) + assert len(operation.response.generated_memories) == 1 + assert ( + operation.response.generated_memories[0].action + == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED + ) + assert operation.response.generated_memories[0].memory.name == memory.name + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.memories.generate", +) + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_generate_and_rollback_memories_async(client): + memory_bank = client.memory_banks.create() + try: + await client.aio.memory_banks.memories.generate( + name=memory_bank.name, + scope={"user_id": "test-user-id"}, + direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( + direct_memories=[ + types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( + fact="I am a software engineer." + ), + types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory( + fact="I like to write replay tests." + ), + ] + ), + config=types.GenerateMemoriesConfig(wait_for_completion=True), + ) + memories_pager = await client.aio.memory_banks.memories.list( + name=memory_bank.name + ) + memory_list = [item async for item in memories_pager] + assert len(memory_list) >= 1 + + revisions_pager = await client.aio.memory_banks.memories.revisions.list( + name=memory_list[0].name + ) + memory_revisions = [item async for item in revisions_pager] + assert len(memory_revisions) >= 1 + revision_name = memory_revisions[0].name + + # Update the memory. + client.memory_banks.memories._update( + name=memory_list[0].name, + fact="This is temporary", + scope={"user_id": "test-user-id"}, + ) + memory = await client.aio.memory_banks.memories.get(name=memory_list[0].name) + assert memory.fact == "This is temporary" + + # Rollback to the revision with the original fact that was created by the + # generation request. + await client.aio.memory_banks.memories.rollback( + name=memory_list[0].name, + target_revision_id=revision_name.split("/")[-1], + ) + memory = await client.aio.memory_banks.memories.get(name=memory_list[0].name) + assert memory.fact == memory_revisions[0].fact + + finally: + await client.aio.memory_banks.delete( + name=memory_bank.name, force=True + ) + + +def test_generate_with_allowed_topics(client): + memory_bank = client.memory_banks.create() + try: + client.memory_banks.memories.generate( + name=memory_bank.name, + scope={"user_id": "test-user-id"}, + direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( + events=[ + types.GenerateMemoriesRequestDirectContentsSourceEvent( + content=genai_types.Content( + role="user", + parts=[genai_types.Part(text="I am a software engineer.")], + ) + ), + types.GenerateMemoriesRequestDirectContentsSourceEvent( + content=genai_types.Content( + role="user", + parts=[ + genai_types.Part( + text="I like to write replay tests." + ) + ], + ) + ), + ] + ), + config=types.GenerateMemoriesConfig( + allowed_topics=[ + types.MemoryTopicId( + managed_memory_topic=types.ManagedTopicEnum.USER_PREFERENCES + ), + ], + wait_for_completion=True, + ), + ) + assert ( + len( + list( + client.memory_banks.memories.list( + name=memory_bank.name, + ) + ) + ) + == 1 + ) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_memories_get.py similarity index 55% rename from tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py rename to tests/unit/agentplatform/genai/replays/test_memories_get.py index 3ec3093562..1d523c88c7 100644 --- a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_get.py @@ -16,30 +16,32 @@ import pytest -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_get_memory(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - memory = client.agent_engines.memories.get( - name=operation.response.name, - ) - assert isinstance(memory, types.Memory) - assert memory.name == operation.response.name - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + memory_bank = client.memory_banks.create() + try: + operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + assert isinstance(operation, types.MemoryOperation) + memory = client.memory_banks.memories.get( + name=operation.response.name, + ) + assert isinstance(memory, types.Memory) + assert memory.name == operation.response.name + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.get", + test_method="memory_banks.memories.get", ) @@ -48,18 +50,18 @@ def test_get_memory(client): @pytest.mark.asyncio async def test_get_memory_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, + memory_bank = client.memory_banks.create() + operation = await client.aio.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - memory = await client.aio.agent_engines.memories.get( + assert isinstance(operation, types.MemoryOperation) + memory = await client.aio.memory_banks.memories.get( name=operation.response.name, ) assert isinstance(memory, types.Memory) assert memory.name == operation.response.name - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.memory_banks.delete( + name=memory_bank.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_list.py b/tests/unit/agentplatform/genai/replays/test_memories_list.py new file mode 100644 index 0000000000..54f7e2b825 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_list.py @@ -0,0 +1,99 @@ +# 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 + +import pytest + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types + + +def test_list_memories(client): + memory_bank = client.memory_banks.create() + try: + assert not list( + client.memory_banks.memories.list( + name=memory_bank.name, + ) + ) + client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + config={ + "wait_for_completion": True, + }, + ) + client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_2", + scope={"user_id": "456"}, + config={ + "wait_for_completion": True, + }, + ) + memory_list = client.memory_banks.memories.list( + name=memory_bank.name, + config=types.ListMemoriesConfig( + page_size=1, + order_by="create_time asc", + ), + ) + assert len(memory_list) == 1 + assert isinstance(memory_list[0], types.Memory) + assert memory_list[0].fact == "memory_fact" + assert memory_list[0].scope["user_id"] == "123" + 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.memories.list", +) + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_async_list_memories(client): + memory_bank = client.memory_banks.create() + try: + pager = await client.aio.memory_banks.memories.list( + name=memory_bank.name, + ) + assert not [item async for item in pager] + + await client.aio.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_2", + scope={"user_id": "456"}, + config={ + "wait_for_completion": True, + }, + ) + pager = await client.aio.memory_banks.memories.list( + name=memory_bank.name + ) + memory_list = [item async for item in pager] + assert len(memory_list) == 1 + assert isinstance(memory_list[0], types.Memory) + finally: + await client.aio.memory_banks.delete( + name=memory_bank.name, force=True + ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py b/tests/unit/agentplatform/genai/replays/test_memories_private_create.py similarity index 70% rename from tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py rename to tests/unit/agentplatform/genai/replays/test_memories_private_create.py index 5ca959abe2..ed2904d9e4 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_create.py @@ -19,12 +19,16 @@ def test_private_create_memory(client): - ae_memory_operation = client.agent_engines.memories._create( - name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", - fact="memory_fact", - scope={"user_id": "123"}, - ) - assert isinstance(ae_memory_operation, types.AgentEngineMemoryOperation) + memory_bank = client.memory_banks.create() + try: + memory_operation = client.memory_banks.memories._create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + assert isinstance(memory_operation, types.MemoryOperation) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( diff --git a/tests/unit/agentplatform/genai/replays/test_memories_private_generate.py b/tests/unit/agentplatform/genai/replays/test_memories_private_generate.py new file mode 100644 index 0000000000..1e30f82466 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_generate.py @@ -0,0 +1,56 @@ +# 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 + +import time + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types +from google.genai import types as genai_types + + +def test_private_generate_memory(client): + memory_bank = client.memory_banks.create() + try: + generate_memories_operation = client.memory_banks.memories._generate( + name=memory_bank.name, + scope={"user_id": "123"}, + direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( + events=[ + types.GenerateMemoriesRequestDirectContentsSourceEvent( + content=genai_types.Content( + role="model", + parts=[ + genai_types.Part( + text="I am a software engineer focusing in security" + ) + ], + ) + ) + ] + ), + ) + assert isinstance(generate_memories_operation, types.GenerateMemoriesOperation) + finally: + # Wait for the operation to complete before trying to delete. + time.sleep(10) + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.memories._generate", +) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_private_get_generate_memories_operation.py b/tests/unit/agentplatform/genai/replays/test_memories_private_get_generate_memories_operation.py new file mode 100644 index 0000000000..4c642b4203 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_get_generate_memories_operation.py @@ -0,0 +1,55 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types +from google.genai import types as genai_types + + +def test_private_get_generate_memories_operation(client): + memory_bank = client.memory_banks.create() + try: + generate_memories_operation = client.memory_banks.memories.generate( + name=memory_bank.name, + scope={"user_id": "123"}, + direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( + events=[ + types.GenerateMemoriesRequestDirectContentsSourceEvent( + content=genai_types.Content( + role="model", + parts=[ + genai_types.Part( + text="I am writing tests." + ) + ], + ) + ) + ] + ), + ) + memory_operation = client.memory_banks.memories._get_generate_memories_operation( + operation_name=generate_memories_operation.name + ) + assert isinstance(memory_operation, types.GenerateMemoriesOperation) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.memories._get_generate_memories_operation", +) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py b/tests/unit/agentplatform/genai/replays/test_memories_private_get_memory_operation.py similarity index 58% rename from tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py rename to tests/unit/agentplatform/genai/replays/test_memories_private_get_memory_operation.py index 23c92f7821..e31342b88f 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_get_memory_operation.py @@ -19,14 +19,24 @@ def test_private_get_memory_operation(client): - memory_operation = client.agent_engines.memories._get_memory_operation( - operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/operations/1044963283964002304" - ) - assert isinstance(memory_operation, types.AgentEngineMemoryOperation) + memory_bank = client.memory_banks.create() + try: + memory_operation = client.memory_banks.memories._create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + operation_name = memory_operation.response.name + memory_operation = client.memory_banks.memories._get_memory_operation( + operation_name=operation_name + ) + assert isinstance(memory_operation, types.MemoryOperation) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._get_memory_operation", + test_method="memory_banks.memories._get_memory_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py b/tests/unit/agentplatform/genai/replays/test_memories_private_list.py similarity index 61% rename from tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py rename to tests/unit/agentplatform/genai/replays/test_memories_private_list.py index 662ad020fb..3837f25813 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_list.py @@ -19,14 +19,22 @@ def test_private_list_memory(client): - ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" - memory_list = client.agent_engines.memories._list(name=ae_name) - assert isinstance(memory_list, types.ListReasoningEnginesMemoriesResponse) - assert isinstance(memory_list.memories[0], types.Memory) + memory_bank = client.memory_banks.create() + try: + _ = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + memory_list = client.memory_banks.memories._list(name=memory_bank.name) + assert isinstance(memory_list, types.ListMemoriesResponse) + assert isinstance(memory_list.memories[0], types.Memory) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._list", + test_method="memory_banks.memories._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py b/tests/unit/agentplatform/genai/replays/test_memories_private_purge.py similarity index 66% rename from tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py rename to tests/unit/agentplatform/genai/replays/test_memories_private_purge.py index d77cb3a795..92cca4b615 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_purge.py @@ -19,17 +19,21 @@ def test_private_purge(client): - ae_memory_purge_operation = client.agent_engines.memories._purge( - name="projects/964831358985/locations/us-central1/reasoningEngines/6086402690647064576", - filter="scope.user_id=123", - ) - assert isinstance( - ae_memory_purge_operation, types.AgentEnginePurgeMemoriesOperation - ) + memory_bank = client.memory_banks.create() + try: + memory_purge_operation = client.memory_banks.memories._purge( + name=memory_bank.name, + filter="scope.user_id=123", + ) + assert isinstance( + memory_purge_operation, types.PurgeMemoriesOperation + ) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._purge", + test_method="memory_banks.memories._purge", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py b/tests/unit/agentplatform/genai/replays/test_memories_private_retrieve.py similarity index 54% rename from tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py rename to tests/unit/agentplatform/genai/replays/test_memories_private_retrieve.py index 4098b11d9b..d0c72671b0 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_retrieve.py @@ -19,20 +19,28 @@ def test_private_retrieve(client): - ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" - retrieved_memories = client.agent_engines.memories._retrieve( - name=ae_name, - scope={"user_id": "123"}, - ) - assert isinstance(retrieved_memories, types.RetrieveMemoriesResponse) - assert isinstance( - retrieved_memories.retrieved_memories[0], - types.RetrieveMemoriesResponseRetrievedMemory, - ) + memory_bank = client.memory_banks.create() + try: + _ = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + retrieved_memories = client.memory_banks.memories._retrieve( + name=memory_bank.name, + scope={"user_id": "123"}, + ) + assert isinstance(retrieved_memories, types.RetrieveMemoriesResponse) + assert isinstance( + retrieved_memories.retrieved_memories[0], + types.RetrieveMemoriesResponseRetrievedMemory, + ) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._retrieve", + test_method="memory_banks.memories._retrieve", ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_private_rollback.py b/tests/unit/agentplatform/genai/replays/test_memories_private_rollback.py new file mode 100644 index 0000000000..ddf157e6e3 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_rollback.py @@ -0,0 +1,52 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types + + +def test_private_rollback(client): + memory_bank = client.memory_banks.create() + try: + memory_operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + memory_name = memory_operation.response.name + client.memory_banks.memories._update( + name=memory_name, + fact="memory_fact_updated", + scope={"user_id": "123"}, + ) + memory_revision_list = client.memory_banks.memories.revisions.list( + name=memory_name + ) + target_revision_name = memory_revision_list[0].name + rollback_operation = client.memory_banks.memories._rollback( + name=memory_name, + target_revision_id=target_revision_name.split("/")[-1], + ) + assert isinstance(rollback_operation, types.RollbackMemoryOperation) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.memories._rollback", +) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py b/tests/unit/agentplatform/genai/replays/test_memories_private_update.py similarity index 56% rename from tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py rename to tests/unit/agentplatform/genai/replays/test_memories_private_update.py index c6dbc1f095..69cdbe5643 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_update.py @@ -19,17 +19,26 @@ def test_private_update_memory(client): - memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - memory_update_operation = client.agent_engines.memories._update( - name=memory_name, - fact="memory_fact_updated", - scope={"user_id": "123"}, - ) - assert isinstance(memory_update_operation, types.AgentEngineMemoryOperation) + memory_bank = client.memory_banks.create() + try: + memory_operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + memory_name = memory_operation.response.name + memory_update_operation = client.memory_banks.memories._update( + name=memory_name, + fact="memory_fact_updated", + scope={"user_id": "123"}, + ) + assert isinstance(memory_update_operation, types.MemoryOperation) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._update", + test_method="memory_banks.memories._update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_memories_purge.py similarity index 66% rename from tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py rename to tests/unit/agentplatform/genai/replays/test_memories_purge.py index e123597953..e2ac58dde9 100644 --- a/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_purge.py @@ -22,28 +22,28 @@ def test_purge_memories(client): """Tests purging memories.""" - agent_engine = client.agent_engines.create() + memory_bank = client.memory_banks.create() try: - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact_4", scope={"user_id": "456"}, config={ @@ -51,8 +51,8 @@ def test_purge_memories(client): "metadata": {"my_key": {"string_value": "my_value"}}, }, ) - operation = client.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = client.memory_banks.memories.purge( + name=memory_bank.name, filter="scope.user_id=123", config={"wait_for_completion": True}, ) @@ -62,16 +62,16 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.memory_banks.memories.list( + name=memory_bank.name ) ) ) == 4 ) # Now, actually purge the memories. - operation = client.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = client.memory_banks.memories.purge( + name=memory_bank.name, filter="scope.user_id=123", force=True, config={"wait_for_completion": True}, @@ -81,16 +81,16 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.memory_banks.memories.list( + name=memory_bank.name ) ) ) == 2 ) # Purge memories using filter groups. - operation = client.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = client.memory_banks.memories.purge( + name=memory_bank.name, force=True, filter_groups=[ {"filters": [{"key": "my_key", "value": {"string_value": "my_value"}}]} @@ -104,21 +104,21 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.memory_banks.memories.list( + name=memory_bank.name ) ) ) == 1 ) finally: - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.purge", + test_method="memory_banks.memories.purge", ) @@ -127,29 +127,29 @@ def test_purge_memories(client): @pytest.mark.asyncio async def test_purge_memories_async(client): - agent_engine = client.agent_engines.create() + memory_bank = client.memory_banks.create() try: - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.memory_banks.memories.create( + name=memory_bank.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - operation = await client.aio.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = await client.aio.memory_banks.memories.purge( + name=memory_bank.name, filter="scope.user_id=123", config={"wait_for_completion": True}, ) @@ -159,16 +159,16 @@ async def test_purge_memories_async(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.memory_banks.memories.list( + name=memory_bank.name ) ) ) == 3 ) # Now, actually purge the memories. - operation = await client.aio.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = await client.aio.memory_banks.memories.purge( + name=memory_bank.name, filter="scope.user_id=123", force=True, config={"wait_for_completion": True}, @@ -178,12 +178,12 @@ async def test_purge_memories_async(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.memory_banks.memories.list( + name=memory_bank.name ) ) ) == 1 ) finally: - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.memory_banks.delete(name=memory_bank.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_retrieve.py b/tests/unit/agentplatform/genai/replays/test_memories_retrieve.py new file mode 100644 index 0000000000..056115999d --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memories_retrieve.py @@ -0,0 +1,205 @@ +# 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 + +import datetime +import pytest + + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types +from google.genai import pagers + + +def test_retrieve_memories_with_similarity_search_params(client): + memory_bank = client.memory_banks.create() + try: + assert not list( + client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope={"user_id": "123"}, + similarity_search_params=types.RetrieveMemoriesRequestSimilaritySearchParams( + search_query="memory_fact_1", + ), + ) + ) + client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_1", + scope={"user_id": "123"}, + ) + assert ( + len( + list( + client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope={"user_id": "123"}, + ) + ) + ) + == 1 + ) + assert not list( + client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope={"user_id": "456"}, + ) + ) + client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_2", + scope={"user_id": "123"}, + ) + assert ( + len( + list( + client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope={"user_id": "123"}, + ) + ) + ) + == 2 + ) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) + + +def test_retrieve_memories_with_simple_retrieval_params(client): + memory_bank = client.memory_banks.create() + try: + client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_1", + scope={"user_id": "123"}, + ) + memories = client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope={"user_id": "123"}, + simple_retrieval_params=types.RetrieveMemoriesRequestSimpleRetrievalParams( + page_size=1, + ), + ) + assert isinstance(memories, pagers.Pager) + assert isinstance( + memories.page[0], types.RetrieveMemoriesResponseRetrievedMemory + ) + assert memories.page_size == 1 + + client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_2", + scope={"user_id": "123"}, + ) + memories = client.memory_banks.memories.retrieve( + name=memory_bank.name, scope={"user_id": "123"} + ) + assert memories.page_size == 2 + + memories = client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope={"user_id": "123"}, + config={"filter": 'fact="memory_fact_2"'}, + ) + assert memories.page_size == 1 + assert memories.page[0].memory.fact == "memory_fact_2" + + finally: + # Clean up resources. + client.memory_banks.delete(name=memory_bank.name, force=True) + + +def test_retrieve_memories_with_metadata(client): + memory_bank = client.memory_banks.create() + try: + metadata = { + "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), + "my_double_key": types.MemoryMetadataValue(double_value=123.456), + "my_boolean_key": types.MemoryMetadataValue(bool_value=True), + "my_timestamp_key": types.MemoryMetadataValue( + timestamp_value=datetime.datetime( + 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc + ) + ), + } + scope = {"user_id": "123"} + client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_1", + scope=scope, + ) + operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact_2", + scope=scope, + config={"metadata": metadata}, + ) + memory_name2 = operation.response.name + + results = client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope=scope, + config={ + "filter_groups": [ + { + "filters": [ + { + "key": "my_string_key", + "value": {"string_value": "my_string_value"}, + } + ] + } + ], + }, + ) + assert len(results) == 1 + assert results[0].memory.name == memory_name2 + + 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.memories.retrieve", +) + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_retrieve_memories_async(client): + memory_bank = client.memory_banks.create() + try: + operation = await client.aio.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + assert isinstance(operation, types.MemoryOperation) + pager = await client.aio.memory_banks.memories.retrieve( + name=memory_bank.name, + scope={"user_id": "123"}, + ) + memories = [item async for item in pager] + assert len(memories) == 1 + assert isinstance(memories[0], types.RetrieveMemoriesResponseRetrievedMemory) + finally: + await client.aio.memory_banks.delete( + name=memory_bank.name, force=True + ) diff --git a/tests/unit/agentplatform/genai/replays/test_structured_memories.py b/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py similarity index 76% rename from tests/unit/agentplatform/genai/replays/test_structured_memories.py rename to tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py index 5d1c2a39a5..ac9cc30941 100644 --- a/tests/unit/agentplatform/genai/replays/test_structured_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py @@ -19,10 +19,7 @@ def test_generate_and_retrieve_profile(client): - # TODO: Use prod once available. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com" - ) + # 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 @@ -44,7 +41,7 @@ def test_generate_and_retrieve_profile(client): structured_memory_config_obj = types.StructuredMemoryConfig( **structured_memory_config ) - agent_engine = client.agent_engines.create( + memory_bank = client.agent_engines.create( config={ "context_spec": { "memory_bank_config": { @@ -56,8 +53,8 @@ def test_generate_and_retrieve_profile(client): }, ) try: - agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name) - memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config + 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_customization_config ] @@ -66,16 +63,16 @@ def test_generate_and_retrieve_profile(client): ] scope = {"user_id": "123"} - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + client.memory_banks.memories.generate( + name=memory_bank.api_resource.name, scope=scope, direct_contents_source={ "events": [{"content": {"parts": [{"text": "My name is Kim."}]}}] }, ) memories = list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.memory_banks.memories.retrieve( + name=memory_bank.api_resource.name, scope=scope, config={"memory_types": ["STRUCTURED_PROFILE"]}, ) @@ -83,18 +80,18 @@ def test_generate_and_retrieve_profile(client): assert len(memories) >= 1 assert memories[0].memory.structured_content is not None - response = client.agent_engines.memories.retrieve_profiles( - name=agent_engine.api_resource.name, scope=scope + response = client.memory_banks.memories.retrieve_profiles( + name=memory_bank.api_resource.name, scope=scope ) assert len(response.profiles) == 1 finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.memory_banks.delete(name=memory_bank.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.retrieve_profiles", + test_method="memory_banks.memories.retrieve_profiles", ) 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 new file mode 100644 index 0000000000..fc74f5ecc3 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_memory_banks_ingest_events.py @@ -0,0 +1,123 @@ +# 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 + +import time + +from tests.unit.agentplatform.genai.replays import pytest_helper + + +def test_ingest_events(client): + memory_bank = client.memory_banks.create() + try: + assert not list( + client.memory_banks.memories.list( + name=memory_bank.name, + ) + ) + scope = {"user_id": "test-user-id"} + # Generate memories using source content. This result is non-deterministic, + # because an LLM is used to generate the memories. + client.memory_banks.ingest_events( + name=memory_bank.name, + scope=scope, + direct_contents_source={ + "events": [ + { + "content": { + "role": "user", + "parts": [{"text": "I like dogs."}], + } + } + ] + }, + # `overlap_event_count` re-includes trailing events from one generation + # window in the next so context is preserved across GenerateMemories + # calls. + generation_trigger_config={ + "generation_rule": {"idle_duration": "60s", "overlap_event_count": 1} + }, + ) + memories = list( + client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope=scope, + ) + ) + # Ingest events should be asynchronous by default, so there should be no + # memories immediately after the call. Processing will only start after 60s + # of inactivity. + assert len(memories) == 0 + + result = client.memory_banks.ingest_events( + name=memory_bank.name, + scope=scope, + direct_contents_source={ + "events": [ + { + "content": { + "role": "user", + "parts": [{"text": "I'm a software engineer!"}], + } + } + ] + }, + # `revision_labels`, `metadata`, and `metadata_merge_strategy` are + # applied to the memories generated from the ingested events. + # Because `force_flush` makes generation synchronous, the + # user-provided metadata is observable on the retrieved memory below. + config={ + "wait_for_completion": True, + "force_flush": True, + "revision_labels": {"source": "ingest-events-test"}, + "metadata": {"topic": {"string_value": "jobs"}}, + "metadata_merge_strategy": "OVERWRITE", + }, + ) + memories = client.memory_banks.memories.retrieve( + name=memory_bank.name, + scope=scope, + simple_retrieval_params={ + "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" + + # # 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) + client.memory_banks.delete(name=memory_bank.name, force=True) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="memory_banks.ingest_events", +) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py b/tests/unit/agentplatform/genai/replays/test_memory_revisions_get.py similarity index 53% rename from tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py rename to tests/unit/agentplatform/genai/replays/test_memory_revisions_get.py index f58e365fe7..0058298362 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py +++ b/tests/unit/agentplatform/genai/replays/test_memory_revisions_get.py @@ -19,14 +19,27 @@ def test_get_memory_revisions(client): - memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/revisions/516064922187071488" - memory_revision = client.agent_engines.memories.revisions.get(name=memory_name) - assert isinstance(memory_revision, types.MemoryRevision) - assert memory_revision.name == memory_name + memory_bank = client.memory_banks.create() + try: + operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + memory_revisions = list(client.memory_banks.memories.revisions.list( + name=operation.response.name, + )) + memory_revision_name = memory_revisions[0].name + memory_revision = client.memory_banks.memories.revisions.get( + name=memory_revision_name) + assert isinstance(memory_revision, types.MemoryRevision) + assert memory_revision.name == memory_revision_name + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.revisions.get", + test_method="memory_banks.memories.revisions.get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py b/tests/unit/agentplatform/genai/replays/test_memory_revisions_private_list.py similarity index 58% rename from tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py rename to tests/unit/agentplatform/genai/replays/test_memory_revisions_private_list.py index c827a57a06..7ade5968a0 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_memory_revisions_private_list.py @@ -19,14 +19,23 @@ def test_private_list_memory_revisions(client): - ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - memory_list = client.agent_engines.memories.revisions._list(name=ae_name) - assert isinstance(memory_list, types.ListAgentEngineMemoryRevisionsResponse) - assert isinstance(memory_list.memory_revisions[0], types.MemoryRevision) + memory_bank = client.memory_banks.create() + try: + memory_operation = client.memory_banks.memories.create( + name=memory_bank.name, + fact="memory_fact", + scope={"user_id": "123"}, + ) + memory_list = client.memory_banks.memories.revisions._list( + name=memory_operation.response.name) + assert isinstance(memory_list, types.ListMemoryRevisionsResponse) + assert isinstance(memory_list.memory_revisions[0], types.MemoryRevision) + finally: + client.memory_banks.delete(name=memory_bank.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.revisions._list", + test_method="memory_banks.memories.revisions._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py deleted file mode 100644 index 02a3e034d0..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py +++ /dev/null @@ -1,196 +0,0 @@ -# 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 - -import datetime -import pytest - - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -from google.genai import pagers - - -def test_retrieve_memories_with_similarity_search_params(client): - agent_engine = client.agent_engines.create() - assert not list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope={"user_id": "123"}, - similarity_search_params=types.RetrieveMemoriesRequestSimilaritySearchParams( - search_query="memory_fact_1", - ), - ) - ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_1", - scope={"user_id": "123"}, - ) - assert ( - len( - list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope={"user_id": "123"}, - ) - ) - ) - == 1 - ) - assert not list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope={"user_id": "456"}, - ) - ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_2", - scope={"user_id": "123"}, - ) - assert ( - len( - list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope={"user_id": "123"}, - ) - ) - ) - == 2 - ) - # Clean up resources. - agent_engine.delete(force=True) - - -def test_retrieve_memories_with_simple_retrieval_params(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_1", - scope={"user_id": "123"}, - ) - memories = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope={"user_id": "123"}, - simple_retrieval_params=types.RetrieveMemoriesRequestSimpleRetrievalParams( - page_size=1, - ), - ) - assert isinstance(memories, pagers.Pager) - assert isinstance(memories.page[0], types.RetrieveMemoriesResponseRetrievedMemory) - assert memories.page_size == 1 - - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_2", - scope={"user_id": "123"}, - ) - memories = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, scope={"user_id": "123"} - ) - assert memories.page_size == 2 - - memories = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope={"user_id": "123"}, - config={"filter": 'fact="memory_fact_2"'}, - ) - assert memories.page_size == 1 - assert memories.page[0].memory.fact == "memory_fact_2" - - # Clean up resources. - agent_engine.delete(force=True) - - -def test_retrieve_memories_with_metadata(client): - agent_engine = client.agent_engines.create() - metadata = { - "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), - "my_double_key": types.MemoryMetadataValue(double_value=123.456), - "my_boolean_key": types.MemoryMetadataValue(bool_value=True), - "my_timestamp_key": types.MemoryMetadataValue( - timestamp_value=datetime.datetime( - 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc - ) - ), - } - scope = {"user_id": "123"} - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_1", - scope=scope, - ) - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact_2", - scope=scope, - config={"metadata": metadata}, - ) - memory_name2 = operation.response.name - - results = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope=scope, - config={ - "filter_groups": [ - { - "filters": [ - { - "key": "my_string_key", - "value": {"string_value": "my_string_value"}, - } - ] - } - ], - }, - ) - assert len(results) == 1 - assert results[0].memory.name == memory_name2 - - # Clean up resources. - agent_engine.delete(force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.memories.retrieve", -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_retrieve_memories_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, - fact="memory_fact", - scope={"user_id": "123"}, - ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - pager = await client.aio.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, - scope={"user_id": "123"}, - ) - memories = [item async for item in pager] - assert len(memories) == 1 - assert isinstance(memories[0], types.RetrieveMemoriesResponseRetrievedMemory) - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True - )