feat: Add memory_id to Create Memory

feat: Add raw_event to Append Event

PiperOrigin-RevId: 889309558

Author: vertex-sdk-bot

tests/unit/vertexai/genai/replays/test_append_agent_engine_session_event.py

@@ -41,6 +41,9 @@ def test_append_session_event(client):
             },
             "error_code": "test-error-code",
             "error_message": "test-error-message",
+            "raw_event": {
+                "test-key": "test-value",
+            },
         },
     )
 

tests/unit/vertexai/genai/replays/test_create_agent_engine_memory.py

@@ -87,6 +87,30 @@ def test_create_memory_with_expire_time(client):
     client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
 
 
+def test_create_memory_with_custom_memory_id(client):
+    agent_engine = client.agent_engines.create()
+    assert isinstance(agent_engine, types.AgentEngine)
+    assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+
+    operation = client.agent_engines.memories.create(
+        name=agent_engine.api_resource.name,
+        fact="memory_fact",
+        scope={"user_id": "123"},
+        config=types.AgentEngineMemoryConfig(
+            display_name="my_memory_fact", memory_id="my-memory-id"
+        ),
+    )
+    assert isinstance(operation, types.AgentEngineMemoryOperation)
+    assert operation.response.fact == "memory_fact"
+    assert operation.response.scope == {"user_id": "123"}
+    assert (
+        operation.response.name
+        == f"{agent_engine.api_resource.name}/memories/my-memory-id"
+    )
+    # Clean up resources.
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),

tests/unit/vertexai/genai/replays/test_list_agent_engine_session_events.py

@@ -38,12 +38,18 @@ def test_list_session_events(client):
         author="test-user-123",
         invocation_id="test-invocation-id",
         timestamp=datetime.datetime.fromtimestamp(1234567890, tz=datetime.timezone.utc),
+        config={
+            "raw_event": {
+                "test-key": "test-value",
+            },
+        },
     )
     session_event_list = client.agent_engines.sessions.events.list(
         name=session.name,
     )
     assert len(session_event_list) == 1
     assert isinstance(session_event_list[0], types.SessionEvent)
+    assert session_event_list[0].raw_event == {"test-key": "test-value"}
 
 
 pytestmark = pytest_helper.setup(

vertexai/_genai/memories.py

@@ -90,6 +90,9 @@ def _AgentEngineMemoryConfig_to_vertex(
             {k: v for k, v in getv(from_object, ["metadata"]).items()},
         )
 
+    if getv(from_object, ["memory_id"]) is not None:
+        setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"]))
+
     return to_object
 
 
@@ -423,6 +426,9 @@ def _UpdateAgentEngineMemoryConfig_to_vertex(
             {k: v for k, v in getv(from_object, ["metadata"]).items()},
         )
 
+    if getv(from_object, ["memory_id"]) is not None:
+        setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"]))
+
     if getv(from_object, ["update_mask"]) is not None:
         setv(
             parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"])

vertexai/_genai/session_events.py

@@ -57,6 +57,9 @@ def _AppendAgentEngineSessionEventConfig_to_vertex(
     if getv(from_object, ["event_metadata"]) is not None:
         setv(parent_object, ["eventMetadata"], getv(from_object, ["event_metadata"]))
 
+    if getv(from_object, ["raw_event"]) is not None:
+        setv(parent_object, ["rawEvent"], getv(from_object, ["raw_event"]))
+
     return to_object
 
 

vertexai/_genai/types/common.py

@@ -8122,6 +8122,10 @@ class AgentEngineMemoryConfig(_common.BaseModel):
         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.""",
     )
+    memory_id: Optional[str] = Field(
+        default=None,
+        description="""Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""",
+    )
 
 
 class AgentEngineMemoryConfigDict(TypedDict, total=False):
@@ -8162,6 +8166,9 @@ class AgentEngineMemoryConfigDict(TypedDict, total=False):
     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."""
 
+    memory_id: Optional[str]
+    """Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number."""
+
 
 AgentEngineMemoryConfigOrDict = Union[
     AgentEngineMemoryConfig, AgentEngineMemoryConfigDict
@@ -9539,6 +9546,10 @@ class UpdateAgentEngineMemoryConfig(_common.BaseModel):
         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.""",
     )
+    memory_id: Optional[str] = Field(
+        default=None,
+        description="""Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""",
+    )
     update_mask: Optional[str] = Field(
         default=None,
         description="""The update mask to apply. For the `FieldMask` definition, see
@@ -9584,6 +9595,9 @@ class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False):
     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."""
 
+    memory_id: Optional[str]
+    """Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number."""
+
     update_mask: Optional[str]
     """The update mask to apply. For the `FieldMask` definition, see
       https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
@@ -11455,6 +11469,10 @@ class AppendAgentEngineSessionEventConfig(_common.BaseModel):
     event_metadata: Optional[EventMetadata] = Field(
         default=None, description="""Metadata relating to the session event."""
     )
+    raw_event: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="""Weakly typed raw event data in proto struct format.""",
+    )
 
 
 class AppendAgentEngineSessionEventConfigDict(TypedDict, total=False):
@@ -11478,6 +11496,9 @@ class AppendAgentEngineSessionEventConfigDict(TypedDict, total=False):
     event_metadata: Optional[EventMetadataDict]
     """Metadata relating to the session event."""
 
+    raw_event: Optional[dict[str, Any]]
+    """Weakly typed raw event data in proto struct format."""
+
 
 AppendAgentEngineSessionEventConfigOrDict = Union[
     AppendAgentEngineSessionEventConfig, AppendAgentEngineSessionEventConfigDict
@@ -11647,6 +11668,10 @@ class SessionEvent(_common.BaseModel):
         default=None,
         description="""Required. Timestamp when the event was created on client side.""",
     )
+    raw_event: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="""Optional. Weakly typed raw event data in proto struct format.""",
+    )
 
 
 class SessionEventDict(TypedDict, total=False):
@@ -11679,6 +11704,9 @@ class SessionEventDict(TypedDict, total=False):
     timestamp: Optional[datetime.datetime]
     """Required. Timestamp when the event was created on client side."""
 
+    raw_event: Optional[dict[str, Any]]
+    """Optional. Weakly typed raw event data in proto struct format."""
+
 
 SessionEventOrDict = Union[SessionEvent, SessionEventDict]