feat: GenAI SDK client(memory): Add PurgeMemories

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

@@ -23,10 +23,6 @@
 
 
 def test_generate_and_rollback_memories(client):
-    # TODO(): Use prod endpoint once experiment is fully rolled out.
-    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.memories.list(
@@ -161,10 +157,6 @@ def test_generate_memories_direct_memories_source(client):
 
 @pytest.mark.asyncio
 async def test_generate_and_rollback_memories_async(client):
-    # TODO(): Use prod endpoint once revisions experiment is fully rolled out.
-    client._api_client._http_options.base_url = (
-        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
-    )
     agent_engine = client.agent_engines.create()
     await client.aio.agent_engines.memories.generate(
         name=agent_engine.api_resource.name,

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

@@ -0,0 +1,157 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring
+
+import pytest
+
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+
+
+def test_purge_memories(client):
+    """Tests purging memories."""
+    agent_engine = client.agent_engines.create()
+    try:
+        client.agent_engines.memories.create(
+            name=agent_engine.api_resource.name,
+            fact="memory_fact_1",
+            scope={"user_id": "123"},
+            config={"wait_for_completion": True},
+        )
+        client.agent_engines.memories.create(
+            name=agent_engine.api_resource.name,
+            fact="memory_fact_2",
+            scope={"user_id": "123"},
+            config={"wait_for_completion": True},
+        )
+        client.agent_engines.memories.create(
+            name=agent_engine.api_resource.name,
+            fact="memory_fact_3",
+            scope={"user_id": "456"},
+            config={"wait_for_completion": True},
+        )
+        operation = client.agent_engines.memories.purge(
+            name=agent_engine.api_resource.name,
+            filter="scope.user_id=123",
+            config={"wait_for_completion": True},
+        )
+        assert operation.done
+        assert operation.response.purge_count == 2
+        # Memories were not actually purged, because `force` was False.
+        assert (
+            len(
+                list(
+                    client.agent_engines.memories.list(
+                        name=agent_engine.api_resource.name
+                    )
+                )
+            )
+            == 3
+        )
+        # Now, actually purge the memories.
+        operation = client.agent_engines.memories.purge(
+            name=agent_engine.api_resource.name,
+            filter="scope.user_id=123",
+            force=True,
+            config={"wait_for_completion": True},
+        )
+        assert operation.done
+        assert operation.response.purge_count == 2
+        assert (
+            len(
+                list(
+                    client.agent_engines.memories.list(
+                        name=agent_engine.api_resource.name
+                    )
+                )
+            )
+            == 1
+        )
+    finally:
+        client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.memories.purge",
+)
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_purge_memories_async(client):
+    agent_engine = client.agent_engines.create()
+    try:
+        client.agent_engines.memories.create(
+            name=agent_engine.api_resource.name,
+            fact="memory_fact_1",
+            scope={"user_id": "123"},
+            config={"wait_for_completion": True},
+        )
+        client.agent_engines.memories.create(
+            name=agent_engine.api_resource.name,
+            fact="memory_fact_2",
+            scope={"user_id": "123"},
+            config={"wait_for_completion": True},
+        )
+        client.agent_engines.memories.create(
+            name=agent_engine.api_resource.name,
+            fact="memory_fact_3",
+            scope={"user_id": "456"},
+            config={"wait_for_completion": True},
+        )
+
+        operation = await client.aio.agent_engines.memories.purge(
+            name=agent_engine.api_resource.name,
+            filter="scope.user_id=123",
+            config={"wait_for_completion": True},
+        )
+        assert operation.done
+        assert operation.response.purge_count == 2
+        # Memories were not actually purged, because `force` was False.
+        assert (
+            len(
+                list(
+                    client.agent_engines.memories.list(
+                        name=agent_engine.api_resource.name
+                    )
+                )
+            )
+            == 3
+        )
+        # Now, actually purge the memories.
+        operation = await client.aio.agent_engines.memories.purge(
+            name=agent_engine.api_resource.name,
+            filter="scope.user_id=123",
+            force=True,
+            config={"wait_for_completion": True},
+        )
+        assert operation.done
+        assert operation.response.purge_count == 2
+        assert (
+            len(
+                list(
+                    client.agent_engines.memories.list(
+                        name=agent_engine.api_resource.name
+                    )
+                )
+            )
+            == 1
+        )
+    finally:
+        client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)