From de941a6fc3ca7885a0fc3e42b5ef7c44a783963b Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Wed, 22 Oct 2025 13:28:22 -0700 Subject: [PATCH] feat: GenAI SDK client(memory): Add Memory Topic labels to Memory PiperOrigin-RevId: 822718268 --- .../genai/replays/test_create_agent_engine.py | 28 ++++++++++ vertexai/_genai/memories.py | 10 ++++ vertexai/_genai/types.py | 56 +++++++++++++++++++ 3 files changed, 94 insertions(+) diff --git a/tests/unit/vertexai/genai/replays/test_create_agent_engine.py b/tests/unit/vertexai/genai/replays/test_create_agent_engine.py index 16f18f4abd..0afd3c6d32 100644 --- a/tests/unit/vertexai/genai/replays/test_create_agent_engine.py +++ b/tests/unit/vertexai/genai/replays/test_create_agent_engine.py @@ -17,6 +17,7 @@ import os from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types def test_create_config_lightweight(client): @@ -54,6 +55,29 @@ def test_create_with_context_spec(client): parent = f"projects/{project}/locations/{location}" generation_model = f"{parent}/publishers/google/models/gemini-2.0-flash-001" embedding_model = f"{parent}/publishers/google/models/text-embedding-005" + customization_config = { + "memory_topics": [ + {"managed_memory_topic": {"managed_topic_enum": "USER_PREFERENCES"}} + ], + "generate_memories_examples": [ + { + "conversation_source": { + "events": [ + {"content": {"role": "user", "parts": [{"text": "Hello"}]}} + ] + }, + "generatedMemories": [ + { + "fact": "I like to say hello.", + "topics": [{"managed_memory_topic": "USER_PREFERENCES"}], + } + ], + } + ], + } + memory_bank_customization_config = types.MemoryBankCustomizationConfig( + **customization_config + ) agent_engine = client.agent_engines.create( config={ @@ -64,6 +88,7 @@ def test_create_with_context_spec(client): "embedding_model": embedding_model, }, "ttl_config": {"default_ttl": "120s"}, + "customization_configs": [memory_bank_customization_config], }, }, "http_options": {"api_version": "v1beta1"}, @@ -76,6 +101,9 @@ def test_create_with_context_spec(client): memory_bank_config.similarity_search_config.embedding_model == embedding_model ) assert memory_bank_config.ttl_config.default_ttl == "120s" + assert memory_bank_config.customization_configs == [ + memory_bank_customization_config + ] # Clean up resources. client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/vertexai/_genai/memories.py b/vertexai/_genai/memories.py index 57f76d3c3a..4ac889c513 100644 --- a/vertexai/_genai/memories.py +++ b/vertexai/_genai/memories.py @@ -72,6 +72,11 @@ def _AgentEngineMemoryConfig_to_vertex( 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"])] + ) + return to_object @@ -365,6 +370,11 @@ def _UpdateAgentEngineMemoryConfig_to_vertex( 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, ["update_mask"]) is not None: setv( parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) diff --git a/vertexai/_genai/types.py b/vertexai/_genai/types.py index 17d0221b6d..876ea2de19 100644 --- a/vertexai/_genai/types.py +++ b/vertexai/_genai/types.py @@ -5330,6 +5330,30 @@ class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict ] +class MemoryTopicId(_common.BaseModel): + """The topic ID for a memory.""" + + custom_memory_topic_label: Optional[str] = Field( + default=None, description="""Optional. The custom memory topic label.""" + ) + managed_memory_topic: Optional[ManagedTopicEnum] = Field( + default=None, description="""Optional. The managed memory topic.""" + ) + + +class MemoryTopicIdDict(TypedDict, total=False): + """The topic ID for a memory.""" + + custom_memory_topic_label: Optional[str] + """Optional. The custom memory topic label.""" + + managed_memory_topic: Optional[ManagedTopicEnum] + """Optional. The managed memory topic.""" + + +MemoryTopicIdOrDict = Union[MemoryTopicId, MemoryTopicIdDict] + + class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( _common.BaseModel ): @@ -5338,6 +5362,10 @@ class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( fact: Optional[str] = Field( default=None, description="""Required. The fact to generate a memory from.""" ) + topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. The list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""", + ) class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( @@ -5348,6 +5376,9 @@ class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( fact: Optional[str] """Required. The fact to generate a memory from.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""" + MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryOrDict = Union[ MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory, @@ -6432,6 +6463,9 @@ class AgentEngineMemoryConfig(_common.BaseModel): 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.""" + ) class AgentEngineMemoryConfigDict(TypedDict, total=False): @@ -6466,6 +6500,9 @@ class AgentEngineMemoryConfigDict(TypedDict, total=False): 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.""" + AgentEngineMemoryConfigOrDict = Union[ AgentEngineMemoryConfig, AgentEngineMemoryConfigDict @@ -6573,6 +6610,9 @@ class Memory(_common.BaseModel): default=None, description="""Output only. Timestamp when this Memory was most recently updated.""", ) + topics: Optional[list[MemoryTopicId]] = Field( + default=None, description="""Optional. The Topics of the Memory.""" + ) class MemoryDict(TypedDict, total=False): @@ -6614,6 +6654,9 @@ class MemoryDict(TypedDict, total=False): update_time: Optional[datetime.datetime] """Output only. Timestamp when this Memory was most recently updated.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The Topics of the Memory.""" + MemoryOrDict = Union[Memory, MemoryDict] @@ -6840,6 +6883,10 @@ class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel) default=None, description="""Required. The fact to consolidate with existing memories.""", ) + topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. The topics that the consolidated memories should be associated with.""", + ) class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( @@ -6850,6 +6897,9 @@ class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( fact: Optional[str] """Required. The fact to consolidate with existing memories.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics that the consolidated memories should be associated with.""" + GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, @@ -7683,6 +7733,9 @@ class UpdateAgentEngineMemoryConfig(_common.BaseModel): 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.""" + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -7722,6 +7775,9 @@ class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): 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.""" + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""