diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index 6cb42423c4..d6b1820371 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -441,7 +441,11 @@ class _RequirementsValidationResult(TypedDict): actions: _RequirementsValidationActions -AgentEngineOperationUnion = Union[genai_types.AgentEngineOperation] +AgentEngineOperationUnion = Union[ + genai_types.AgentEngineOperation, + genai_types.AgentEngineMemoryOperation, + genai_types.AgentEngineGenerateMemoriesOperation, +] class GetOperationFunction(Protocol): diff --git a/agentplatform/_genai/_memory_bank_utils.py b/agentplatform/_genai/_memory_bank_utils.py deleted file mode 100644 index 694fa4e3cf..0000000000 --- a/agentplatform/_genai/_memory_bank_utils.py +++ /dev/null @@ -1,142 +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. -# -"""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 2fd26fda80..ef4d6f2a09 100644 --- a/agentplatform/_genai/agent_engines.py +++ b/agentplatform/_genai/agent_engines.py @@ -37,10 +37,12 @@ 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 @@ -1511,6 +1513,7 @@ def _update( return return_value _a2a_tasks = None + _memories = None _sandboxes = None _sessions = None _runtimes = None @@ -1545,6 +1548,21 @@ 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: @@ -2881,6 +2899,145 @@ 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, *, @@ -3771,6 +3928,7 @@ async def _update( return return_value _a2a_tasks = None + _memories = None _sessions = None _runtimes = None @@ -3832,6 +3990,21 @@ 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: @@ -3873,6 +4046,23 @@ 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, *, @@ -3890,6 +4080,23 @@ 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 e93e977980..0dcb0d445e 100644 --- a/agentplatform/_genai/client.py +++ b/agentplatform/_genai/client.py @@ -47,9 +47,6 @@ 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" @@ -95,7 +92,6 @@ 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( @@ -211,12 +207,6 @@ 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. @@ -324,7 +314,6 @@ 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": @@ -464,9 +453,3 @@ 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 6c660c14d2..c83cd240e1 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 _memory_bank_utils +from . import _agent_engines_utils from . import types if typing.TYPE_CHECKING: @@ -41,8 +41,63 @@ logger = logging.getLogger("agentplatform_genai.memories") +logger.setLevel(logging.INFO) -def _CreateMemoryRequestParameters_to_vertex( + +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( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -57,12 +112,12 @@ def _CreateMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _MemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + _AgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteMemoryRequestParameters_to_vertex( +def _DeleteAgentEngineMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -73,7 +128,7 @@ def _DeleteMemoryRequestParameters_to_vertex( return to_object -def _GenerateMemoriesConfig_to_vertex( +def _GenerateAgentEngineMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -130,7 +185,7 @@ def _GenerateMemoriesConfig_to_vertex( return to_object -def _GenerateMemoriesRequestParameters_to_vertex( +def _GenerateAgentEngineMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -163,12 +218,14 @@ def _GenerateMemoriesRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _GenerateMemoriesConfig_to_vertex(getv(from_object, ["config"]), to_object) + _GenerateAgentEngineMemoriesConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) return to_object -def _GetGenerateMemoriesOperationParameters_to_vertex( +def _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -181,7 +238,7 @@ def _GetGenerateMemoriesOperationParameters_to_vertex( return to_object -def _GetMemoryOperationParameters_to_vertex( +def _GetAgentEngineMemoryOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -194,7 +251,7 @@ def _GetMemoryOperationParameters_to_vertex( return to_object -def _GetMemoryRequestParameters_to_vertex( +def _GetAgentEngineMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -205,28 +262,53 @@ def _GetMemoryRequestParameters_to_vertex( return to_object -def _ListMemoriesConfig_to_vertex( +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, ["page_size"]) is not None: - setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) + if getv(from_object, ["force_flush"]) is not None: + setv(parent_object, ["forceFlush"], getv(from_object, ["force_flush"])) - if getv(from_object, ["page_token"]) is not None: - setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) + if getv(from_object, ["revision_labels"]) is not None: + setv(parent_object, ["revisionLabels"], getv(from_object, ["revision_labels"])) - if getv(from_object, ["filter"]) is not None: - setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) + if getv(from_object, ["revision_expire_time"]) is not None: + setv( + parent_object, + ["revisionExpireTime"], + getv(from_object, ["revision_expire_time"]), + ) - if getv(from_object, ["order_by"]) is not None: - setv(parent_object, ["_query", "orderBy"], getv(from_object, ["order_by"])) + 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 _ListMemoriesRequestParameters_to_vertex( +def _IngestEventsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -234,66 +316,68 @@ def _ListMemoriesRequestParameters_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: - _ListMemoriesConfig_to_vertex(getv(from_object, ["config"]), to_object) + _IngestEventsConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _MemoryConfig_to_vertex( +def _ListAgentEngineMemoryConfig_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, ["page_size"]) is not None: + setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) - if getv(from_object, ["expire_time"]) is not None: - setv(parent_object, ["expireTime"], getv(from_object, ["expire_time"])) + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) - if getv(from_object, ["revision_expire_time"]) is not None: - setv( - parent_object, - ["revisionExpireTime"], - getv(from_object, ["revision_expire_time"]), - ) + if getv(from_object, ["filter"]) is not None: + setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) - if getv(from_object, ["revision_ttl"]) is not None: - setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) + if getv(from_object, ["order_by"]) is not None: + setv(parent_object, ["_query", "orderBy"], getv(from_object, ["order_by"])) - if getv(from_object, ["disable_memory_revisions"]) is not None: - setv( - parent_object, - ["disableMemoryRevisions"], - getv(from_object, ["disable_memory_revisions"]), - ) + return 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()}, - ) +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, ["memory_id"]) is not None: - setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"])) + if getv(from_object, ["config"]) is not None: + _ListAgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _PurgeMemoriesRequestParameters_to_vertex( +def _PurgeAgentEngineMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -317,7 +401,7 @@ def _PurgeMemoriesRequestParameters_to_vertex( return to_object -def _RetrieveMemoriesConfig_to_vertex( +def _RetrieveAgentEngineMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -339,7 +423,7 @@ def _RetrieveMemoriesConfig_to_vertex( return to_object -def _RetrieveMemoriesRequestParameters_to_vertex( +def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -365,7 +449,9 @@ def _RetrieveMemoriesRequestParameters_to_vertex( ) if getv(from_object, ["config"]) is not None: - _RetrieveMemoriesConfig_to_vertex(getv(from_object, ["config"]), to_object) + _RetrieveAgentEngineMemoriesConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) return to_object @@ -384,7 +470,7 @@ def _RetrieveMemoryProfilesRequestParameters_to_vertex( return to_object -def _RollbackMemoryRequestParameters_to_vertex( +def _RollbackAgentEngineMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -398,7 +484,7 @@ def _RollbackMemoryRequestParameters_to_vertex( return to_object -def _UpdateMemoryConfig_to_vertex( +def _UpdateAgentEngineMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -456,7 +542,7 @@ def _UpdateMemoryConfig_to_vertex( return to_object -def _UpdateMemoryRequestParameters_to_vertex( +def _UpdateAgentEngineMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -471,7 +557,9 @@ def _UpdateMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _UpdateMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + _UpdateAgentEngineMemoryConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) return to_object @@ -484,13 +572,13 @@ def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.MemoryConfigOrDict] = None, - ) -> types.MemoryOperation: + config: Optional[types.AgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: """ - Creates a new memory in the Memory Bank. + Creates a new memory in the Agent Engine. """ - parameter_model = types._CreateMemoryRequestParameters( + parameter_model = types._CreateAgentEngineMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -503,7 +591,9 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -530,7 +620,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.MemoryOperation._from_response( + return_value = types.AgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -555,21 +645,24 @@ def _create( return return_value def delete( - self, *, name: str, config: Optional[types.DeleteMemoryConfigOrDict] = None - ) -> types.DeleteMemoryOperation: + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, + ) -> types.DeleteAgentEngineMemoryOperation: """ - Delete a memory. + Delete an Agent Engine memory. Args: name (str): - Required. The name of the memory to delete. Format: + Required. The name of the Agent Engine memory to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteRuntimeMemoryConfig): - Optional. Additional configurations for deleting the Agent Runtime. + config (DeleteAgentEngineMemoryConfig): + Optional. Additional configurations for deleting the Agent Engine. """ - parameter_model = types._DeleteMemoryRequestParameters( + parameter_model = types._DeleteAgentEngineMemoryRequestParameters( name=name, config=config, ) @@ -580,7 +673,9 @@ def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -607,7 +702,7 @@ def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteMemoryOperation._from_response( + return_value = types.DeleteAgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -645,13 +740,13 @@ def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateMemoriesConfigOrDict] = None, - ) -> types.GenerateMemoriesOperation: + config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEngineGenerateMemoriesOperation: """ - Generates memories for a Memory Bank. + Generates memories for an Agent Engine. """ - parameter_model = types._GenerateMemoriesRequestParameters( + parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -666,7 +761,9 @@ def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:generate".format_map(request_url_dict) @@ -693,7 +790,7 @@ def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.GenerateMemoriesOperation._from_response( + return_value = types.AgentEngineGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -718,10 +815,13 @@ def _generate( return return_value def get( - self, *, name: str, config: Optional[types.GetMemoryConfigOrDict] = None + self, + *, + name: str, + config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, ) -> types.Memory: """ - Gets a memory. + Gets an agent engine memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -730,7 +830,7 @@ def get( """ - parameter_model = types._GetMemoryRequestParameters( + parameter_model = types._GetAgentEngineMemoryRequestParameters( name=name, config=config, ) @@ -741,7 +841,9 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -792,14 +894,101 @@ 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.ListMemoriesConfigOrDict] = None - ) -> types.ListMemoriesResponse: + self, + *, + name: str, + config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + ) -> types.ListReasoningEnginesMemoriesResponse: """ - Lists memories. + Lists Agent Engine memories. """ - parameter_model = types._ListMemoriesRequestParameters( + parameter_model = types._ListAgentEngineMemoryRequestParameters( name=name, config=config, ) @@ -810,7 +999,9 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -837,7 +1028,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListMemoriesResponse._from_response( + return_value = types.ListReasoningEnginesMemoriesResponse._from_response( response=response_dict, kwargs=( { @@ -865,9 +1056,9 @@ def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, - ) -> types.MemoryOperation: - parameter_model = types._GetMemoryOperationParameters( + config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: + parameter_model = types._GetAgentEngineMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -878,7 +1069,9 @@ def _get_memory_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetMemoryOperationParameters_to_vertex(parameter_model) + request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{operationName}".format_map(request_url_dict) @@ -905,7 +1098,7 @@ def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.MemoryOperation._from_response( + return_value = types.AgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -933,9 +1126,9 @@ def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, - ) -> types.GenerateMemoriesOperation: - parameter_model = types._GetGenerateMemoriesOperationParameters( + config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + ) -> types.AgentEngineGenerateMemoriesOperation: + parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -946,7 +1139,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 = _GetGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -975,7 +1168,7 @@ def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.GenerateMemoriesOperation._from_response( + return_value = types.AgentEngineGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1010,13 +1203,13 @@ def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories from a Memory Bank. + Retrieves memories for an Agent Engine. """ - parameter_model = types._RetrieveMemoriesRequestParameters( + parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -1030,7 +1223,9 @@ def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:retrieve".format_map(request_url_dict) @@ -1089,13 +1284,13 @@ def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for a Memory Bank. + Retrieves memory profiles for an Agent Engine. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.memory_banks.memories.retrieve_profiles( + result = client.agent_engines.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -1188,13 +1383,13 @@ def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackMemoryConfigOrDict] = None, - ) -> types.RollbackMemoryOperation: + config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineRollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackMemoryRequestParameters( + parameter_model = types._RollbackAgentEngineMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -1206,7 +1401,9 @@ def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}:rollback".format_map(request_url_dict) @@ -1233,7 +1430,7 @@ def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.RollbackMemoryOperation._from_response( + return_value = types.AgentEngineRollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1263,13 +1460,13 @@ def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateMemoryConfigOrDict] = None, - ) -> types.MemoryOperation: + config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: """ - Updates a memory. + Updates an Agent Engine memory. """ - parameter_model = types._UpdateMemoryRequestParameters( + parameter_model = types._UpdateAgentEngineMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -1282,7 +1479,9 @@ def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -1309,7 +1508,7 @@ def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.MemoryOperation._from_response( + return_value = types.AgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1342,13 +1541,13 @@ def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeMemoriesConfigOrDict] = None, - ) -> types.PurgeMemoriesOperation: + config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEnginePurgeMemoriesOperation: """ - Purges memories from a Memory Bank. + Purges memories from an Agent Engine. """ - parameter_model = types._PurgeMemoriesRequestParameters( + parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -1362,7 +1561,9 @@ def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:purge".format_map(request_url_dict) @@ -1389,7 +1590,7 @@ def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.PurgeMemoriesOperation._from_response( + return_value = types.AgentEnginePurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1418,7 +1619,18 @@ def _purge( @property def revisions(self) -> "memory_revisions_module.MemoryRevisions": if self._revisions is None: - self._revisions = importlib.import_module(".memory_revisions", __package__) + 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 return self._revisions.MemoryRevisions(self._api_client) # type: ignore[no-any-return] def create( @@ -1427,9 +1639,9 @@ def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.MemoryConfigOrDict] = None, - ) -> types.MemoryOperation: - """Creates a new memory in Memory Bank. + config: Optional[types.AgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: + """Creates a new memory in the Agent Engine. Args: name (str): @@ -1438,16 +1650,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 (MemoryConfigOrDict): + config (AgentEngineMemoryConfigOrDict): Optional. The configuration for the memory. Returns: - MemoryOperation: The operation for creating the memory. + AgentEngineMemoryOperation: The operation for creating the memory. """ if config is None: - config = types.MemoryConfig() + config = types.AgentEngineMemoryConfig() elif isinstance(config, dict): - config = types.MemoryConfig.model_validate(config) + config = types.AgentEngineMemoryConfig.model_validate(config) operation = self._create( name=name, fact=fact, @@ -1456,7 +1668,7 @@ def create( ) if config.wait_for_completion: if not operation.done: - operation = _memory_bank_utils._await_operation( + operation = _agent_engines_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1485,13 +1697,13 @@ def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateMemoriesConfigOrDict] = None, - ) -> types.GenerateMemoriesOperation: - """Generates memories with Memory Bank. + config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEngineGenerateMemoriesOperation: + """Generates memories for the agent engine. Args: name (str): - Required. The name of the memory bank to generate memories with. + Required. The name of the agent engine to generate memories for. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -1512,13 +1724,13 @@ def generate( Optional. The configuration for the memories to generate. Returns: - GenerateMemoriesOperation: + AgentEngineGenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateMemoriesConfig() + config = types.GenerateAgentEngineMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateMemoriesConfig.model_validate(config) + config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) operation = self._generate( name=name, vertex_session_source=vertex_session_source, @@ -1528,7 +1740,7 @@ def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = _memory_bank_utils._await_operation( + operation = _agent_engines_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -1541,14 +1753,14 @@ def list( self, *, name: str, - config: Optional[types.ListMemoriesConfigOrDict] = None, + config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, ) -> Iterator[types.Memory]: - """Lists Memory Bank memories. + """Lists Agent Engine memories. Args: name (str): - Required. The name of the memory bank to list memories for. - config (ListMemoriesConfig): + Required. The name of the agent engine to list memories for. + config (ListAgentEngineMemoryConfig): Optional. The configuration for the memories to list. Returns: @@ -1573,13 +1785,13 @@ def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: - """Retrieves memories from Memory Bank. + """Retrieves memories for the agent. Args: name (str): - Required. The name of the memory bank to retrieve memories from. + Required. The name of the agent engine to retrieve memories for. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -1589,7 +1801,7 @@ def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveMemoriesConfig): + config (RetrieveAgentEngineMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -1620,8 +1832,8 @@ def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackMemoryConfigOrDict] = None, - ) -> types.RollbackMemoryOperation: + config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineRollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -1629,24 +1841,24 @@ def rollback( Required. The name of the memory to rollback. target_revision_id (str): Required. The revision ID to roll back to - config (RollbackMemoryConfig): + config (RollbackAgentEngineMemoryConfig): Optional. The configuration for the rollback. Returns: - RollbackMemoryOperation: + AgentEngineRollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackMemoryConfig() + config = types.RollbackAgentEngineMemoryConfig() elif isinstance(config, dict): - config = types.RollbackMemoryConfig.model_validate(config) + config = types.RollbackAgentEngineMemoryConfig.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 = _memory_bank_utils._await_operation( + operation = _agent_engines_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1662,13 +1874,13 @@ def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeMemoriesConfigOrDict] = None, - ) -> types.PurgeMemoriesOperation: - """Purges memories from Memory Bank. + config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEnginePurgeMemoriesOperation: + """Purges memories from an Agent Engine. Args: name (str): - Required. The name of the Memory Bank to purge memories from. + Required. The name of the Agent Engine to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -1678,17 +1890,17 @@ def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeMemoriesConfig): + config (PurgeAgentEngineMemoriesConfig): Optional. The configuration for the purge operation. Returns: - PurgeMemoriesOperation: + AgentEnginePurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeMemoriesConfig() + config = types.PurgeAgentEngineMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeMemoriesConfig.model_validate(config) + config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) operation = self._purge( name=name, filter=filter, @@ -1697,7 +1909,7 @@ def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = _memory_bank_utils._await_operation( + operation = _agent_engines_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1706,6 +1918,89 @@ 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): @@ -1715,13 +2010,13 @@ async def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.MemoryConfigOrDict] = None, - ) -> types.MemoryOperation: + config: Optional[types.AgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: """ - Creates a new memory in the Memory Bank. + Creates a new memory in the Agent Engine. """ - parameter_model = types._CreateMemoryRequestParameters( + parameter_model = types._CreateAgentEngineMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -1734,7 +2029,9 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -1763,7 +2060,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.MemoryOperation._from_response( + return_value = types.AgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1788,21 +2085,24 @@ async def _create( return return_value async def delete( - self, *, name: str, config: Optional[types.DeleteMemoryConfigOrDict] = None - ) -> types.DeleteMemoryOperation: + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, + ) -> types.DeleteAgentEngineMemoryOperation: """ - Delete a memory. + Delete an Agent Engine memory. Args: name (str): - Required. The name of the memory to delete. Format: + Required. The name of the Agent Engine memory to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteRuntimeMemoryConfig): - Optional. Additional configurations for deleting the Agent Runtime. + config (DeleteAgentEngineMemoryConfig): + Optional. Additional configurations for deleting the Agent Engine. """ - parameter_model = types._DeleteMemoryRequestParameters( + parameter_model = types._DeleteAgentEngineMemoryRequestParameters( name=name, config=config, ) @@ -1813,7 +2113,9 @@ async def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -1842,7 +2144,7 @@ async def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteMemoryOperation._from_response( + return_value = types.DeleteAgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1880,13 +2182,13 @@ async def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateMemoriesConfigOrDict] = None, - ) -> types.GenerateMemoriesOperation: + config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEngineGenerateMemoriesOperation: """ - Generates memories for a Memory Bank. + Generates memories for an Agent Engine. """ - parameter_model = types._GenerateMemoriesRequestParameters( + parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -1901,7 +2203,9 @@ async def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:generate".format_map(request_url_dict) @@ -1930,7 +2234,7 @@ async def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.GenerateMemoriesOperation._from_response( + return_value = types.AgentEngineGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1955,10 +2259,13 @@ async def _generate( return return_value async def get( - self, *, name: str, config: Optional[types.GetMemoryConfigOrDict] = None + self, + *, + name: str, + config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, ) -> types.Memory: """ - Gets a memory. + Gets an agent engine memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -1967,7 +2274,7 @@ async def get( """ - parameter_model = types._GetMemoryRequestParameters( + parameter_model = types._GetAgentEngineMemoryRequestParameters( name=name, config=config, ) @@ -1978,7 +2285,9 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -2031,14 +2340,103 @@ 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.ListMemoriesConfigOrDict] = None - ) -> types.ListMemoriesResponse: + self, + *, + name: str, + config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + ) -> types.ListReasoningEnginesMemoriesResponse: """ - Lists memories. + Lists Agent Engine memories. """ - parameter_model = types._ListMemoriesRequestParameters( + parameter_model = types._ListAgentEngineMemoryRequestParameters( name=name, config=config, ) @@ -2049,7 +2447,9 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories".format_map(request_url_dict) @@ -2078,7 +2478,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListMemoriesResponse._from_response( + return_value = types.ListReasoningEnginesMemoriesResponse._from_response( response=response_dict, kwargs=( { @@ -2106,9 +2506,9 @@ async def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, - ) -> types.MemoryOperation: - parameter_model = types._GetMemoryOperationParameters( + config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: + parameter_model = types._GetAgentEngineMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -2119,7 +2519,9 @@ 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 = _GetMemoryOperationParameters_to_vertex(parameter_model) + request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{operationName}".format_map(request_url_dict) @@ -2148,7 +2550,7 @@ async def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.MemoryOperation._from_response( + return_value = types.AgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2176,9 +2578,9 @@ async def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetMemoryBankOperationConfigOrDict] = None, - ) -> types.GenerateMemoriesOperation: - parameter_model = types._GetGenerateMemoriesOperationParameters( + config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + ) -> types.AgentEngineGenerateMemoriesOperation: + parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -2189,7 +2591,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 = _GetGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2220,7 +2622,7 @@ async def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.GenerateMemoriesOperation._from_response( + return_value = types.AgentEngineGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2255,13 +2657,13 @@ async def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories from a Memory Bank. + Retrieves memories for an Agent Engine. """ - parameter_model = types._RetrieveMemoriesRequestParameters( + parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -2275,7 +2677,9 @@ async def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:retrieve".format_map(request_url_dict) @@ -2336,13 +2740,13 @@ async def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for a Memory Bank. + Retrieves memory profiles for an Agent Engine. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.memory_banks.memories.retrieve_profiles( + result = client.agent_engines.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -2437,13 +2841,13 @@ async def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackMemoryConfigOrDict] = None, - ) -> types.RollbackMemoryOperation: + config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineRollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackMemoryRequestParameters( + parameter_model = types._RollbackAgentEngineMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -2455,7 +2859,9 @@ async def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}:rollback".format_map(request_url_dict) @@ -2484,7 +2890,7 @@ async def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.RollbackMemoryOperation._from_response( + return_value = types.AgentEngineRollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2514,13 +2920,13 @@ async def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateMemoryConfigOrDict] = None, - ) -> types.MemoryOperation: + config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: """ - Updates a memory. + Updates an Agent Engine memory. """ - parameter_model = types._UpdateMemoryRequestParameters( + parameter_model = types._UpdateAgentEngineMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -2533,7 +2939,9 @@ async def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateMemoryRequestParameters_to_vertex(parameter_model) + request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -2562,7 +2970,7 @@ async def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.MemoryOperation._from_response( + return_value = types.AgentEngineMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2595,13 +3003,13 @@ async def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeMemoriesConfigOrDict] = None, - ) -> types.PurgeMemoriesOperation: + config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEnginePurgeMemoriesOperation: """ - Purges memories from a Memory Bank. + Purges memories from an Agent Engine. """ - parameter_model = types._PurgeMemoriesRequestParameters( + parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -2615,7 +3023,9 @@ async def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeMemoriesRequestParameters_to_vertex(parameter_model) + request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( + parameter_model + ) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/memories:purge".format_map(request_url_dict) @@ -2644,7 +3054,7 @@ async def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.PurgeMemoriesOperation._from_response( + return_value = types.AgentEnginePurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2673,7 +3083,18 @@ async def _purge( @property def revisions(self) -> "memory_revisions_module.AsyncMemoryRevisions": if self._revisions is None: - self._revisions = importlib.import_module(".memory_revisions", __package__) + 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 return self._revisions.AsyncMemoryRevisions(self._api_client) # type: ignore[no-any-return] async def create( @@ -2682,9 +3103,9 @@ async def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.MemoryConfigOrDict] = None, - ) -> types.MemoryOperation: - """Creates a new memory in Memory Bank. + config: Optional[types.AgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineMemoryOperation: + """Creates a new memory in the Agent Engine. Args: name (str): @@ -2693,16 +3114,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 (MemoryConfigOrDict): + config (AgentEngineMemoryConfigOrDict): Optional. The configuration for the memory. Returns: - MemoryOperation: The operation for creating the memory. + AgentEngineMemoryOperation: The operation for creating the memory. """ if config is None: - config = types.MemoryConfig() + config = types.AgentEngineMemoryConfig() elif isinstance(config, dict): - config = types.MemoryConfig.model_validate(config) + config = types.AgentEngineMemoryConfig.model_validate(config) operation = await self._create( name=name, fact=fact, @@ -2711,7 +3132,7 @@ async def create( ) if config.wait_for_completion: if not operation.done: - operation = await _memory_bank_utils._await_async_operation( + operation = await _agent_engines_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -2740,13 +3161,13 @@ async def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateMemoriesConfigOrDict] = None, - ) -> types.GenerateMemoriesOperation: - """Generates memories for Memory Bank. + config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEngineGenerateMemoriesOperation: + """Generates memories for the agent engine. Args: name (str): - Required. The name of the Memory Bank to generate memories for. + Required. The name of the agent engine to generate memories for. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -2767,13 +3188,13 @@ async def generate( Optional. The configuration for the memories to generate. Returns: - GenerateMemoriesOperation: + AgentEngineGenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateMemoriesConfig() + config = types.GenerateAgentEngineMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateMemoriesConfig.model_validate(config) + config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) operation = await self._generate( name=name, vertex_session_source=vertex_session_source, @@ -2783,7 +3204,7 @@ async def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _memory_bank_utils._await_async_operation( + operation = await _agent_engines_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -2796,14 +3217,14 @@ async def list( self, *, name: str, - config: Optional[types.ListMemoriesConfigOrDict] = None, + config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, ) -> AsyncPager[types.Memory]: - """Lists memories in Memory Bank. + """Lists Agent Engine memories. Args: name (str): - Required. The name of the Memory Bank to list memories for. - config (ListMemoriesConfig): + Required. The name of the agent engine to list memories for. + config (ListAgentEngineMemoryConfig): Optional. The configuration for the memories to list. Returns: @@ -2828,13 +3249,13 @@ async def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, ) -> AsyncPager[types.RetrieveMemoriesResponseRetrievedMemory]: - """Retrieves memories from Memory Bank. + """Retrieves memories for the agent. Args: name (str): - Required. The name of the Memory Bank to retrieve memories for. + Required. The name of the agent engine to retrieve memories for. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -2844,7 +3265,7 @@ async def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveMemoriesConfig): + config (RetrieveAgentEngineMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -2875,8 +3296,8 @@ async def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackMemoryConfigOrDict] = None, - ) -> types.RollbackMemoryOperation: + config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, + ) -> types.AgentEngineRollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -2884,24 +3305,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 (RollbackMemoryConfig): + config (RollbackAgentEngineMemoryConfig): Optional. The configuration for the rollback. Returns: - RollbackMemoryOperation: + AgentEngineRollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackMemoryConfig() + config = types.RollbackAgentEngineMemoryConfig() elif isinstance(config, dict): - config = types.RollbackMemoryConfig.model_validate(config) + config = types.RollbackAgentEngineMemoryConfig.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 _memory_bank_utils._await_async_operation( + operation = await _agent_engines_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -2917,13 +3338,13 @@ async def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeMemoriesConfigOrDict] = None, - ) -> types.PurgeMemoriesOperation: - """Purges memories from a Memory Bank. + config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, + ) -> types.AgentEnginePurgeMemoriesOperation: + """Purges memories from an Agent Engine. Args: name (str): - Required. The name of the Memory Bank to purge memories from. + Required. The name of the Agent Engine to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -2933,17 +3354,17 @@ async def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeMemoriesConfig): + config (PurgeAgentEngineMemoriesConfig): Optional. The configuration for the purge operation. Returns: - MemoryBankPurgeMemoriesOperation: + AgentEnginePurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeMemoriesConfig() + config = types.PurgeAgentEngineMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeMemoriesConfig.model_validate(config) + config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) operation = await self._purge( name=name, filter=filter, @@ -2952,7 +3373,7 @@ async def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _memory_bank_utils._await_async_operation( + operation = await _agent_engines_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -2960,3 +3381,86 @@ 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 deleted file mode 100644 index c0b0cee277..0000000000 --- a/agentplatform/_genai/memory_banks.py +++ /dev/null @@ -1,1021 +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. -# - -# 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. - """ - 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. - """ - 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 73917a973b..c1a1e307b1 100644 --- a/agentplatform/_genai/memory_revisions.py +++ b/agentplatform/_genai/memory_revisions.py @@ -31,8 +31,10 @@ logger = logging.getLogger("agentplatform_genai.memoryrevisions") +logger.setLevel(logging.INFO) -def _GetMemoryRevisionRequestParameters_to_vertex( + +def _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -43,7 +45,7 @@ def _GetMemoryRevisionRequestParameters_to_vertex( return to_object -def _ListMemoryRevisionsConfig_to_vertex( +def _ListAgentEngineMemoryRevisionsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -61,7 +63,7 @@ def _ListMemoryRevisionsConfig_to_vertex( return to_object -def _ListMemoryRevisionsRequestParameters_to_vertex( +def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -70,7 +72,9 @@ def _ListMemoryRevisionsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListMemoryRevisionsConfig_to_vertex(getv(from_object, ["config"]), to_object) + _ListAgentEngineMemoryRevisionsConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) return to_object @@ -78,23 +82,26 @@ def _ListMemoryRevisionsRequestParameters_to_vertex( class MemoryRevisions(_api_module.BaseModule): def get( - self, *, name: str, config: Optional[types.GetMemoryRevisionConfigOrDict] = None + self, + *, + name: str, + config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, ) -> types.MemoryRevision: """ - Gets a memory revision. + Gets an agent engine memory revision. Args: - name (str): Required. The name of the memory revision to get. Format: + name (str): Required. The name of the Agent Engine memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetMemoryRevisionConfig): - Optional. Additional configurations for getting the memory revision. + config (GetAgentEngineMemoryRevisionConfig): + Optional. Additional configurations for getting the Agent Engine memory revision. Returns: - MemoryRevision: The requested memory revision. + AgentEngineMemoryRevision: The requested Agent Engine memory revision. """ - parameter_model = types._GetMemoryRevisionRequestParameters( + parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( name=name, config=config, ) @@ -105,7 +112,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -162,23 +169,23 @@ def _list( self, *, name: str, - config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, - ) -> types.ListMemoryRevisionsResponse: + config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + ) -> types.ListAgentEngineMemoryRevisionsResponse: """ - Lists memory revisions. + Lists Agent Engine memory revisions. Args: - name (str): Required. The name of the memory to list revisions for. Format: + name (str): Required. The name of the Agent Engine memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListMemoryRevisionsConfig): - Optional. Additional configurations for listing the memory revisions. + config (ListAgentEngineMemoryRevisionsConfig): + Optional. Additional configurations for listing the Agent Engine memory revisions. Returns: - LisemoryRevisionsResponse: The requested memory revisions. + ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. """ - parameter_model = types._ListMemoryRevisionsRequestParameters( + parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -189,7 +196,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -218,7 +225,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListMemoryRevisionsResponse._from_response( + return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -246,14 +253,14 @@ def list( self, *, name: str, - config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, ) -> Iterator[types.MemoryRevision]: - """Lists memory revisions. + """Lists Agent Engine memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListMemoryRevisionsConfigOrDict): + config (ListAgentEngineMemoryRevisionsConfigOrDict): Optional. The configuration for the memories to list revisions. Returns: @@ -271,23 +278,26 @@ def list( class AsyncMemoryRevisions(_api_module.BaseModule): async def get( - self, *, name: str, config: Optional[types.GetMemoryRevisionConfigOrDict] = None + self, + *, + name: str, + config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, ) -> types.MemoryRevision: """ - Gets a memory revision. + Gets an agent engine memory revision. Args: - name (str): Required. The name of the memory revision to get. Format: + name (str): Required. The name of the Agent Engine memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetMemoryRevisionConfig): - Optional. Additional configurations for getting the memory revision. + config (GetAgentEngineMemoryRevisionConfig): + Optional. Additional configurations for getting the Agent Engine memory revision. Returns: - MemoryRevision: The requested memory revision. + AgentEngineMemoryRevision: The requested Agent Engine memory revision. """ - parameter_model = types._GetMemoryRevisionRequestParameters( + parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( name=name, config=config, ) @@ -298,7 +308,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -357,23 +367,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, - ) -> types.ListMemoryRevisionsResponse: + config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + ) -> types.ListAgentEngineMemoryRevisionsResponse: """ - Lists memory revisions. + Lists Agent Engine memory revisions. Args: - name (str): Required. The name of the memory to list revisions for. Format: + name (str): Required. The name of the Agent Engine memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListMemoryRevisionsConfig): - Optional. Additional configurations for listing the memory revisions. + config (ListAgentEngineMemoryRevisionsConfig): + Optional. Additional configurations for listing the Agent Engine memory revisions. Returns: - LisemoryRevisionsResponse: The requested memory revisions. + ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. """ - parameter_model = types._ListMemoryRevisionsRequestParameters( + parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -384,7 +394,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -415,7 +425,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListMemoryRevisionsResponse._from_response( + return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -443,14 +453,14 @@ async def list( self, *, name: str, - config: Optional[types.ListMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, ) -> AsyncPager[types.MemoryRevision]: - """Lists memory revisions. + """Lists Agent Engine memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListMemoryRevisionsConfigOrDict): + config (ListAgentEngineMemoryRevisionsConfigOrDict): 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 9e9278030a..6d719ddd84 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -29,6 +29,7 @@ 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 @@ -40,8 +41,6 @@ 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 @@ -50,6 +49,7 @@ 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,8 +58,6 @@ 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 @@ -71,10 +69,14 @@ 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 @@ -97,12 +99,7 @@ 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 @@ -123,6 +120,8 @@ 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 @@ -134,8 +133,6 @@ 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 @@ -147,25 +144,25 @@ from .common import _ListSkillsRequestParameters from .common import _OptimizeRequestParameters from .common import _OptimizeRequestParameters -from .common import _PurgeMemoriesRequestParameters +from .common import _PurgeAgentEngineMemoriesRequestParameters from .common import _QueryAgentEngineRequestParameters from .common import _QueryAgentEngineRuntimeRevisionRequestParameters from .common import _RecommendSpecRequestParameters from .common import _RestoreVersionRequestParameters -from .common import _RetrieveMemoriesRequestParameters +from .common import _RetrieveAgentEngineMemoriesRequestParameters from .common import _RetrieveMemoryProfilesRequestParameters from .common import _RetrieveRagContextsRequestParameters from .common import _RetrieveSkillsRequestParameters -from .common import _RollbackMemoryRequestParameters +from .common import _RollbackAgentEngineMemoryRequestParameters 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 @@ -183,10 +180,25 @@ 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 @@ -350,9 +362,6 @@ 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 @@ -405,6 +414,12 @@ 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 @@ -441,18 +456,6 @@ 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 @@ -667,6 +670,9 @@ 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 @@ -679,12 +685,6 @@ 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,6 +716,12 @@ 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 @@ -761,15 +767,6 @@ 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 @@ -877,6 +874,15 @@ 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 @@ -931,18 +937,6 @@ 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 @@ -976,6 +970,9 @@ 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 @@ -1043,7 +1040,6 @@ 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 @@ -1087,17 +1083,9 @@ 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 @@ -1115,9 +1103,6 @@ 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 @@ -1284,12 +1269,9 @@ from .common import PublisherModelResourceReference from .common import PublisherModelResourceReferenceDict from .common import PublisherModelResourceReferenceOrDict -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 PurgeAgentEngineMemoriesConfig +from .common import PurgeAgentEngineMemoriesConfigDict +from .common import PurgeAgentEngineMemoriesConfigOrDict from .common import PurgeMemoriesResponse from .common import PurgeMemoriesResponseDict from .common import PurgeMemoriesResponseOrDict @@ -1564,6 +1546,9 @@ 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 @@ -1573,9 +1558,6 @@ 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 @@ -1600,12 +1582,9 @@ from .common import RetrieveSkillsResponse from .common import RetrieveSkillsResponseDict from .common import RetrieveSkillsResponseOrDict -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 RollbackAgentEngineMemoryConfig +from .common import RollbackAgentEngineMemoryConfigDict +from .common import RollbackAgentEngineMemoryConfigOrDict from .common import RougeInput from .common import RougeInputDict from .common import RougeInputOrDict @@ -1943,15 +1922,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 @@ -2660,57 +2639,27 @@ "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", - "IngestEventsConfig", - "IngestEventsConfigDict", - "IngestEventsConfigOrDict", - "MemoryBankIngestEventsOperation", - "MemoryBankIngestEventsOperationDict", - "MemoryBankIngestEventsOperationOrDict", - "GetMemoryBankOperationConfig", - "GetMemoryBankOperationConfigDict", - "GetMemoryBankOperationConfigOrDict", - "MemoryConfig", - "MemoryConfigDict", - "MemoryConfigOrDict", + "AgentEngineMemoryConfig", + "AgentEngineMemoryConfigDict", + "AgentEngineMemoryConfigOrDict", "MemoryStructuredContent", "MemoryStructuredContentDict", "MemoryStructuredContentOrDict", "Memory", "MemoryDict", "MemoryOrDict", - "MemoryOperation", - "MemoryOperationDict", - "MemoryOperationOrDict", - "DeleteMemoryConfig", - "DeleteMemoryConfigDict", - "DeleteMemoryConfigOrDict", - "DeleteMemoryOperation", - "DeleteMemoryOperationDict", - "DeleteMemoryOperationOrDict", + "AgentEngineMemoryOperation", + "AgentEngineMemoryOperationDict", + "AgentEngineMemoryOperationOrDict", + "DeleteAgentEngineMemoryConfig", + "DeleteAgentEngineMemoryConfigDict", + "DeleteAgentEngineMemoryConfigOrDict", + "DeleteAgentEngineMemoryOperation", + "DeleteAgentEngineMemoryOperationDict", + "DeleteAgentEngineMemoryOperationOrDict", "GenerateMemoriesRequestVertexSessionSource", "GenerateMemoriesRequestVertexSessionSourceDict", "GenerateMemoriesRequestVertexSessionSourceOrDict", @@ -2726,27 +2675,39 @@ "GenerateMemoriesRequestDirectMemoriesSource", "GenerateMemoriesRequestDirectMemoriesSourceDict", "GenerateMemoriesRequestDirectMemoriesSourceOrDict", - "GenerateMemoriesConfig", - "GenerateMemoriesConfigDict", - "GenerateMemoriesConfigOrDict", + "GenerateAgentEngineMemoriesConfig", + "GenerateAgentEngineMemoriesConfigDict", + "GenerateAgentEngineMemoriesConfigOrDict", "GenerateMemoriesResponseGeneratedMemory", "GenerateMemoriesResponseGeneratedMemoryDict", "GenerateMemoriesResponseGeneratedMemoryOrDict", "GenerateMemoriesResponse", "GenerateMemoriesResponseDict", "GenerateMemoriesResponseOrDict", - "GenerateMemoriesOperation", - "GenerateMemoriesOperationDict", - "GenerateMemoriesOperationOrDict", - "GetMemoryConfig", - "GetMemoryConfigDict", - "GetMemoryConfigOrDict", - "ListMemoriesConfig", - "ListMemoriesConfigDict", - "ListMemoriesConfigOrDict", - "ListMemoriesResponse", - "ListMemoriesResponseDict", - "ListMemoriesResponseOrDict", + "AgentEngineGenerateMemoriesOperation", + "AgentEngineGenerateMemoriesOperationDict", + "AgentEngineGenerateMemoriesOperationOrDict", + "GetAgentEngineMemoryConfig", + "GetAgentEngineMemoryConfigDict", + "GetAgentEngineMemoryConfigOrDict", + "IngestionDirectContentsSourceEvent", + "IngestionDirectContentsSourceEventDict", + "IngestionDirectContentsSourceEventOrDict", + "IngestionDirectContentsSource", + "IngestionDirectContentsSourceDict", + "IngestionDirectContentsSourceOrDict", + "IngestEventsConfig", + "IngestEventsConfigDict", + "IngestEventsConfigOrDict", + "MemoryBankIngestEventsOperation", + "MemoryBankIngestEventsOperationDict", + "MemoryBankIngestEventsOperationOrDict", + "ListAgentEngineMemoryConfig", + "ListAgentEngineMemoryConfigDict", + "ListAgentEngineMemoryConfigOrDict", + "ListReasoningEnginesMemoriesResponse", + "ListReasoningEnginesMemoriesResponseDict", + "ListReasoningEnginesMemoriesResponseOrDict", "RetrieveMemoriesRequestSimilaritySearchParams", "RetrieveMemoriesRequestSimilaritySearchParamsDict", "RetrieveMemoriesRequestSimilaritySearchParamsOrDict", @@ -2759,9 +2720,9 @@ "MemoryConjunctionFilter", "MemoryConjunctionFilterDict", "MemoryConjunctionFilterOrDict", - "RetrieveMemoriesConfig", - "RetrieveMemoriesConfigDict", - "RetrieveMemoriesConfigOrDict", + "RetrieveAgentEngineMemoriesConfig", + "RetrieveAgentEngineMemoriesConfigDict", + "RetrieveAgentEngineMemoriesConfigOrDict", "RetrieveMemoriesResponseRetrievedMemory", "RetrieveMemoriesResponseRetrievedMemoryDict", "RetrieveMemoriesResponseRetrievedMemoryOrDict", @@ -2777,39 +2738,39 @@ "RetrieveProfilesResponse", "RetrieveProfilesResponseDict", "RetrieveProfilesResponseOrDict", - "RollbackMemoryConfig", - "RollbackMemoryConfigDict", - "RollbackMemoryConfigOrDict", - "RollbackMemoryOperation", - "RollbackMemoryOperationDict", - "RollbackMemoryOperationOrDict", - "UpdateMemoryConfig", - "UpdateMemoryConfigDict", - "UpdateMemoryConfigOrDict", - "PurgeMemoriesConfig", - "PurgeMemoriesConfigDict", - "PurgeMemoriesConfigOrDict", + "RollbackAgentEngineMemoryConfig", + "RollbackAgentEngineMemoryConfigDict", + "RollbackAgentEngineMemoryConfigOrDict", + "AgentEngineRollbackMemoryOperation", + "AgentEngineRollbackMemoryOperationDict", + "AgentEngineRollbackMemoryOperationOrDict", + "UpdateAgentEngineMemoryConfig", + "UpdateAgentEngineMemoryConfigDict", + "UpdateAgentEngineMemoryConfigOrDict", + "PurgeAgentEngineMemoriesConfig", + "PurgeAgentEngineMemoriesConfigDict", + "PurgeAgentEngineMemoriesConfigOrDict", "PurgeMemoriesResponse", "PurgeMemoriesResponseDict", "PurgeMemoriesResponseOrDict", - "PurgeMemoriesOperation", - "PurgeMemoriesOperationDict", - "PurgeMemoriesOperationOrDict", - "GetMemoryRevisionConfig", - "GetMemoryRevisionConfigDict", - "GetMemoryRevisionConfigOrDict", + "AgentEnginePurgeMemoriesOperation", + "AgentEnginePurgeMemoriesOperationDict", + "AgentEnginePurgeMemoriesOperationOrDict", + "GetAgentEngineMemoryRevisionConfig", + "GetAgentEngineMemoryRevisionConfigDict", + "GetAgentEngineMemoryRevisionConfigOrDict", "IntermediateExtractedMemory", "IntermediateExtractedMemoryDict", "IntermediateExtractedMemoryOrDict", "MemoryRevision", "MemoryRevisionDict", "MemoryRevisionOrDict", - "ListMemoryRevisionsConfig", - "ListMemoryRevisionsConfigDict", - "ListMemoryRevisionsConfigOrDict", - "ListMemoryRevisionsResponse", - "ListMemoryRevisionsResponseDict", - "ListMemoryRevisionsResponseOrDict", + "ListAgentEngineMemoryRevisionsConfig", + "ListAgentEngineMemoryRevisionsConfigDict", + "ListAgentEngineMemoryRevisionsConfigOrDict", + "ListAgentEngineMemoryRevisionsResponse", + "ListAgentEngineMemoryRevisionsResponseDict", + "ListAgentEngineMemoryRevisionsResponseOrDict", "AskContextsConfig", "AskContextsConfigDict", "AskContextsConfigOrDict", @@ -3838,24 +3799,21 @@ "_GetAgentEngineOperationParameters", "_QueryAgentEngineRequestParameters", "_UpdateAgentEngineRequestParameters", - "_CreateMemoryBankRequestParameters", - "_DeleteMemoryBankRequestParameters", + "_CreateAgentEngineMemoryRequestParameters", + "_DeleteAgentEngineMemoryRequestParameters", + "_GenerateAgentEngineMemoriesRequestParameters", + "_GetAgentEngineMemoryRequestParameters", "_IngestEventsRequestParameters", - "_GetMemoryBankOperationParameters", - "_CreateMemoryRequestParameters", - "_DeleteMemoryRequestParameters", - "_GenerateMemoriesRequestParameters", - "_GetMemoryRequestParameters", - "_ListMemoriesRequestParameters", - "_GetMemoryOperationParameters", - "_GetGenerateMemoriesOperationParameters", - "_RetrieveMemoriesRequestParameters", + "_ListAgentEngineMemoryRequestParameters", + "_GetAgentEngineMemoryOperationParameters", + "_GetAgentEngineGenerateMemoriesOperationParameters", + "_RetrieveAgentEngineMemoriesRequestParameters", "_RetrieveMemoryProfilesRequestParameters", - "_RollbackMemoryRequestParameters", - "_UpdateMemoryRequestParameters", - "_PurgeMemoriesRequestParameters", - "_GetMemoryRevisionRequestParameters", - "_ListMemoryRevisionsRequestParameters", + "_RollbackAgentEngineMemoryRequestParameters", + "_UpdateAgentEngineMemoryRequestParameters", + "_PurgeAgentEngineMemoriesRequestParameters", + "_GetAgentEngineMemoryRevisionRequestParameters", + "_ListAgentEngineMemoryRevisionsRequestParameters", "_AskContextsRequestParameters", "_CreateRagCorpusRequestParameters", "_GetCorpusOperationParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index b0343429c1..24265d7dab 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -10184,424 +10184,353 @@ class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False): ] -class CreateMemoryBankConfig(_common.BaseModel): - """Config for create memory bank.""" +class MemoryMetadataValue(_common.BaseModel): + """The metadata values for memories.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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.""", ) -class CreateMemoryBankConfigDict(TypedDict, total=False): - """Config for create memory bank.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -CreateMemoryBankConfigOrDict = Union[CreateMemoryBankConfig, CreateMemoryBankConfigDict] - - -class _CreateMemoryBankRequestParameters(_common.BaseModel): - """Parameters for creating memory banks.""" +class MemoryMetadataValueDict(TypedDict, total=False): + """The metadata values for memories.""" - config: Optional[CreateMemoryBankConfig] = Field(default=None, description="""""") + bool_value: Optional[bool] + """Represents a boolean value.""" + double_value: Optional[float] + """Represents a double value.""" -class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): - """Parameters for creating memory banks.""" + string_value: Optional[str] + """Represents a string value.""" - config: Optional[CreateMemoryBankConfigDict] - """""" + timestamp_value: Optional[datetime.datetime] + """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" -_CreateMemoryBankRequestParametersOrDict = Union[ - _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict -] +MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] -class MemoryBank(_common.BaseModel): - """A memory bank.""" +class AgentEngineMemoryConfig(_common.BaseModel): + """Config for creating a Memory.""" - name: Optional[str] = Field( - default=None, - description="""Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""", + 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. - -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( + The expiration time is computed: now + TTL.""", + ) + expire_time: Optional[datetime.datetime] = 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. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", ) - metadata: Optional[dict[str, Any]] = Field( + revision_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. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", ) - done: Optional[bool] = Field( + revision_ttl: Optional[str] = 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. The TTL for the revision. The expiration time is computed: now + TTL.""", ) - error: Optional[dict[str, Any]] = Field( + disable_memory_revisions: Optional[bool] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Optional. Input only. If true, no revision will be created for this request.""", ) - response: Optional[MemoryBank] = Field( - default=None, description="""The created Memory Bank.""" + topics: Optional[list[MemoryTopicId]] = Field( + default=None, description="""Optional. The topics of the memory.""" ) - - -class MemoryBankOperationDict(TypedDict, total=False): - """Operation that has an memory bank as a response.""" - - 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.""" - - 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.""" - - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" - - response: Optional[MemoryBankDict] - """The created Memory Bank.""" - - -MemoryBankOperationOrDict = Union[MemoryBankOperation, MemoryBankOperationDict] - - -class DeleteMemoryBankConfig(_common.BaseModel): - """Config for delete memory bank.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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 DeleteMemoryBankConfigDict(TypedDict, total=False): - """Config for delete memory bank.""" +class AgentEngineMemoryConfigDict(TypedDict, total=False): + """Config for creating a Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + display_name: Optional[str] + """The display name of the memory.""" -DeleteMemoryBankConfigOrDict = Union[DeleteMemoryBankConfig, DeleteMemoryBankConfigDict] + description: Optional[str] + """The description of the memory.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -class _DeleteMemoryBankRequestParameters(_common.BaseModel): - """Parameters for deleting a memory bank.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. - 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="""""") + 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.""" -class _DeleteMemoryBankRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a memory bank.""" + 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.""" - name: Optional[str] - """Name of the memory bank.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - force: Optional[bool] - """If set to true, any child resources will also be deleted.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. If true, no revision will be created for this request.""" - config: Optional[DeleteMemoryBankConfigDict] - """""" + 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.""" + + 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.""" -_DeleteMemoryBankRequestParametersOrDict = Union[ - _DeleteMemoryBankRequestParameters, _DeleteMemoryBankRequestParametersDict +AgentEngineMemoryConfigOrDict = Union[ + AgentEngineMemoryConfig, AgentEngineMemoryConfigDict ] -class DeleteMemoryBankOperation(_common.BaseModel): - """Operation for deleting a memory bank.""" +class _CreateAgentEngineMemoryRequestParameters(_common.BaseModel): + """Parameters for creating Agent Engine Memories.""" name: 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}`.""", - ) - metadata: Optional[dict[str, Any]] = 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="""Name of the agent engine to create the memory under.""", ) - done: Optional[bool] = Field( + fact: Optional[str] = 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="""The fact of the memory. + + This is the semantic knowledge extracted from the source content).""", ) - error: Optional[dict[str, Any]] = Field( + scope: Optional[dict[str, str]] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + 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 '*'.""", ) + config: Optional[AgentEngineMemoryConfig] = Field(default=None, description="""""") -class DeleteMemoryBankOperationDict(TypedDict, total=False): - """Operation for deleting a memory bank.""" +class _CreateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Engine 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}`.""" + """Name of the agent engine to create the memory under.""" - 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.""" + fact: Optional[str] + """The fact 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.""" + This is the semantic knowledge extracted from the source content).""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + scope: Optional[dict[str, str]] + """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 '*'.""" + + config: Optional[AgentEngineMemoryConfigDict] + """""" -DeleteMemoryBankOperationOrDict = Union[ - DeleteMemoryBankOperation, DeleteMemoryBankOperationDict +_CreateAgentEngineMemoryRequestParametersOrDict = Union[ + _CreateAgentEngineMemoryRequestParameters, + _CreateAgentEngineMemoryRequestParametersDict, ] -class IngestionDirectContentsSourceEvent(_common.BaseModel): - """The direct contents source event for ingesting events.""" +class MemoryStructuredContent(_common.BaseModel): + """Represents the structured value of the memory.""" - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. The content of the event.""" - ) - event_id: Optional[str] = Field( + data: Optional[dict[str, Any]] = 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. Represents the structured value of the memory.""", ) - event_time: Optional[datetime.datetime] = Field( + schema_id: Optional[str] = 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="""Required. Represents the schema ID for which this structured memory belongs to.""", ) -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 MemoryStructuredContentDict(TypedDict, total=False): + """Represents the structured value of the memory.""" - 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.""" + data: Optional[dict[str, Any]] + """Required. Represents the structured value of the memory.""" - 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.""" + schema_id: Optional[str] + """Required. Represents the schema ID for which this structured memory belongs to.""" -IngestionDirectContentsSourceEventOrDict = Union[ - IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict +MemoryStructuredContentOrDict = Union[ + MemoryStructuredContent, MemoryStructuredContentDict ] -class IngestionDirectContentsSource(_common.BaseModel): - """The direct contents source for ingesting events.""" +class Memory(_common.BaseModel): + """A memory.""" - events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( - default=None, description="""Required. The events to ingest.""" + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Represents the timestamp when this Memory was created.""", ) - - -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.""" + description: Optional[str] = Field( + default=None, + description="""Optional. Represents the description of the Memory.""", ) - double_value: Optional[float] = Field( - default=None, description="""Represents a double value.""" + disable_memory_revisions: Optional[bool] = Field( + default=None, + description="""Optional. Input only. Indicates whether no revision will be created for this request.""", ) - string_value: Optional[str] = Field( - default=None, description="""Represents a string value.""" - ) - timestamp_value: Optional[datetime.datetime] = Field( + display_name: Optional[str] = Field( default=None, - description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", + description="""Optional. Represents the display name of the Memory.""", ) - - -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.""" + 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.""", ) - wait_for_completion: Optional[bool] = Field( - default=False, - description="""Waits for the underlying memory generation operation to complete - before returning. Defaults to false.""", + fact: Optional[str] = Field( + default=None, + description="""Optional. Represents semantic knowledge extracted from the source content.""", ) - force_flush: Optional[bool] = Field( + metadata: Optional[dict[str, MemoryMetadataValue]] = 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="""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]] = Field( + name: Optional[str] = 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="""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.""", + ) + memory_type: Optional[MemoryType] = Field( + default=None, + description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", + ) + structured_content: Optional[MemoryStructuredContent] = Field( + default=None, + description="""Optional. Represents the structured content of the memory.""", ) -class IngestEventsConfigDict(TypedDict, total=False): - """Config for ingesting events.""" - - 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.""" - - 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`.""" - - 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.""" +class MemoryDict(TypedDict, total=False): + """A 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.""" + create_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was created.""" - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + description: Optional[str] + """Optional. Represents the description of the Memory.""" disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revisions will be created for this request.""" - - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" - - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] - """Optional. The strategy to use when applying metadata to existing memories.""" - + """Optional. Input only. Indicates whether no revision will be created for this request.""" -IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] + 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.""" -class _IngestEventsRequestParameters(_common.BaseModel): - """Parameters for ingesting events to Memory Bank.""" + fact: Optional[str] + """Optional. Represents semantic knowledge extracted from the source content.""" - 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. + 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.""" - 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="""""") + 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. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" -class _IngestEventsRequestParametersDict(TypedDict, total=False): - """Parameters for ingesting events to Memory Bank.""" + 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.""" - name: Optional[str] - """Name of the Memory Bank to ingest events into.""" + revision_ttl: Optional[str] + """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" - stream_id: Optional[str] - """The ID of the stream to ingest events into.""" + 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 '*'.""" - direct_contents_source: Optional[IngestionDirectContentsSourceDict] - """The direct memories source of the events that should be ingested.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. Represents the Topics of the Memory.""" - scope: Optional[dict[str, str]] - """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 '*'.""" + update_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was most recently updated.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """The configuration for the memory generation trigger.""" + 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.""" - config: Optional[IngestEventsConfigDict] - """""" + structured_content: Optional[MemoryStructuredContentDict] + """Optional. Represents the structured content of the memory.""" -_IngestEventsRequestParametersOrDict = Union[ - _IngestEventsRequestParameters, _IngestEventsRequestParametersDict -] +MemoryOrDict = Union[Memory, MemoryDict] -class MemoryBankIngestEventsOperation(_common.BaseModel): - """Operation that ingests events into a memory bank.""" +class AgentEngineMemoryOperation(_common.BaseModel): + """Operation that has an agent engine memory as a response.""" name: Optional[str] = Field( default=None, @@ -10619,10 +10548,13 @@ class MemoryBankIngestEventsOperation(_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 MemoryBankIngestEventsOperationDict(TypedDict, total=False): - """Operation that ingests events into a memory bank.""" +class AgentEngineMemoryOperationDict(TypedDict, total=False): + """Operation that has an agent engine memory as a response.""" 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}`.""" @@ -10636,449 +10568,465 @@ class MemoryBankIngestEventsOperationDict(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.""" -MemoryBankIngestEventsOperationOrDict = Union[ - MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict + +AgentEngineMemoryOperationOrDict = Union[ + AgentEngineMemoryOperation, AgentEngineMemoryOperationDict ] -class GetMemoryBankOperationConfig(_common.BaseModel): +class DeleteAgentEngineMemoryConfig(_common.BaseModel): + """Config for deleting an Agent Engine Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetMemoryBankOperationConfigDict(TypedDict, total=False): +class DeleteAgentEngineMemoryConfigDict(TypedDict, total=False): + """Config for deleting an Agent Engine Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetMemoryBankOperationConfigOrDict = Union[ - GetMemoryBankOperationConfig, GetMemoryBankOperationConfigDict +DeleteAgentEngineMemoryConfigOrDict = Union[ + DeleteAgentEngineMemoryConfig, DeleteAgentEngineMemoryConfigDict ] -class _GetMemoryBankOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a memory bank as a response.""" +class _DeleteAgentEngineMemoryRequestParameters(_common.BaseModel): + """Parameters for deleting agent engines.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + name: Optional[str] = Field( + default=None, description="""Name of the agent engine memory to delete.""" ) - config: Optional[GetMemoryBankOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + config: Optional[DeleteAgentEngineMemoryConfig] = Field( + default=None, description="""""" ) -class _GetMemoryBankOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a memory bank as a response.""" +class _DeleteAgentEngineMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent engines.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + name: Optional[str] + """Name of the agent engine memory to delete.""" - config: Optional[GetMemoryBankOperationConfigDict] - """Used to override the default configuration.""" + config: Optional[DeleteAgentEngineMemoryConfigDict] + """""" -_GetMemoryBankOperationParametersOrDict = Union[ - _GetMemoryBankOperationParameters, _GetMemoryBankOperationParametersDict +_DeleteAgentEngineMemoryRequestParametersOrDict = Union[ + _DeleteAgentEngineMemoryRequestParameters, + _DeleteAgentEngineMemoryRequestParametersDict, ] -class MemoryConfig(_common.BaseModel): - """Config for creating a Memory.""" - - 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. +class DeleteAgentEngineMemoryOperation(_common.BaseModel): + """Operation for deleting agent engines.""" - The expiration time is computed: now + TTL.""", - ) - expire_time: Optional[datetime.datetime] = 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.""", - ) - 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.""", - ) - revision_ttl: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + 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}`.""", ) - disable_memory_revisions: Optional[bool] = Field( + metadata: Optional[dict[str, Any]] = 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="""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.""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + done: Optional[bool] = 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="""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.""", ) - memory_id: Optional[str] = Field( + error: Optional[dict[str, Any]] = 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.""", + description="""The error result of the operation in case of failure or cancellation.""", ) -class MemoryConfigDict(TypedDict, total=False): - """Config for creating a Memory.""" +class DeleteAgentEngineMemoryOperationDict(TypedDict, total=False): + """Operation for deleting agent engines.""" - 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. - The expiration time is computed: now + TTL.""" +DeleteAgentEngineMemoryOperationOrDict = Union[ + DeleteAgentEngineMemoryOperation, DeleteAgentEngineMemoryOperationDict +] - 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.""" - 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.""" +class GenerateMemoriesRequestVertexSessionSource(_common.BaseModel): + """The vertex session source for generating memories.""" - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + 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.""", + ) - disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revision will be created for this request.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics of the memory.""" +class GenerateMemoriesRequestVertexSessionSourceDict(TypedDict, total=False): + """The vertex session source for generating memories.""" - 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.""" + end_time: Optional[datetime.datetime] + """Optional. End time (exclusive) of the time range. If not set, the end 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.""" + 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.""" -MemoryConfigOrDict = Union[MemoryConfig, MemoryConfigDict] +GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ + GenerateMemoriesRequestVertexSessionSource, + GenerateMemoriesRequestVertexSessionSourceDict, +] -class _CreateMemoryRequestParameters(_common.BaseModel): - """Parameters for creating Memories.""" +class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): - name: Optional[str] = Field( + content: Optional[genai_types.Content] = Field( default=None, - description="""Name of the memory bank to create the memory under.""", + description="""Required. A single piece of content from which to generate memories.""", ) - 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. - 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="""""") +class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): + content: Optional[genai_types.Content] + """Required. A single piece of content from which to generate memories.""" -class _CreateMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for creating Memories.""" - name: Optional[str] - """Name of the memory bank to create the memory under.""" +GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ + GenerateMemoriesRequestDirectContentsSourceEvent, + GenerateMemoriesRequestDirectContentsSourceEventDict, +] - fact: Optional[str] - """The fact of the memory. - This is the semantic knowledge extracted from the source content).""" +class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): + """The direct contents source for generating memories.""" - scope: Optional[dict[str, str]] - """The scope of the memory. + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( + default=None, + description="""Required. The source content (i.e. chat history) to generate memories from.""", + ) - 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[MemoryConfigDict] - """""" +class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for generating memories.""" + + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] + """Required. The source content (i.e. chat history) to generate memories from.""" -_CreateMemoryRequestParametersOrDict = Union[ - _CreateMemoryRequestParameters, _CreateMemoryRequestParametersDict +GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ + GenerateMemoriesRequestDirectContentsSource, + GenerateMemoriesRequestDirectContentsSourceDict, ] -class MemoryStructuredContent(_common.BaseModel): - """Represents the structured value of the memory.""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): + """A direct memory to upload to Memory Bank.""" - data: Optional[dict[str, Any]] = Field( + fact: Optional[str] = Field( default=None, - description="""Required. Represents the structured value of the memory.""", + description="""Required. The fact to consolidate with existing memories.""", ) - schema_id: Optional[str] = Field( + topics: Optional[list[MemoryTopicId]] = Field( default=None, - description="""Required. Represents the schema ID for which this structured memory belongs to.""", + description="""Optional. The topics that the consolidated memories should be associated with.""", ) -class MemoryStructuredContentDict(TypedDict, total=False): - """Represents the structured value of the memory.""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( + TypedDict, total=False +): + """A direct memory to upload to Memory Bank.""" - data: Optional[dict[str, Any]] - """Required. Represents the structured value of the memory.""" + fact: Optional[str] + """Required. The fact to consolidate with existing memories.""" - schema_id: Optional[str] - """Required. Represents the schema ID for which this structured memory belongs to.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics that the consolidated memories should be associated with.""" -MemoryStructuredContentOrDict = Union[ - MemoryStructuredContent, MemoryStructuredContentDict +GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, + GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, ] -class Memory(_common.BaseModel): - """A memory.""" +class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): + """The direct memories source for generating memories.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp when this Memory was created.""", - ) - description: Optional[str] = Field( - default=None, - description="""Optional. Represents the description of the Memory.""", - ) - disable_memory_revisions: Optional[bool] = Field( - default=None, - description="""Optional. Input only. Indicates whether no revision will be created for this request.""", - ) - display_name: Optional[str] = Field( + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemory] + ] = Field( default=None, - description="""Optional. Represents the display name of the Memory.""", + description="""Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""", ) - 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.""", + + +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.""" ) - fact: Optional[str] = Field( + disable_consolidation: Optional[bool] = Field( default=None, - description="""Optional. Represents semantic knowledge extracted from the source content.""", + 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.""", ) - 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.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) - name: Optional[str] = Field( + revision_labels: Optional[dict[str, str]] = Field( default=None, - description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", + description="""Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""", ) revision_expire_time: Optional[datetime.datetime] = 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.""", - ) - 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.""", + description="""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] = Field( default=None, - description="""Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""", - ) - scope: Optional[dict[str, str]] = 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. The TTL for the revision. The expiration time is computed: now + TTL.""", ) - ttl: Optional[str] = Field( + disable_memory_revisions: Optional[bool] = Field( default=None, - description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", + description="""Optional. Input only. If true, no revisions will be created for this request.""", ) - update_time: Optional[datetime.datetime] = Field( + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( default=None, - description="""Output only. Represents the timestamp when this Memory was most recently updated.""", + description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", ) - memory_type: Optional[MemoryType] = Field( + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = 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. The strategy to use when applying metadata to existing memories.""", ) - structured_content: Optional[MemoryStructuredContent] = Field( + allowed_topics: Optional[list[MemoryTopicId]] = Field( default=None, - description="""Optional. Represents the structured content of the memory.""", + description="""Optional. Restricts memory generation to a subset of memory topics.""", ) -class MemoryDict(TypedDict, total=False): - """A memory.""" - - create_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was created.""" - - description: Optional[str] - """Optional. Represents the description of the Memory.""" - - disable_memory_revisions: Optional[bool] - """Optional. Input only. Indicates whether no revision will be created for this request.""" - - 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.""" +class GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False): + """Config for generating memories.""" - fact: Optional[str] - """Optional. Represents semantic knowledge extracted from the source content.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" + disable_consolidation: Optional[bool] + """Whether to disable consolidation of memories. - name: Optional[str] - """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{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.""" - 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.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" 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. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" + """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]] - """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 '*'.""" + 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.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. Represents the Topics of the Memory.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - ttl: Optional[str] - """Optional. Input only. Represents the TTL for this resource. 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.""" - update_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was most recently updated.""" + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" - 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.""" + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] + """Optional. The strategy to use when applying metadata to existing memories.""" - structured_content: Optional[MemoryStructuredContentDict] - """Optional. Represents the structured content of the memory.""" + allowed_topics: Optional[list[MemoryTopicIdDict]] + """Optional. Restricts memory generation to a subset of memory topics.""" -MemoryOrDict = Union[Memory, MemoryDict] +GenerateAgentEngineMemoriesConfigOrDict = Union[ + GenerateAgentEngineMemoriesConfig, GenerateAgentEngineMemoriesConfigDict +] -class MemoryOperation(_common.BaseModel): - """Operation that has a memory as a response.""" +class _GenerateAgentEngineMemoriesRequestParameters(_common.BaseModel): + """Parameters for generating agent engine memories.""" name: 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="""Name of the agent engine to generate memories for.""", ) - metadata: Optional[dict[str, Any]] = Field( + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = 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="""The vertex session source of the memories that should be generated.""", ) - done: Optional[bool] = 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.""", + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( + Field( + default=None, + description="""The direct contents source of the memories that should be generated.""", + ) ) - error: Optional[dict[str, Any]] = Field( + 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 error result of the operation in case of failure or cancellation.""", + 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 '*'.""", + ) + config: Optional[GenerateAgentEngineMemoriesConfig] = Field( + default=None, description="""""" ) - response: Optional[Memory] = Field(default=None, description="""The Memory.""") -class MemoryOperationDict(TypedDict, total=False): - """Operation that has a memory as a response.""" +class _GenerateAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for generating agent engine 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}`.""" + """Name of the agent engine to generate memories for.""" - 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.""" + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] + """The vertex session source 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.""" + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] + """The direct contents source of the memories that should be generated.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] + """The direct memories source of the memories that should be generated.""" - response: Optional[MemoryDict] - """The Memory.""" + scope: Optional[dict[str, str]] + """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 '*'.""" + config: Optional[GenerateAgentEngineMemoriesConfigDict] + """""" -MemoryOperationOrDict = Union[MemoryOperation, MemoryOperationDict] +_GenerateAgentEngineMemoriesRequestParametersOrDict = Union[ + _GenerateAgentEngineMemoriesRequestParameters, + _GenerateAgentEngineMemoriesRequestParametersDict, +] -class DeleteMemoryConfig(_common.BaseModel): - """Config for deleting a Memory.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" +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 DeleteMemoryConfigDict(TypedDict, total=False): - """Config for deleting a Memory.""" +class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): + """A memory that was generated.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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}`""" -DeleteMemoryConfigOrDict = Union[DeleteMemoryConfig, DeleteMemoryConfigDict] +GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ + GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict +] -class _DeleteMemoryRequestParameters(_common.BaseModel): - """Parameters for deleting memories.""" - name: Optional[str] = Field( - default=None, description="""Name of the memory to delete.""" - ) - config: Optional[DeleteMemoryConfig] = Field(default=None, description="""""") +class GenerateMemoriesResponse(_common.BaseModel): + """The response for generating memories.""" + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( + default=None, description="""The generated memories.""" + ) -class _DeleteMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for deleting memories.""" - name: Optional[str] - """Name of the memory to delete.""" +class GenerateMemoriesResponseDict(TypedDict, total=False): + """The response for generating memories.""" - config: Optional[DeleteMemoryConfigDict] - """""" + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] + """The generated memories.""" -_DeleteMemoryRequestParametersOrDict = Union[ - _DeleteMemoryRequestParameters, _DeleteMemoryRequestParametersDict +GenerateMemoriesResponseOrDict = Union[ + GenerateMemoriesResponse, GenerateMemoriesResponseDict ] -class DeleteMemoryOperation(_common.BaseModel): - """Operation for deleting memories.""" +class AgentEngineGenerateMemoriesOperation(_common.BaseModel): + """Operation that generates memories for an agent engine.""" name: Optional[str] = Field( default=None, @@ -11096,10 +11044,13 @@ class DeleteMemoryOperation(_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 DeleteMemoryOperationDict(TypedDict, total=False): - """Operation for deleting memories.""" +class AgentEngineGenerateMemoriesOperationDict(TypedDict, total=False): + """Operation that generates memories for an agent engine.""" 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}`.""" @@ -11113,163 +11064,129 @@ class DeleteMemoryOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - -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.""" + response: Optional[GenerateMemoriesResponseDict] + """The response for generating memories.""" -GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ - GenerateMemoriesRequestVertexSessionSource, - GenerateMemoriesRequestVertexSessionSourceDict, +AgentEngineGenerateMemoriesOperationOrDict = Union[ + AgentEngineGenerateMemoriesOperation, AgentEngineGenerateMemoriesOperationDict ] -class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): +class GetAgentEngineMemoryConfig(_common.BaseModel): + """Config for getting an Agent Engine Memory.""" - content: Optional[genai_types.Content] = Field( - default=None, - description="""Required. A single piece of content from which to generate memories.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): +class GetAgentEngineMemoryConfigDict(TypedDict, total=False): + """Config for getting an Agent Engine Memory.""" - content: Optional[genai_types.Content] - """Required. A single piece of content from which to generate memories.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ - GenerateMemoriesRequestDirectContentsSourceEvent, - GenerateMemoriesRequestDirectContentsSourceEventDict, +GetAgentEngineMemoryConfigOrDict = Union[ + GetAgentEngineMemoryConfig, GetAgentEngineMemoryConfigDict ] -class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): - """The direct contents source for generating memories.""" +class _GetAgentEngineMemoryRequestParameters(_common.BaseModel): + """Parameters for getting an agent engine.""" - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( - default=None, - description="""Required. The source content (i.e. chat history) to generate memories from.""", + name: Optional[str] = Field( + default=None, description="""Name of the agent engine.""" + ) + config: Optional[GetAgentEngineMemoryConfig] = Field( + default=None, description="""""" ) -class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for generating memories.""" +class _GetAgentEngineMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent engine.""" - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] - """Required. The source content (i.e. chat history) to generate memories from.""" + name: Optional[str] + """Name of the agent engine.""" + + config: Optional[GetAgentEngineMemoryConfigDict] + """""" -GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ - GenerateMemoriesRequestDirectContentsSource, - GenerateMemoriesRequestDirectContentsSourceDict, +_GetAgentEngineMemoryRequestParametersOrDict = Union[ + _GetAgentEngineMemoryRequestParameters, _GetAgentEngineMemoryRequestParametersDict ] -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): - """A direct memory to upload to Memory Bank.""" +class IngestionDirectContentsSourceEvent(_common.BaseModel): + """The direct contents source event for ingesting events.""" - fact: Optional[str] = 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. The fact to consolidate with existing memories.""", + 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.""", ) - topics: Optional[list[MemoryTopicId]] = Field( + event_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. The topics that the consolidated memories should be associated with.""", + 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 GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( - TypedDict, total=False -): - """A direct memory to upload to Memory Bank.""" +class IngestionDirectContentsSourceEventDict(TypedDict, total=False): + """The direct contents source event for ingesting events.""" - fact: Optional[str] - """Required. The fact to consolidate with existing memories.""" + content: Optional[genai_types.Content] + """Required. The content of the event.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics that the consolidated memories should be associated with.""" + 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.""" -GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, - GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, -] +IngestionDirectContentsSourceEventOrDict = Union[ + IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict +] -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 IngestionDirectContentsSource(_common.BaseModel): + """The direct contents source for ingesting events.""" + + events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( + default=None, description="""Required. The events to ingest.""" ) -class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): - """The direct memories source for generating memories.""" +class IngestionDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for ingesting events.""" - direct_memories: Optional[ - list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] - ] - """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" + events: Optional[list[IngestionDirectContentsSourceEventDict]] + """Required. The events to ingest.""" -GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSource, - GenerateMemoriesRequestDirectMemoriesSourceDict, +IngestionDirectContentsSourceOrDict = Union[ + IngestionDirectContentsSource, IngestionDirectContentsSourceDict ] -class GenerateMemoriesConfig(_common.BaseModel): - """Config for generating memories.""" +class IngestEventsConfig(_common.BaseModel): + """Config for ingesting events.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - disable_consolidation: 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.""", - ) wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + default=False, + description="""Waits for the underlying memory generation operation to complete + before returning. Defaults to false.""", + ) + force_flush: 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`.""", ) revision_labels: Optional[dict[str, str]] = Field( default=None, @@ -11295,28 +11212,20 @@ class GenerateMemoriesConfig(_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 GenerateMemoriesConfigDict(TypedDict, total=False): - """Config for generating memories.""" +class IngestEventsConfigDict(TypedDict, total=False): + """Config for ingesting events.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - disable_consolidation: Optional[bool] - """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] - """Waits for the operation to complete before returning.""" + """Waits for the underlying memory generation operation to complete + before returning. Defaults to false.""" + + 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`.""" 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.""" @@ -11336,143 +11245,69 @@ class GenerateMemoriesConfigDict(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.""" - -GenerateMemoriesConfigOrDict = Union[GenerateMemoriesConfig, GenerateMemoriesConfigDict] +IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] -class _GenerateMemoriesRequestParameters(_common.BaseModel): - """Parameters for generating memories.""" +class _IngestEventsRequestParameters(_common.BaseModel): + """Parameters for purging agent engine memories.""" name: Optional[str] = Field( - default=None, - description="""Name of the Memory Bank to generate memories with.""", - ) - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( - default=None, - description="""The vertex session source of the memories that should be generated.""", + default=None, description="""Name of the Agent Engine to ingest events into.""" ) - direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( - Field( - default=None, - description="""The direct contents source of the memories that should be generated.""", - ) + stream_id: Optional[str] = Field( + default=None, description="""The ID of the stream to ingest events into.""" ) - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( - Field( - default=None, - description="""The direct memories source of the memories that should be generated.""", - ) + 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. + description="""The scope of the memories that should be generated from the stream. - 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 + Memories will be consolidated across memories with the same scope. Scope values cannot contain the wildcard character '*'.""", ) - config: Optional[GenerateMemoriesConfig] = Field(default=None, description="""""") + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""The configuration for the memory generation trigger.""", + ) + config: Optional[IngestEventsConfig] = Field(default=None, description="""""") -class _GenerateMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for generating memories.""" +class _IngestEventsRequestParametersDict(TypedDict, total=False): + """Parameters for purging agent engine memories.""" name: Optional[str] - """Name of the Memory Bank to generate memories with.""" + """Name of the Agent Engine to ingest events into.""" - 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.""" + stream_id: Optional[str] + """The ID of the stream to ingest events into.""" - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] - """The direct memories source of the memories that should be generated.""" + direct_contents_source: Optional[IngestionDirectContentsSourceDict] + """The direct memories source of the events that should be ingested.""" scope: Optional[dict[str, str]] - """The scope of the memories that should be generated. + """The scope of the memories that should be generated from the stream. - 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 + Memories will be consolidated across memories with the same scope. Scope values cannot contain the wildcard character '*'.""" - config: Optional[GenerateMemoriesConfigDict] - """""" - - -_GenerateMemoriesRequestParametersOrDict = Union[ - _GenerateMemoriesRequestParameters, _GenerateMemoriesRequestParametersDict -] - - -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.""" + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """The configuration for the memory generation trigger.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] - """The generated memories.""" + config: Optional[IngestEventsConfigDict] + """""" -GenerateMemoriesResponseOrDict = Union[ - GenerateMemoriesResponse, GenerateMemoriesResponseDict +_IngestEventsRequestParametersOrDict = Union[ + _IngestEventsRequestParameters, _IngestEventsRequestParametersDict ] -class GenerateMemoriesOperation(_common.BaseModel): - """Operation that generates memories with a Memory Bank.""" +class MemoryBankIngestEventsOperation(_common.BaseModel): + """Operation that ingests events into a memory bank.""" name: Optional[str] = Field( default=None, @@ -11490,13 +11325,10 @@ class GenerateMemoriesOperation(_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 GenerateMemoriesOperationDict(TypedDict, total=False): - """Operation that generates memories with a Memory Bank.""" +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}`.""" @@ -11510,57 +11342,14 @@ class GenerateMemoriesOperationDict(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.""" - - -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 +MemoryBankIngestEventsOperationOrDict = Union[ + MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict ] -class ListMemoriesConfig(_common.BaseModel): - """Config for listing memories.""" +class ListAgentEngineMemoryConfig(_common.BaseModel): + """Config for listing agent engine memories.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -11585,8 +11374,8 @@ class ListMemoriesConfig(_common.BaseModel): ) -class ListMemoriesConfigDict(TypedDict, total=False): - """Config for listing memories.""" +class ListAgentEngineMemoryConfigDict(TypedDict, total=False): + """Config for listing agent engine memories.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -11612,47 +11401,51 @@ class ListMemoriesConfigDict(TypedDict, total=False): * `update_time`""" -ListMemoriesConfigOrDict = Union[ListMemoriesConfig, ListMemoriesConfigDict] +ListAgentEngineMemoryConfigOrDict = Union[ + ListAgentEngineMemoryConfig, ListAgentEngineMemoryConfigDict +] -class _ListMemoriesRequestParameters(_common.BaseModel): - """Parameters for listing memories.""" +class _ListAgentEngineMemoryRequestParameters(_common.BaseModel): + """Parameters for listing agent engines.""" name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank.""" + default=None, description="""Name of the agent engine.""" + ) + config: Optional[ListAgentEngineMemoryConfig] = Field( + default=None, description="""""" ) - config: Optional[ListMemoriesConfig] = Field(default=None, description="""""") -class _ListMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for listing memories.""" +class _ListAgentEngineMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent engines.""" name: Optional[str] - """Name of the Memory Bank.""" + """Name of the agent engine.""" - config: Optional[ListMemoriesConfigDict] + config: Optional[ListAgentEngineMemoryConfigDict] """""" -_ListMemoriesRequestParametersOrDict = Union[ - _ListMemoriesRequestParameters, _ListMemoriesRequestParametersDict +_ListAgentEngineMemoryRequestParametersOrDict = Union[ + _ListAgentEngineMemoryRequestParameters, _ListAgentEngineMemoryRequestParametersDict ] -class ListMemoriesResponse(_common.BaseModel): - """Response for listing memories.""" +class ListReasoningEnginesMemoriesResponse(_common.BaseModel): + """Response for listing agent engine 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 memories.""" + default=None, description="""List of agent engine memories.""" ) -class ListMemoriesResponseDict(TypedDict, total=False): - """Response for listing memories.""" +class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): + """Response for listing agent engine memories.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -11661,61 +11454,65 @@ class ListMemoriesResponseDict(TypedDict, total=False): """""" memories: Optional[list[MemoryDict]] - """List of memories.""" + """List of agent engine memories.""" -ListMemoriesResponseOrDict = Union[ListMemoriesResponse, ListMemoriesResponseDict] +ListReasoningEnginesMemoriesResponseOrDict = Union[ + ListReasoningEnginesMemoriesResponse, ListReasoningEnginesMemoriesResponseDict +] -class _GetMemoryOperationParameters(_common.BaseModel): +class _GetAgentEngineMemoryOperationParameters(_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[GetMemoryBankOperationConfig] = Field( + config: Optional[GetAgentEngineOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetMemoryOperationParametersDict(TypedDict, total=False): +class _GetAgentEngineMemoryOperationParametersDict(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[GetMemoryBankOperationConfigDict] + config: Optional[GetAgentEngineOperationConfigDict] """Used to override the default configuration.""" -_GetMemoryOperationParametersOrDict = Union[ - _GetMemoryOperationParameters, _GetMemoryOperationParametersDict +_GetAgentEngineMemoryOperationParametersOrDict = Union[ + _GetAgentEngineMemoryOperationParameters, + _GetAgentEngineMemoryOperationParametersDict, ] -class _GetGenerateMemoriesOperationParameters(_common.BaseModel): +class _GetAgentEngineGenerateMemoriesOperationParameters(_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[GetMemoryBankOperationConfig] = Field( + config: Optional[GetAgentEngineOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetGenerateMemoriesOperationParametersDict(TypedDict, total=False): +class _GetAgentEngineGenerateMemoriesOperationParametersDict(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[GetMemoryBankOperationConfigDict] + config: Optional[GetAgentEngineOperationConfigDict] """Used to override the default configuration.""" -_GetGenerateMemoriesOperationParametersOrDict = Union[ - _GetGenerateMemoriesOperationParameters, _GetGenerateMemoriesOperationParametersDict +_GetAgentEngineGenerateMemoriesOperationParametersOrDict = Union[ + _GetAgentEngineGenerateMemoriesOperationParameters, + _GetAgentEngineGenerateMemoriesOperationParametersDict, ] @@ -11836,7 +11633,7 @@ class MemoryConjunctionFilterDict(TypedDict, total=False): ] -class RetrieveMemoriesConfig(_common.BaseModel): +class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11879,7 +11676,7 @@ class RetrieveMemoriesConfig(_common.BaseModel): ) -class RetrieveMemoriesConfigDict(TypedDict, total=False): +class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] @@ -11918,15 +11715,17 @@ class RetrieveMemoriesConfigDict(TypedDict, total=False): retrieve memories matching any of the specified `MemoryType` values.""" -RetrieveMemoriesConfigOrDict = Union[RetrieveMemoriesConfig, RetrieveMemoriesConfigDict] +RetrieveAgentEngineMemoriesConfigOrDict = Union[ + RetrieveAgentEngineMemoriesConfig, RetrieveAgentEngineMemoriesConfigDict +] -class _RetrieveMemoriesRequestParameters(_common.BaseModel): - """Parameters for retrieving memories.""" +class _RetrieveAgentEngineMemoriesRequestParameters(_common.BaseModel): + """Parameters for retrieving agent engine memories.""" name: Optional[str] = Field( default=None, - description="""Name of the Memory Bank to retrieve memories from.""", + description="""Name of the agent engine to retrieve memories from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -11948,14 +11747,16 @@ class _RetrieveMemoriesRequestParameters(_common.BaseModel): description="""Parameters for simple (non-similarity search) retrieval.""", ) ) - config: Optional[RetrieveMemoriesConfig] = Field(default=None, description="""""") + config: Optional[RetrieveAgentEngineMemoriesConfig] = Field( + default=None, description="""""" + ) -class _RetrieveMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving memories.""" +class _RetrieveAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving agent engine memories.""" name: Optional[str] - """Name of the Memory Bank to retrieve memories from.""" + """Name of the agent engine to retrieve memories from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -11972,12 +11773,13 @@ class _RetrieveMemoriesRequestParametersDict(TypedDict, total=False): simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] """Parameters for simple (non-similarity search) retrieval.""" - config: Optional[RetrieveMemoriesConfigDict] + config: Optional[RetrieveAgentEngineMemoriesConfigDict] """""" -_RetrieveMemoriesRequestParametersOrDict = Union[ - _RetrieveMemoriesRequestParameters, _RetrieveMemoriesRequestParametersDict +_RetrieveAgentEngineMemoriesRequestParametersOrDict = Union[ + _RetrieveAgentEngineMemoriesRequestParameters, + _RetrieveAgentEngineMemoriesRequestParametersDict, ] @@ -12056,11 +11858,11 @@ class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): - """Parameters for retrieving memory profiles.""" + """Parameters for retrieving agent engine memory profiles.""" name: Optional[str] = Field( default=None, - description="""Name of the Memory Bank to retrieve memory profiles from.""", + description="""Name of the agent engine to retrieve memory profiles from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -12076,10 +11878,10 @@ class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving memory profiles.""" + """Parameters for retrieving agent engine memory profiles.""" name: Optional[str] - """Name of the Memory Bank to retrieve memory profiles from.""" + """Name of the agent engine to retrieve memory profiles from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -12150,7 +11952,7 @@ class RetrieveProfilesResponseDict(TypedDict, total=False): ] -class RollbackMemoryConfig(_common.BaseModel): +class RollbackAgentEngineMemoryConfig(_common.BaseModel): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -12162,7 +11964,7 @@ class RollbackMemoryConfig(_common.BaseModel): ) -class RollbackMemoryConfigDict(TypedDict, total=False): +class RollbackAgentEngineMemoryConfigDict(TypedDict, total=False): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] @@ -12172,40 +11974,45 @@ class RollbackMemoryConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -RollbackMemoryConfigOrDict = Union[RollbackMemoryConfig, RollbackMemoryConfigDict] +RollbackAgentEngineMemoryConfigOrDict = Union[ + RollbackAgentEngineMemoryConfig, RollbackAgentEngineMemoryConfigDict +] -class _RollbackMemoryRequestParameters(_common.BaseModel): - """Parameters for generating memories.""" +class _RollbackAgentEngineMemoryRequestParameters(_common.BaseModel): + """Parameters for generating agent engine memories.""" name: Optional[str] = Field( - default=None, description="""Name of the memory to rollback.""" + default=None, description="""Name of the agent engine memory to rollback.""" ) target_revision_id: Optional[str] = Field( default=None, description="""The ID of the revision to rollback to.""" ) - config: Optional[RollbackMemoryConfig] = Field(default=None, description="""""") + config: Optional[RollbackAgentEngineMemoryConfig] = Field( + default=None, description="""""" + ) -class _RollbackMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for generating memories.""" +class _RollbackAgentEngineMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for generating agent engine memories.""" name: Optional[str] - """Name of the memory to rollback.""" + """Name of the agent engine memory to rollback.""" target_revision_id: Optional[str] """The ID of the revision to rollback to.""" - config: Optional[RollbackMemoryConfigDict] + config: Optional[RollbackAgentEngineMemoryConfigDict] """""" -_RollbackMemoryRequestParametersOrDict = Union[ - _RollbackMemoryRequestParameters, _RollbackMemoryRequestParametersDict +_RollbackAgentEngineMemoryRequestParametersOrDict = Union[ + _RollbackAgentEngineMemoryRequestParameters, + _RollbackAgentEngineMemoryRequestParametersDict, ] -class RollbackMemoryOperation(_common.BaseModel): +class AgentEngineRollbackMemoryOperation(_common.BaseModel): """Operation that rolls back a memory.""" name: Optional[str] = Field( @@ -12226,7 +12033,7 @@ class RollbackMemoryOperation(_common.BaseModel): ) -class RollbackMemoryOperationDict(TypedDict, total=False): +class AgentEngineRollbackMemoryOperationDict(TypedDict, total=False): """Operation that rolls back a memory.""" name: Optional[str] @@ -12242,13 +12049,13 @@ class RollbackMemoryOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -RollbackMemoryOperationOrDict = Union[ - RollbackMemoryOperation, RollbackMemoryOperationDict +AgentEngineRollbackMemoryOperationOrDict = Union[ + AgentEngineRollbackMemoryOperation, AgentEngineRollbackMemoryOperationDict ] -class UpdateMemoryConfig(_common.BaseModel): - """Config for updating a memory.""" +class UpdateAgentEngineMemoryConfig(_common.BaseModel): + """Config for updating agent engine memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -12303,8 +12110,8 @@ class UpdateMemoryConfig(_common.BaseModel): ) -class UpdateMemoryConfigDict(TypedDict, total=False): - """Config for updating a memory.""" +class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): + """Config for updating agent engine memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -12349,14 +12156,16 @@ class UpdateMemoryConfigDict(TypedDict, total=False): https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" -UpdateMemoryConfigOrDict = Union[UpdateMemoryConfig, UpdateMemoryConfigDict] +UpdateAgentEngineMemoryConfigOrDict = Union[ + UpdateAgentEngineMemoryConfig, UpdateAgentEngineMemoryConfigDict +] -class _UpdateMemoryRequestParameters(_common.BaseModel): - """Parameters for updating memories.""" +class _UpdateAgentEngineMemoryRequestParameters(_common.BaseModel): + """Parameters for updating agent engine memories.""" name: Optional[str] = Field( - default=None, description="""Name of the memory to update.""" + default=None, description="""Name of the agent engine memory to update.""" ) fact: Optional[str] = Field( default=None, @@ -12372,14 +12181,16 @@ class _UpdateMemoryRequestParameters(_common.BaseModel): creating or generating memories. Up to 5 key-value pairs are accepted, and scope values cannot contain the wildcard character '*'.""", ) - config: Optional[UpdateMemoryConfig] = Field(default=None, description="""""") + config: Optional[UpdateAgentEngineMemoryConfig] = Field( + default=None, description="""""" + ) -class _UpdateMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for updating memories.""" +class _UpdateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent engine memories.""" name: Optional[str] - """Name of the memory to update.""" + """Name of the agent engine memory to update.""" fact: Optional[str] """The updated fact of the memory. @@ -12393,16 +12204,17 @@ class _UpdateMemoryRequestParametersDict(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[UpdateMemoryConfigDict] + config: Optional[UpdateAgentEngineMemoryConfigDict] """""" -_UpdateMemoryRequestParametersOrDict = Union[ - _UpdateMemoryRequestParameters, _UpdateMemoryRequestParametersDict +_UpdateAgentEngineMemoryRequestParametersOrDict = Union[ + _UpdateAgentEngineMemoryRequestParameters, + _UpdateAgentEngineMemoryRequestParametersDict, ] -class PurgeMemoriesConfig(_common.BaseModel): +class PurgeAgentEngineMemoriesConfig(_common.BaseModel): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -12414,7 +12226,7 @@ class PurgeMemoriesConfig(_common.BaseModel): ) -class PurgeMemoriesConfigDict(TypedDict, total=False): +class PurgeAgentEngineMemoriesConfigDict(TypedDict, total=False): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] @@ -12424,14 +12236,16 @@ class PurgeMemoriesConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -PurgeMemoriesConfigOrDict = Union[PurgeMemoriesConfig, PurgeMemoriesConfigDict] +PurgeAgentEngineMemoriesConfigOrDict = Union[ + PurgeAgentEngineMemoriesConfig, PurgeAgentEngineMemoriesConfigDict +] -class _PurgeMemoriesRequestParameters(_common.BaseModel): - """Parameters for purging memories.""" +class _PurgeAgentEngineMemoriesRequestParameters(_common.BaseModel): + """Parameters for purging agent engine memories.""" name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank to purge memories from.""" + default=None, description="""Name of the Agent Engine to purge memories from.""" ) filter: Optional[str] = Field( default=None, @@ -12459,14 +12273,16 @@ class _PurgeMemoriesRequestParameters(_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[PurgeMemoriesConfig] = Field(default=None, description="""""") + config: Optional[PurgeAgentEngineMemoriesConfig] = Field( + default=None, description="""""" + ) -class _PurgeMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for purging memories.""" +class _PurgeAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for purging agent engine memories.""" name: Optional[str] - """Name of the Memory Bank to purge memories from.""" + """Name of the Agent Engine to purge memories from.""" filter: Optional[str] """The standard list filter to determine which memories to purge. @@ -12491,12 +12307,13 @@ class _PurgeMemoriesRequestParametersDict(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[PurgeMemoriesConfigDict] + config: Optional[PurgeAgentEngineMemoriesConfigDict] """""" -_PurgeMemoriesRequestParametersOrDict = Union[ - _PurgeMemoriesRequestParameters, _PurgeMemoriesRequestParametersDict +_PurgeAgentEngineMemoriesRequestParametersOrDict = Union[ + _PurgeAgentEngineMemoriesRequestParameters, + _PurgeAgentEngineMemoriesRequestParametersDict, ] @@ -12518,8 +12335,8 @@ class PurgeMemoriesResponseDict(TypedDict, total=False): PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] -class PurgeMemoriesOperation(_common.BaseModel): - """Operation that purges memories from a Memory Bank.""" +class AgentEnginePurgeMemoriesOperation(_common.BaseModel): + """Operation that purges memories from an agent engine.""" name: Optional[str] = Field( default=None, @@ -12542,8 +12359,8 @@ class PurgeMemoriesOperation(_common.BaseModel): ) -class PurgeMemoriesOperationDict(TypedDict, total=False): - """Operation that purges memories from a Memory Bank.""" +class AgentEnginePurgeMemoriesOperationDict(TypedDict, total=False): + """Operation that purges memories from an agent engine.""" 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}`.""" @@ -12561,50 +12378,55 @@ class PurgeMemoriesOperationDict(TypedDict, total=False): """The response for purging memories.""" -PurgeMemoriesOperationOrDict = Union[PurgeMemoriesOperation, PurgeMemoriesOperationDict] +AgentEnginePurgeMemoriesOperationOrDict = Union[ + AgentEnginePurgeMemoriesOperation, AgentEnginePurgeMemoriesOperationDict +] -class GetMemoryRevisionConfig(_common.BaseModel): - """Config for getting a Memory Revision.""" +class GetAgentEngineMemoryRevisionConfig(_common.BaseModel): + """Config for getting an Agent Engine Memory Revision.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetMemoryRevisionConfigDict(TypedDict, total=False): - """Config for getting a Memory Revision.""" +class GetAgentEngineMemoryRevisionConfigDict(TypedDict, total=False): + """Config for getting an Agent Engine Memory Revision.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetMemoryRevisionConfigOrDict = Union[ - GetMemoryRevisionConfig, GetMemoryRevisionConfigDict +GetAgentEngineMemoryRevisionConfigOrDict = Union[ + GetAgentEngineMemoryRevisionConfig, GetAgentEngineMemoryRevisionConfigDict ] -class _GetMemoryRevisionRequestParameters(_common.BaseModel): - """Parameters for getting a memory revision.""" +class _GetAgentEngineMemoryRevisionRequestParameters(_common.BaseModel): + """Parameters for getting an Agent Engine memory revision.""" name: Optional[str] = Field( - default=None, description="""Name of the Memory Revision.""" + default=None, description="""Name of the agent engine.""" + ) + config: Optional[GetAgentEngineMemoryRevisionConfig] = Field( + default=None, description="""""" ) - config: Optional[GetMemoryRevisionConfig] = Field(default=None, description="""""") -class _GetMemoryRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting a memory revision.""" +class _GetAgentEngineMemoryRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an Agent Engine memory revision.""" name: Optional[str] - """Name of the Memory Revision.""" + """Name of the agent engine.""" - config: Optional[GetMemoryRevisionConfigDict] + config: Optional[GetAgentEngineMemoryRevisionConfigDict] """""" -_GetMemoryRevisionRequestParametersOrDict = Union[ - _GetMemoryRevisionRequestParameters, _GetMemoryRevisionRequestParametersDict +_GetAgentEngineMemoryRevisionRequestParametersOrDict = Union[ + _GetAgentEngineMemoryRevisionRequestParameters, + _GetAgentEngineMemoryRevisionRequestParametersDict, ] @@ -12704,8 +12526,8 @@ class MemoryRevisionDict(TypedDict, total=False): MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] -class ListMemoryRevisionsConfig(_common.BaseModel): - """Config for listing memory revisions.""" +class ListAgentEngineMemoryRevisionsConfig(_common.BaseModel): + """Config for listing Agent Engine memory revisions.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -12719,8 +12541,8 @@ class ListMemoryRevisionsConfig(_common.BaseModel): ) -class ListMemoryRevisionsConfigDict(TypedDict, total=False): - """Config for listing memory revisions.""" +class ListAgentEngineMemoryRevisionsConfigDict(TypedDict, total=False): + """Config for listing Agent Engine memory revisions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -12736,37 +12558,40 @@ class ListMemoryRevisionsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListMemoryRevisionsConfigOrDict = Union[ - ListMemoryRevisionsConfig, ListMemoryRevisionsConfigDict +ListAgentEngineMemoryRevisionsConfigOrDict = Union[ + ListAgentEngineMemoryRevisionsConfig, ListAgentEngineMemoryRevisionsConfigDict ] -class _ListMemoryRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing memory revisions.""" +class _ListAgentEngineMemoryRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing Agent Engine memory revisions.""" - name: Optional[str] = Field(default=None, description="""Name of the memory""") - config: Optional[ListMemoryRevisionsConfig] = Field( + name: Optional[str] = Field( + default=None, description="""Name of the Agent Engine memory""" + ) + config: Optional[ListAgentEngineMemoryRevisionsConfig] = Field( default=None, description="""""" ) -class _ListMemoryRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing memory revisions.""" +class _ListAgentEngineMemoryRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing Agent Engine memory revisions.""" name: Optional[str] - """Name of the memory""" + """Name of the Agent Engine memory""" - config: Optional[ListMemoryRevisionsConfigDict] + config: Optional[ListAgentEngineMemoryRevisionsConfigDict] """""" -_ListMemoryRevisionsRequestParametersOrDict = Union[ - _ListMemoryRevisionsRequestParameters, _ListMemoryRevisionsRequestParametersDict +_ListAgentEngineMemoryRevisionsRequestParametersOrDict = Union[ + _ListAgentEngineMemoryRevisionsRequestParameters, + _ListAgentEngineMemoryRevisionsRequestParametersDict, ] -class ListMemoryRevisionsResponse(_common.BaseModel): - """Response for listing memory revisions.""" +class ListAgentEngineMemoryRevisionsResponse(_common.BaseModel): + """Response for listing agent engine memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" @@ -12777,8 +12602,8 @@ class ListMemoryRevisionsResponse(_common.BaseModel): ) -class ListMemoryRevisionsResponseDict(TypedDict, total=False): - """Response for listing memory revisions.""" +class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): + """Response for listing agent engine memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -12790,8 +12615,8 @@ class ListMemoryRevisionsResponseDict(TypedDict, total=False): """List of memory revisions.""" -ListMemoryRevisionsResponseOrDict = Union[ - ListMemoryRevisionsResponse, ListMemoryRevisionsResponseDict +ListAgentEngineMemoryRevisionsResponseOrDict = Union[ + ListAgentEngineMemoryRevisionsResponse, ListAgentEngineMemoryRevisionsResponseDict ] diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py new file mode 100644 index 0000000000..e4b33ad2d5 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py @@ -0,0 +1,32 @@ +# 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 new file mode 100644 index 0000000000..bca6fca8c9 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py @@ -0,0 +1,32 @@ +# 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_memories_private_create.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py similarity index 70% rename from tests/unit/agentplatform/genai/replays/test_memories_private_create.py rename to tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py index ed2904d9e4..5ca959abe2 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py @@ -19,16 +19,12 @@ def test_private_create_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) - finally: - client.memory_banks.delete(name=memory_bank.name, force=True) + 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) pytestmark = pytest_helper.setup( 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 new file mode 100644 index 0000000000..7b294a01cf --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py @@ -0,0 +1,35 @@ +# 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 new file mode 100644 index 0000000000..bdcd2f14a4 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py @@ -0,0 +1,32 @@ +# 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_memories_private_get_memory_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py similarity index 58% rename from tests/unit/agentplatform/genai/replays/test_memories_private_get_memory_operation.py rename to tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py index e31342b88f..23c92f7821 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_get_memory_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py @@ -19,24 +19,14 @@ def test_private_get_memory_operation(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"}, - ) - 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) + 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) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories._get_memory_operation", + test_method="agent_engines.memories._get_memory_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py similarity index 61% rename from tests/unit/agentplatform/genai/replays/test_memories_private_list.py rename to tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py index 3837f25813..662ad020fb 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py @@ -19,22 +19,14 @@ def test_private_list_memory(client): - 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) + 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) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories._list", + test_method="agent_engines.memories._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_private_purge.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py similarity index 66% rename from tests/unit/agentplatform/genai/replays/test_memories_private_purge.py rename to tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py index 92cca4b615..d77cb3a795 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_purge.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py @@ -19,21 +19,17 @@ def test_private_purge(client): - 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) + 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 + ) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories._purge", + test_method="agent_engines.memories._purge", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py new file mode 100644 index 0000000000..4098b11d9b --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py @@ -0,0 +1,38 @@ +# 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_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, + ) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories._retrieve", +) 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 new file mode 100644 index 0000000000..85329da47e --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py @@ -0,0 +1,33 @@ +# 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_memories_private_update.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py similarity index 56% rename from tests/unit/agentplatform/genai/replays/test_memories_private_update.py rename to tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py index 69cdbe5643..c6dbc1f095 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_update.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py @@ -19,26 +19,17 @@ def test_private_update_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"}, - ) - 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) + 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) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories._update", + test_method="agent_engines.memories._update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py new file mode 100644 index 0000000000..f58e365fe7 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py @@ -0,0 +1,32 @@ +# 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_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 + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="agent_engines.memories.revisions.get", +) diff --git a/tests/unit/agentplatform/genai/replays/test_memory_revisions_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py similarity index 58% rename from tests/unit/agentplatform/genai/replays/test_memory_revisions_private_list.py rename to tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py index 7ade5968a0..c827a57a06 100644 --- a/tests/unit/agentplatform/genai/replays/test_memory_revisions_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py @@ -19,23 +19,14 @@ def test_private_list_memory_revisions(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_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) + 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) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories.revisions._list", + test_method="agent_engines.memories.revisions._list", ) 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 new file mode 100644 index 0000000000..f321a921c0 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py @@ -0,0 +1,118 @@ +# 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 new file mode 100644 index 0000000000..9d3dea4c72 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py @@ -0,0 +1,62 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +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 new file mode 100644 index 0000000000..4942e3ac9e --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py @@ -0,0 +1,381 @@ +# 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_memories_get.py b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py similarity index 55% rename from tests/unit/agentplatform/genai/replays/test_memories_get.py rename to tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py index 1d523c88c7..3ec3093562 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_get.py +++ b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py @@ -16,32 +16,30 @@ import pytest -from agentplatform._genai import types from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types def test_get_memory(client): - 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) + 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) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories.get", + test_method="agent_engines.memories.get", ) @@ -50,18 +48,18 @@ def test_get_memory(client): @pytest.mark.asyncio async def test_get_memory_async(client): - memory_bank = client.memory_banks.create() - operation = await client.aio.memory_banks.memories.create( - name=memory_bank.name, + 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.MemoryOperation) - memory = await client.aio.memory_banks.memories.get( + assert isinstance(operation, types.AgentEngineMemoryOperation) + memory = await client.aio.agent_engines.memories.get( name=operation.response.name, ) assert isinstance(memory, types.Memory) assert memory.name == operation.response.name - await client.aio.memory_banks.delete( - name=memory_bank.name, force=True + await client.aio.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 new file mode 100644 index 0000000000..ff7ad46583 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py @@ -0,0 +1,118 @@ +# 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 new file mode 100644 index 0000000000..4373b4bf25 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py @@ -0,0 +1,96 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +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 deleted file mode 100644 index a820a2f012..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_create.py +++ /dev/null @@ -1,121 +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): - 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 deleted file mode 100644 index 18f31d4991..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_delete.py +++ /dev/null @@ -1,67 +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): - 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 deleted file mode 100644 index 5646cedcf5..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_generate.py +++ /dev/null @@ -1,394 +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): - 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_memories_list.py b/tests/unit/agentplatform/genai/replays/test_memories_list.py deleted file mode 100644 index 54f7e2b825..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_list.py +++ /dev/null @@ -1,99 +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): - 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_memories_private_generate.py b/tests/unit/agentplatform/genai/replays/test_memories_private_generate.py deleted file mode 100644 index 1e30f82466..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_generate.py +++ /dev/null @@ -1,56 +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 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 deleted file mode 100644 index 4c642b4203..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_get_generate_memories_operation.py +++ /dev/null @@ -1,55 +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 -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_memories_private_retrieve.py b/tests/unit/agentplatform/genai/replays/test_memories_private_retrieve.py deleted file mode 100644 index d0c72671b0..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_retrieve.py +++ /dev/null @@ -1,46 +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_retrieve(client): - 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="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 deleted file mode 100644 index ddf157e6e3..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_rollback.py +++ /dev/null @@ -1,52 +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): - 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_memories_retrieve.py b/tests/unit/agentplatform/genai/replays/test_memories_retrieve.py deleted file mode 100644 index 056115999d..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memories_retrieve.py +++ /dev/null @@ -1,205 +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): - 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_memory_banks_ingest_events.py b/tests/unit/agentplatform/genai/replays/test_memory_banks_ingest_events.py deleted file mode 100644 index fc74f5ecc3..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memory_banks_ingest_events.py +++ /dev/null @@ -1,123 +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 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_memory_revisions_get.py b/tests/unit/agentplatform/genai/replays/test_memory_revisions_get.py deleted file mode 100644 index 0058298362..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_memory_revisions_get.py +++ /dev/null @@ -1,45 +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_revisions(client): - 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="memory_banks.memories.revisions.get", -) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_purge.py b/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py similarity index 66% rename from tests/unit/agentplatform/genai/replays/test_memories_purge.py rename to tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py index e2ac58dde9..e123597953 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_purge.py +++ b/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py @@ -22,28 +22,28 @@ def test_purge_memories(client): """Tests purging memories.""" - memory_bank = client.memory_banks.create() + agent_engine = client.agent_engines.create() try: - client.memory_banks.memories.create( - name=memory_bank.name, + client.agent_engines.memories.create( + name=agent_engine.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.memory_banks.memories.create( - name=memory_bank.name, + client.agent_engines.memories.create( + name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.memory_banks.memories.create( - name=memory_bank.name, + client.agent_engines.memories.create( + name=agent_engine.api_resource.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - client.memory_banks.memories.create( - name=memory_bank.name, + client.agent_engines.memories.create( + name=agent_engine.api_resource.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.memory_banks.memories.purge( - name=memory_bank.name, + operation = client.agent_engines.memories.purge( + name=agent_engine.api_resource.name, filter="scope.user_id=123", config={"wait_for_completion": True}, ) @@ -62,16 +62,16 @@ def test_purge_memories(client): assert ( len( list( - client.memory_banks.memories.list( - name=memory_bank.name + client.agent_engines.memories.list( + name=agent_engine.api_resource.name ) ) ) == 4 ) # Now, actually purge the memories. - operation = client.memory_banks.memories.purge( - name=memory_bank.name, + operation = client.agent_engines.memories.purge( + name=agent_engine.api_resource.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.memory_banks.memories.list( - name=memory_bank.name + client.agent_engines.memories.list( + name=agent_engine.api_resource.name ) ) ) == 2 ) # Purge memories using filter groups. - operation = client.memory_banks.memories.purge( - name=memory_bank.name, + operation = client.agent_engines.memories.purge( + name=agent_engine.api_resource.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.memory_banks.memories.list( - name=memory_bank.name + client.agent_engines.memories.list( + name=agent_engine.api_resource.name ) ) ) == 1 ) finally: - client.memory_banks.delete(name=memory_bank.name, force=True) + client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories.purge", + test_method="agent_engines.memories.purge", ) @@ -127,29 +127,29 @@ def test_purge_memories(client): @pytest.mark.asyncio async def test_purge_memories_async(client): - memory_bank = client.memory_banks.create() + agent_engine = client.agent_engines.create() try: - client.memory_banks.memories.create( - name=memory_bank.name, + client.agent_engines.memories.create( + name=agent_engine.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.memory_banks.memories.create( - name=memory_bank.name, + client.agent_engines.memories.create( + name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.memory_banks.memories.create( - name=memory_bank.name, + client.agent_engines.memories.create( + name=agent_engine.api_resource.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - operation = await client.aio.memory_banks.memories.purge( - name=memory_bank.name, + operation = await client.aio.agent_engines.memories.purge( + name=agent_engine.api_resource.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.memory_banks.memories.list( - name=memory_bank.name + client.agent_engines.memories.list( + name=agent_engine.api_resource.name ) ) ) == 3 ) # Now, actually purge the memories. - operation = await client.aio.memory_banks.memories.purge( - name=memory_bank.name, + operation = await client.aio.agent_engines.memories.purge( + name=agent_engine.api_resource.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.memory_banks.memories.list( - name=memory_bank.name + client.agent_engines.memories.list( + name=agent_engine.api_resource.name ) ) ) == 1 ) finally: - client.memory_banks.delete(name=memory_bank.name, force=True) + client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) 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 new file mode 100644 index 0000000000..02a3e034d0 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py @@ -0,0 +1,196 @@ +# 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 + ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py b/tests/unit/agentplatform/genai/replays/test_structured_memories.py similarity index 76% rename from tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py rename to tests/unit/agentplatform/genai/replays/test_structured_memories.py index ac9cc30941..5d1c2a39a5 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py +++ b/tests/unit/agentplatform/genai/replays/test_structured_memories.py @@ -19,7 +19,10 @@ def test_generate_and_retrieve_profile(client): - # TODO: Switch to Memory Bank for creation once it supports configs. + # TODO: Use prod once available. + client._api_client._http_options.base_url = ( + "https://us-central1-autopush-aiplatform.sandbox.googleapis.com" + ) customization_config = {"disable_natural_language_memories": True} memory_bank_customization_config = types.MemoryBankCustomizationConfig( **customization_config @@ -41,7 +44,7 @@ def test_generate_and_retrieve_profile(client): structured_memory_config_obj = types.StructuredMemoryConfig( **structured_memory_config ) - memory_bank = client.agent_engines.create( + agent_engine = client.agent_engines.create( config={ "context_spec": { "memory_bank_config": { @@ -53,8 +56,8 @@ def test_generate_and_retrieve_profile(client): }, ) try: - memory_bank = client.agent_engines.get(name=memory_bank.api_resource.name) - memory_bank_config = memory_bank.api_resource.context_spec.memory_bank_config + agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name) + memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config assert memory_bank_config.customization_configs == [ memory_bank_customization_config ] @@ -63,16 +66,16 @@ def test_generate_and_retrieve_profile(client): ] scope = {"user_id": "123"} - client.memory_banks.memories.generate( - name=memory_bank.api_resource.name, + client.agent_engines.memories.generate( + name=agent_engine.api_resource.name, scope=scope, direct_contents_source={ "events": [{"content": {"parts": [{"text": "My name is Kim."}]}}] }, ) memories = list( - client.memory_banks.memories.retrieve( - name=memory_bank.api_resource.name, + client.agent_engines.memories.retrieve( + name=agent_engine.api_resource.name, scope=scope, config={"memory_types": ["STRUCTURED_PROFILE"]}, ) @@ -80,18 +83,18 @@ def test_generate_and_retrieve_profile(client): assert len(memories) >= 1 assert memories[0].memory.structured_content is not None - response = client.memory_banks.memories.retrieve_profiles( - name=memory_bank.api_resource.name, scope=scope + response = client.agent_engines.memories.retrieve_profiles( + name=agent_engine.api_resource.name, scope=scope ) assert len(response.profiles) == 1 finally: # Clean up resources. - client.memory_banks.delete(name=memory_bank.api_resource.name, force=True) + client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="memory_banks.memories.retrieve_profiles", + test_method="agent_engines.retrieve_profiles", )