feat: Initial implementation of the SDK for Memory Revisions

FUTURE_COPYBARA_INTEGRATE_REVIEW=#5883 from googleapis:release-please--branches--main 048acfd
PiperOrigin-RevId: 815928067

Author: vertex-sdk-bot

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

@@ -19,14 +19,19 @@
 from google.genai import types as genai_types
 
 
-def test_generate_memories(client):
+def test_generate_and_rollback_memories(client):
+    client._api_client._http_options.base_url = (
+        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
+    )
     agent_engine = client.agent_engines.create()
     assert not list(
-        client.agent_engines.list_memories(
+        client.agent_engines.memories.list(
             name=agent_engine.api_resource.name,
         )
     )
-    client.agent_engines.generate_memories(
+    # Generate memories using source content. This result is non-deterministic,
+    # because an LLM is used to generate the memories.
+    client.agent_engines.memories.generate(
         name=agent_engine.api_resource.name,
         scope={"user_id": "test-user-id"},
         direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource(
@@ -43,50 +48,99 @@ def test_generate_memories(client):
                 )
             ]
         ),
+        config=types.GenerateAgentEngineMemoriesConfig(
+            revision_labels={"key": "value"}
+        ),
     )
-    assert (
-        len(
-            list(
-                client.agent_engines.list_memories(
-                    name=agent_engine.api_resource.name,
-                )
-            )
+    memories = 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)
+    assert len(memories) >= 1
 
+    # Every action that modifies a memory creates a new revision.
+    memory_revisions = list(
+        client.agent_engines.memories.revisions.list(
+            name=memories[0].name,
+        )
+    )
+    assert len(memory_revisions) >= 1
+    # The revision's labels depend on the generation request's revision labels.
+    assert memory_revisions[0].labels == {"key": "value"}
+    revision_name = memory_revisions[0].name
 
-def test_generate_memories_direct_memories_source(client):
-    agent_engine = client.agent_engines.create()
-    client.agent_engines.generate_memories(
+    # Update the memory.
+    client.agent_engines.memories._update(
+        name=memories[0].name,
+        fact="This is temporary",
+        scope={"user_id": "test-user-id"},
+    )
+    memory = client.agent_engines.memories.get(name=memories[0].name)
+    assert memory.fact == "This is temporary"
+
+    # Rollback to the revision with the original fact that was created by the
+    # generation request.
+    client.agent_engines.memories.rollback(
+        name=memories[0].name,
+        target_revision_id=revision_name.split("/")[-1],
+    )
+    memory = client.agent_engines.memories.get(name=memories[0].name)
+    assert memory.fact == memory_revisions[0].fact
+
+    # Update the memory again using generation. We use the original source
+    # content to ensure that the original memory is updated. The response should
+    # refer to the previous revision.
+    response = client.agent_engines.memories.generate(
         name=agent_engine.api_resource.name,
         scope={"user_id": "test-user-id"},
-        direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource(
-            direct_memories=[
-                types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(
-                    fact="I am a software engineer."
-                ),
-                types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(
-                    fact="I like to write replay tests."
-                ),
+        direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource(
+            events=[
+                types.GenerateMemoriesRequestDirectContentsSourceEvent(
+                    content=genai_types.Content(
+                        role="model",
+                        parts=[genai_types.Part(text=memory_revisions[0].fact)],
+                    )
+                )
             ]
         ),
-        config=types.GenerateAgentEngineMemoriesConfig(wait_for_completion=True),
-    )
-    assert (
-        len(
-            list(
-                client.agent_engines.list_memories(
-                    name=agent_engine.api_resource.name,
-                )
-            )
-        )
-        >= 1
     )
+    # The memory was updated, so the previous revision is set.
+    assert response.response.generated_memories[0].previous_revision is not None
+
     client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
 
 
+# def test_generate_memories_direct_memories_source(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_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource(
+#             direct_memories=[
+#                 types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(
+#                     fact="I am a software engineer."
+#                 ),
+#                 types.GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(
+#                     fact="I like to write replay tests."
+#                 ),
+#             ]
+#         ),
+#         config=types.GenerateAgentEngineMemoriesConfig(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)
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),