From a8948c4dfcc65b73ea579b93f8c36baa65817f25 Mon Sep 17 00:00:00 2001
From: A Vertex SDK engineer <vertex-sdk-bot@google.com>
Date: Thu, 2 Apr 2026 14:54:43 -0700
Subject: [PATCH] feat: Add consolidation customization to Memory Bank

PiperOrigin-RevId: 893716901
---
 .../genai/replays/test_create_agent_engine.py |  1 +
 vertexai/_genai/types/__init__.py             |  6 ++++
 vertexai/_genai/types/common.py               | 31 +++++++++++++++++++
 3 files changed, 38 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 ab9a029e90..16dd0c6c68 100644
--- a/tests/unit/vertexai/genai/replays/test_create_agent_engine.py
+++ b/tests/unit/vertexai/genai/replays/test_create_agent_engine.py
@@ -85,6 +85,7 @@ def test_create_with_context_spec(client):
             }
         ],
         "enable_third_person_memories": True,
+        "consolidation_config": {"revisions_per_candidate_count": 5},
     }
     memory_bank_customization_config = types.MemoryBankCustomizationConfig(
         **customization_config
diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py
index 937e8fe43f..241344b847 100644
--- a/vertexai/_genai/types/__init__.py
+++ b/vertexai/_genai/types/__init__.py
@@ -607,6 +607,9 @@
 from .common import MapInstanceOrDict
 from .common import Memory
 from .common import MemoryBankCustomizationConfig
+from .common import MemoryBankCustomizationConfigConsolidationConfig
+from .common import MemoryBankCustomizationConfigConsolidationConfigDict
+from .common import MemoryBankCustomizationConfigConsolidationConfigOrDict
 from .common import MemoryBankCustomizationConfigDict
 from .common import MemoryBankCustomizationConfigGenerateMemoriesExample
 from .common import (
@@ -1577,6 +1580,9 @@
     "MemoryBankCustomizationConfigMemoryTopic",
     "MemoryBankCustomizationConfigMemoryTopicDict",
     "MemoryBankCustomizationConfigMemoryTopicOrDict",
+    "MemoryBankCustomizationConfigConsolidationConfig",
+    "MemoryBankCustomizationConfigConsolidationConfigDict",
+    "MemoryBankCustomizationConfigConsolidationConfigOrDict",
     "MemoryBankCustomizationConfig",
     "MemoryBankCustomizationConfigDict",
     "MemoryBankCustomizationConfigOrDict",
diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py
index d549c5490c..a438ee1690 100644
--- a/vertexai/_genai/types/common.py
+++ b/vertexai/_genai/types/common.py
@@ -6214,6 +6214,28 @@ class MemoryBankCustomizationConfigMemoryTopicDict(TypedDict, total=False):
 ]
 
 
+class MemoryBankCustomizationConfigConsolidationConfig(_common.BaseModel):
+    """Represents configuration for customizing how memories are consolidated."""
+
+    revisions_per_candidate_count: Optional[int] = Field(
+        default=None,
+        description="""Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""",
+    )
+
+
+class MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=False):
+    """Represents configuration for customizing how memories are consolidated."""
+
+    revisions_per_candidate_count: Optional[int]
+    """Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered."""
+
+
+MemoryBankCustomizationConfigConsolidationConfigOrDict = Union[
+    MemoryBankCustomizationConfigConsolidationConfig,
+    MemoryBankCustomizationConfigConsolidationConfigDict,
+]
+
+
 class MemoryBankCustomizationConfig(_common.BaseModel):
     """Configuration for organizing memories for a particular scope."""
 
@@ -6235,6 +6257,12 @@ class MemoryBankCustomizationConfig(_common.BaseModel):
         default=None,
         description="""Optional. The scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""",
     )
+    consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfig] = (
+        Field(
+            default=None,
+            description="""Optional. Represents configuration for customizing how memories are consolidated together.""",
+        )
+    )
 
 
 class MemoryBankCustomizationConfigDict(TypedDict, total=False):
@@ -6254,6 +6282,9 @@ class MemoryBankCustomizationConfigDict(TypedDict, total=False):
     scope_keys: Optional[list[str]]
     """Optional. The scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank."""
 
+    consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfigDict]
+    """Optional. Represents configuration for customizing how memories are consolidated together."""
+
 
 MemoryBankCustomizationConfigOrDict = Union[
     MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict