feat: Add subset_topics to Memory Bank GenerateMemories

PiperOrigin-RevId: 894126692

Author: vertex-sdk-bot

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)

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
 
 

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