From 727b8e0c9a426281f8fe302b48e9836a87645fa6 Mon Sep 17 00:00:00 2001
From: A Vertex SDK engineer <vertex-sdk-bot@google.com>
Date: Thu, 9 Apr 2026 16:40:26 -0700
Subject: [PATCH] feat: Add subset_topics to Memory Bank GenerateMemories

PiperOrigin-RevId: 897366335
---
 .../test_generate_agent_engine_memories.py    | 43 +++++++++++++++++++
 vertexai/_genai/memories.py                   |  7 +++
 vertexai/_genai/types/common.py               |  7 +++
 3 files changed, 57 insertions(+)

diff --git a/tests/unit/vertexai/genai/replays/test_generate_agent_engine_memories.py b/tests/unit/vertexai/genai/replays/test_generate_agent_engine_memories.py
index f06ae99b2a..5a432b9196 100644
--- a/tests/unit/vertexai/genai/replays/test_generate_agent_engine_memories.py
+++ b/tests/unit/vertexai/genai/replays/test_generate_agent_engine_memories.py
@@ -336,3 +336,46 @@ async def test_generate_and_rollback_memories_async(client):
     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/vertexai/_genai/memories.py b/vertexai/_genai/memories.py
index 23aa9c0016..4a48350272 100644
--- a/vertexai/_genai/memories.py
+++ b/vertexai/_genai/memories.py
@@ -174,6 +174,13 @@ def _GenerateAgentEngineMemoriesConfig_to_vertex(
             getv(from_object, ["metadata_merge_strategy"]),
         )
 
+    if getv(from_object, ["allowed_topics"]) is not None:
+        setv(
+            parent_object,
+            ["allowedTopics"],
+            [item for item in getv(from_object, ["allowed_topics"])],
+        )
+
     return to_object
 
 
diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py
index b67be65045..cc115e03cd 100644
--- a/vertexai/_genai/types/common.py
+++ b/vertexai/_genai/types/common.py
@@ -9428,6 +9428,10 @@ class GenerateAgentEngineMemoriesConfig(_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 GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False):
@@ -9465,6 +9469,9 @@ class GenerateAgentEngineMemoriesConfigDict(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."""
+
 
 GenerateAgentEngineMemoriesConfigOrDict = Union[
     GenerateAgentEngineMemoriesConfig, GenerateAgentEngineMemoriesConfigDict