feat: Add metadata and metadata filtering to memories

PiperOrigin-RevId: 854287666

Author: vertex-sdk-bot

vertexai/_genai/memories.py

@@ -77,6 +77,9 @@ def _AgentEngineMemoryConfig_to_vertex(
             parent_object, ["topics"], [item for item in getv(from_object, ["topics"])]
         )
 
+    if getv(from_object, ["metadata"]) is not None:
+        setv(parent_object, ["metadata"], getv(from_object, ["metadata"]))
+
     return to_object
 
 
@@ -153,6 +156,16 @@ def _GenerateAgentEngineMemoriesConfig_to_vertex(
             getv(from_object, ["disable_memory_revisions"]),
         )
 
+    if getv(from_object, ["metadata"]) is not None:
+        setv(parent_object, ["metadata"], getv(from_object, ["metadata"]))
+
+    if getv(from_object, ["metadata_merge_strategy"]) is not None:
+        setv(
+            parent_object,
+            ["metadataMergeStrategy"],
+            getv(from_object, ["metadata_merge_strategy"]),
+        )
+
     return to_object
 
 
@@ -316,6 +329,13 @@ def _RetrieveAgentEngineMemoriesConfig_to_vertex(
     if getv(from_object, ["filter"]) is not None:
         setv(parent_object, ["filter"], getv(from_object, ["filter"]))
 
+    if getv(from_object, ["filter_groups"]) is not None:
+        setv(
+            parent_object,
+            ["filterGroups"],
+            [item for item in getv(from_object, ["filter_groups"])],
+        )
+
     return to_object
 
 
@@ -413,6 +433,9 @@ def _UpdateAgentEngineMemoryConfig_to_vertex(
             parent_object, ["topics"], [item for item in getv(from_object, ["topics"])]
         )
 
+    if getv(from_object, ["metadata"]) is not None:
+        setv(parent_object, ["metadata"], getv(from_object, ["metadata"]))
+
     if getv(from_object, ["update_mask"]) is not None:
         setv(
             parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"])

vertexai/_genai/types/__init__.py

@@ -578,7 +578,17 @@
 from .common import MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict
 from .common import MemoryBankCustomizationConfigMemoryTopicOrDict
 from .common import MemoryBankCustomizationConfigOrDict
+from .common import MemoryConjunctionFilter
+from .common import MemoryConjunctionFilterDict
+from .common import MemoryConjunctionFilterOrDict
 from .common import MemoryDict
+from .common import MemoryFilter
+from .common import MemoryFilterDict
+from .common import MemoryFilterOrDict
+from .common import MemoryMetadataMergeStrategy
+from .common import MemoryMetadataValue
+from .common import MemoryMetadataValueDict
+from .common import MemoryMetadataValueOrDict
 from .common import MemoryOrDict
 from .common import MemoryRevision
 from .common import MemoryRevisionDict
@@ -613,6 +623,7 @@
 from .common import ObservabilityEvalCase
 from .common import ObservabilityEvalCaseDict
 from .common import ObservabilityEvalCaseOrDict
+from .common import Operator
 from .common import OptimizeConfig
 from .common import OptimizeConfigDict
 from .common import OptimizeConfigOrDict
@@ -1523,6 +1534,15 @@
     "RetrieveMemoriesRequestSimpleRetrievalParams",
     "RetrieveMemoriesRequestSimpleRetrievalParamsDict",
     "RetrieveMemoriesRequestSimpleRetrievalParamsOrDict",
+    "MemoryMetadataValue",
+    "MemoryMetadataValueDict",
+    "MemoryMetadataValueOrDict",
+    "MemoryFilter",
+    "MemoryFilterDict",
+    "MemoryFilterOrDict",
+    "MemoryConjunctionFilter",
+    "MemoryConjunctionFilterDict",
+    "MemoryConjunctionFilterOrDict",
     "RetrieveAgentEngineMemoriesConfig",
     "RetrieveAgentEngineMemoriesConfigDict",
     "RetrieveAgentEngineMemoriesConfigOrDict",
@@ -1909,6 +1929,7 @@
     "IdentityType",
     "AgentServerMode",
     "ManagedTopicEnum",
+    "Operator",
     "Language",
     "MachineConfig",
     "State",
@@ -1917,6 +1938,7 @@
     "RubricContentType",
     "EvaluationRunState",
     "OptimizeTarget",
+    "MemoryMetadataMergeStrategy",
     "GenerateMemoriesResponseGeneratedMemoryAction",
     "PromptOptimizerMethod",
     "PromptData",

vertexai/_genai/types/common.py

@@ -243,6 +243,19 @@ class ManagedTopicEnum(_common.CaseInSensitiveEnum):
     """Information that the user explicitly requested to remember or forget."""
 
 
+class Operator(_common.CaseInSensitiveEnum):
+    """Operator to apply to the filter. If not set, then EQUAL will be used."""
+
+    OPERATOR_UNSPECIFIED = "OPERATOR_UNSPECIFIED"
+    """Unspecified operator. Defaults to EQUAL."""
+    EQUAL = "EQUAL"
+    """Equal to."""
+    GREATER_THAN = "GREATER_THAN"
+    """Greater than."""
+    LESS_THAN = "LESS_THAN"
+    """Less than."""
+
+
 class Language(_common.CaseInSensitiveEnum):
     """The coding language supported in this environment."""
 
@@ -345,6 +358,19 @@ class OptimizeTarget(_common.CaseInSensitiveEnum):
     """The prompt optimizer based on user provided examples with target responses."""
 
 
+class MemoryMetadataMergeStrategy(_common.CaseInSensitiveEnum):
+    """The strategy to use when applying metadata to existing memories during consolidation."""
+
+    METADATA_MERGE_STRATEGY_UNSPECIFIED = "METADATA_MERGE_STRATEGY_UNSPECIFIED"
+    """The metadata merge strategy is unspecified."""
+    OVERWRITE = "OVERWRITE"
+    """Replace the metadata of the updated memories with the new metadata."""
+    MERGE = "MERGE"
+    """Append new metadata to the existing metadata. If there are duplicate keys, the existing values will be overwritten."""
+    REQUIRE_EXACT_MATCH = "REQUIRE_EXACT_MATCH"
+    """Restrict consolidation to memories that have exactly the same metadata as the request. If a memory doesn't have the same metadata, it is not eligible for consolidation."""
+
+
 class GenerateMemoriesResponseGeneratedMemoryAction(_common.CaseInSensitiveEnum):
     """The action to take."""
 
@@ -6622,6 +6648,10 @@ class AgentEngineMemoryConfig(_common.BaseModel):
     topics: Optional[list[MemoryTopicId]] = Field(
         default=None, description="""Optional. The topics of the memory."""
     )
+    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.""",
+    )
 
 
 class AgentEngineMemoryConfigDict(TypedDict, total=False):
@@ -6659,6 +6689,9 @@ class AgentEngineMemoryConfigDict(TypedDict, total=False):
     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."""
+
 
 AgentEngineMemoryConfigOrDict = Union[
     AgentEngineMemoryConfig, AgentEngineMemoryConfigDict
@@ -6769,6 +6802,10 @@ class Memory(_common.BaseModel):
     topics: Optional[list[MemoryTopicId]] = Field(
         default=None, description="""Optional. The Topics of the Memory."""
     )
+    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.""",
+    )
 
 
 class MemoryDict(TypedDict, total=False):
@@ -6813,6 +6850,9 @@ class MemoryDict(TypedDict, total=False):
     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."""
+
 
 MemoryOrDict = Union[Memory, MemoryDict]
 
@@ -7124,6 +7164,14 @@ class GenerateAgentEngineMemoriesConfig(_common.BaseModel):
         default=None,
         description="""Optional. Input only. If true, no revisions will be created for this request.""",
     )
+    metadata: Optional[dict[str, "MemoryMetadataValue"]] = Field(
+        default=None,
+        description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""",
+    )
+    metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field(
+        default=None,
+        description="""Optional. The strategy to use when applying metadata to existing memories.""",
+    )
 
 
 class GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False):
@@ -7155,6 +7203,12 @@ class GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False):
     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."""
+
 
 GenerateAgentEngineMemoriesConfigOrDict = Union[
     GenerateAgentEngineMemoriesConfig, GenerateAgentEngineMemoriesConfigDict
@@ -7612,6 +7666,95 @@ class RetrieveMemoriesRequestSimpleRetrievalParamsDict(TypedDict, total=False):
 ]
 
 
+class MemoryMetadataValue(_common.BaseModel):
+    """Memory metadata."""
+
+    timestamp_value: Optional[datetime.datetime] = Field(
+        default=None,
+        description="""Timestamp value. When filtering on timestamp values, only the seconds field will be compared.""",
+    )
+    double_value: Optional[float] = Field(default=None, description="""Double value.""")
+    bool_value: Optional[bool] = Field(default=None, description="""Boolean value.""")
+    string_value: Optional[str] = Field(default=None, description="""String value.""")
+
+
+class MemoryMetadataValueDict(TypedDict, total=False):
+    """Memory metadata."""
+
+    timestamp_value: Optional[datetime.datetime]
+    """Timestamp value. When filtering on timestamp values, only the seconds field will be compared."""
+
+    double_value: Optional[float]
+    """Double value."""
+
+    bool_value: Optional[bool]
+    """Boolean value."""
+
+    string_value: Optional[str]
+    """String value."""
+
+
+MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict]
+
+
+class MemoryFilter(_common.BaseModel):
+    """Filter to apply when retrieving memories."""
+
+    op: Optional[Operator] = Field(
+        default=None,
+        description="""Operator to apply to the filter. If not set, then EQUAL will be used.""",
+    )
+    negate: Optional[bool] = Field(
+        default=None, description="""If true, the filter will be negated."""
+    )
+    key: Optional[str] = Field(
+        default=None,
+        description="""Key of the filter. For example, "author" would apply to `metadata` entries with the key "author".""",
+    )
+    value: Optional[MemoryMetadataValue] = Field(
+        default=None, description="""Value to compare to."""
+    )
+
+
+class MemoryFilterDict(TypedDict, total=False):
+    """Filter to apply when retrieving memories."""
+
+    op: Optional[Operator]
+    """Operator to apply to the filter. If not set, then EQUAL will be used."""
+
+    negate: Optional[bool]
+    """If true, the filter will be negated."""
+
+    key: Optional[str]
+    """Key of the filter. For example, "author" would apply to `metadata` entries with the key "author"."""
+
+    value: Optional[MemoryMetadataValueDict]
+    """Value to compare to."""
+
+
+MemoryFilterOrDict = Union[MemoryFilter, MemoryFilterDict]
+
+
+class MemoryConjunctionFilter(_common.BaseModel):
+    """The conjunction filter for memories."""
+
+    filters: Optional[list[MemoryFilter]] = Field(
+        default=None, description="""Filters that will combined using AND logic."""
+    )
+
+
+class MemoryConjunctionFilterDict(TypedDict, total=False):
+    """The conjunction filter for memories."""
+
+    filters: Optional[list[MemoryFilterDict]]
+    """Filters that will combined using AND logic."""
+
+
+MemoryConjunctionFilterOrDict = Union[
+    MemoryConjunctionFilter, MemoryConjunctionFilterDict
+]
+
+
 class RetrieveAgentEngineMemoriesConfig(_common.BaseModel):
     """Config for retrieving memories."""
 
@@ -7629,6 +7772,23 @@ class RetrieveAgentEngineMemoriesConfig(_common.BaseModel):
        * `update_time`
       """,
     )
+    filter_groups: Optional[list[MemoryConjunctionFilter]] = Field(
+        default=None,
+        description="""Metadata filters that will be applied to the retrieved memories'
+      `metadata` using OR logic. Filters are defined using disjunctive normal
+      form (OR of ANDs).
+
+      For example:
+      `filter_groups: [{filters: [{key: "author", value: {string_value: "agent
+      `123"}, op: EQUAL}]}, {filters: [{key: "label", value: {string_value:
+      "travel"}, op: EQUAL}, {key: "author", value: {string_value: "agent 321"},
+      op: EQUAL}]}]`
+
+      would be equivalent to the logical expression:
+      `(metadata.author = "agent 123" OR (metadata.label = "travel" AND
+      metadata.author = "agent 321"))`.
+      """,
+    )
 
 
 class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False):
@@ -7647,6 +7807,22 @@ class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False):
        * `update_time`
       """
 
+    filter_groups: Optional[list[MemoryConjunctionFilterDict]]
+    """Metadata filters that will be applied to the retrieved memories'
+      `metadata` using OR logic. Filters are defined using disjunctive normal
+      form (OR of ANDs).
+
+      For example:
+      `filter_groups: [{filters: [{key: "author", value: {string_value: "agent
+      `123"}, op: EQUAL}]}, {filters: [{key: "label", value: {string_value:
+      "travel"}, op: EQUAL}, {key: "author", value: {string_value: "agent 321"},
+      op: EQUAL}]}]`
+
+      would be equivalent to the logical expression:
+      `(metadata.author = "agent 123" OR (metadata.label = "travel" AND
+      metadata.author = "agent 321"))`.
+      """
+
 
 RetrieveAgentEngineMemoriesConfigOrDict = Union[
     RetrieveAgentEngineMemoriesConfig, RetrieveAgentEngineMemoriesConfigDict
@@ -7913,6 +8089,10 @@ class UpdateAgentEngineMemoryConfig(_common.BaseModel):
     topics: Optional[list[MemoryTopicId]] = Field(
         default=None, description="""Optional. The topics of the memory."""
     )
+    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.""",
+    )
     update_mask: Optional[str] = Field(
         default=None,
         description="""The update mask to apply. For the `FieldMask` definition, see
@@ -7955,6 +8135,9 @@ class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False):
     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."""
+
     update_mask: Optional[str]
     """The update mask to apply. For the `FieldMask` definition, see
       https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""