Skip to content

Commit 6415a97

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: GenAI SDK client(memory): Add PurgeMemories
PiperOrigin-RevId: 831917031
1 parent 46285bf commit 6415a97

File tree

5 files changed

+543
-8
lines changed

5 files changed

+543
-8
lines changed

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,6 @@
2323

2424

2525
def test_generate_and_rollback_memories(client):
26-
# TODO(): Use prod endpoint once experiment is fully rolled out.
27-
client._api_client._http_options.base_url = (
28-
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
29-
)
3026
agent_engine = client.agent_engines.create()
3127
assert not list(
3228
client.agent_engines.memories.list(
@@ -161,10 +157,6 @@ def test_generate_memories_direct_memories_source(client):
161157

162158
@pytest.mark.asyncio
163159
async def test_generate_and_rollback_memories_async(client):
164-
# TODO(): Use prod endpoint once revisions experiment is fully rolled out.
165-
client._api_client._http_options.base_url = (
166-
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
167-
)
168160
agent_engine = client.agent_engines.create()
169161
await client.aio.agent_engines.memories.generate(
170162
name=agent_engine.api_resource.name,

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

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
16+
17+
import pytest
18+
19+
20+
from tests.unit.vertexai.genai.replays import pytest_helper
21+
22+
23+
def test_purge_memories(client):
24+
"""Tests purging memories."""
25+
agent_engine = client.agent_engines.create()
26+
try:
27+
client.agent_engines.memories.create(
28+
name=agent_engine.api_resource.name,
29+
fact="memory_fact_1",
30+
scope={"user_id": "123"},
31+
config={"wait_for_completion": True},
32+
)
33+
client.agent_engines.memories.create(
34+
name=agent_engine.api_resource.name,
35+
fact="memory_fact_2",
36+
scope={"user_id": "123"},
37+
config={"wait_for_completion": True},
38+
)
39+
client.agent_engines.memories.create(
40+
name=agent_engine.api_resource.name,
41+
fact="memory_fact_3",
42+
scope={"user_id": "456"},
43+
config={"wait_for_completion": True},
44+
)
45+
operation = client.agent_engines.memories.purge(
46+
name=agent_engine.api_resource.name,
47+
filter="scope.user_id=123",
48+
config={"wait_for_completion": True},
49+
)
50+
assert operation.done
51+
assert operation.response.purge_count == 2
52+
# Memories were not actually purged, because `force` was False.
53+
assert (
54+
len(
55+
list(
56+
client.agent_engines.memories.list(
57+
name=agent_engine.api_resource.name
58+
)
59+
)
60+
)
61+
== 3
62+
)
63+
# Now, actually purge the memories.
64+
operation = client.agent_engines.memories.purge(
65+
name=agent_engine.api_resource.name,
66+
filter="scope.user_id=123",
67+
force=True,
68+
config={"wait_for_completion": True},
69+
)
70+
assert operation.done
71+
assert operation.response.purge_count == 2
72+
assert (
73+
len(
74+
list(
75+
client.agent_engines.memories.list(
76+
name=agent_engine.api_resource.name
77+
)
78+
)
79+
)
80+
== 1
81+
)
82+
finally:
83+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
84+
85+
86+
pytestmark = pytest_helper.setup(
87+
file=__file__,
88+
globals_for_file=globals(),
89+
test_method="agent_engines.memories.purge",
90+
)
91+
92+
93+
pytest_plugins = ("pytest_asyncio",)
94+
95+
96+
@pytest.mark.asyncio
97+
async def test_purge_memories_async(client):
98+
agent_engine = client.agent_engines.create()
99+
try:
100+
client.agent_engines.memories.create(
101+
name=agent_engine.api_resource.name,
102+
fact="memory_fact_1",
103+
scope={"user_id": "123"},
104+
config={"wait_for_completion": True},
105+
)
106+
client.agent_engines.memories.create(
107+
name=agent_engine.api_resource.name,
108+
fact="memory_fact_2",
109+
scope={"user_id": "123"},
110+
config={"wait_for_completion": True},
111+
)
112+
client.agent_engines.memories.create(
113+
name=agent_engine.api_resource.name,
114+
fact="memory_fact_3",
115+
scope={"user_id": "456"},
116+
config={"wait_for_completion": True},
117+
)
118+
119+
operation = await client.aio.agent_engines.memories.purge(
120+
name=agent_engine.api_resource.name,
121+
filter="scope.user_id=123",
122+
config={"wait_for_completion": True},
123+
)
124+
assert operation.done
125+
assert operation.response.purge_count == 2
126+
# Memories were not actually purged, because `force` was False.
127+
assert (
128+
len(
129+
list(
130+
client.agent_engines.memories.list(
131+
name=agent_engine.api_resource.name
132+
)
133+
)
134+
)
135+
== 3
136+
)
137+
# Now, actually purge the memories.
138+
operation = await client.aio.agent_engines.memories.purge(
139+
name=agent_engine.api_resource.name,
140+
filter="scope.user_id=123",
141+
force=True,
142+
config={"wait_for_completion": True},
143+
)
144+
assert operation.done
145+
assert operation.response.purge_count == 2
146+
assert (
147+
len(
148+
list(
149+
client.agent_engines.memories.list(
150+
name=agent_engine.api_resource.name
151+
)
152+
)
153+
)
154+
== 1
155+
)
156+
finally:
157+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)

0 commit comments

Comments
 (0)