Skip to content

Commit e8d18b6

Browse files
yeesiancopybara-github
yeesian
authored and
copybara-github
committed
feat: GenAI SDK client - Initial release of Agent Engine Memories SDK
PiperOrigin-RevId: 776202806
1 parent a506b94 commit e8d18b6

10 files changed

+4322
-627
lines changed

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import os
1818

1919
from tests.unit.vertexai.genai.replays import pytest_helper
20-
from vertexai._genai import types
2120

2221

2322
def test_create_config_lightweight(client):
@@ -39,13 +38,6 @@ def test_create_config_lightweight(client):
3938
}
4039

4140

42-
def test_create_operation(client):
43-
operation = client.agent_engines.create(
44-
config=types.AgentEngineConfig(return_agent=False)
45-
)
46-
assert "reasoningEngines" in operation.name
47-
48-
4941
pytestmark = pytest_helper.setup(
5042
file=__file__,
5143
globals_for_file=globals(),

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
19+
20+
def test_create_memory(client):
21+
agent_engine = client.agent_engines.create()
22+
operation = client.agent_engines.create_memory(
23+
name=agent_engine.api_resource.name,
24+
fact="memory_fact",
25+
scope={"user_id": "123"},
26+
)
27+
assert operation.response.fact == "memory_fact"
28+
assert operation.response.scope == {"user_id": "123"}
29+
assert operation.response.name.startswith(agent_engine.api_resource.name)
30+
31+
32+
pytestmark = pytest_helper.setup(
33+
file=__file__,
34+
globals_for_file=globals(),
35+
test_method="agent_engines.create_memory",
36+
)

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
19+
20+
def test_delete_memory(client):
21+
agent_engine = client.agent_engines.create()
22+
operation = client.agent_engines.create_memory(
23+
name=agent_engine.api_resource.name,
24+
fact="memory_fact",
25+
scope={"user_id": "123"},
26+
)
27+
memory = operation.response
28+
operation = client.agent_engines.delete_memory(name=memory.name)
29+
assert operation.name.startswith(memory.name + "/operations/")
30+
31+
32+
pytestmark = pytest_helper.setup(
33+
file=__file__,
34+
globals_for_file=globals(),
35+
test_method="agent_engines.delete_memory",
36+
)

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

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
19+
20+
def test_generate_memories(client):
21+
agent_engine = client.agent_engines.create()
22+
assert not list(
23+
client.agent_engines.list_memories(
24+
name=agent_engine.api_resource.name,
25+
)
26+
)
27+
client.agent_engines.generate_memories(
28+
name=agent_engine.api_resource.name,
29+
direct_contents_source={
30+
"events": [
31+
{
32+
"content": {
33+
"role": "model",
34+
"parts": [
35+
{"text": "I am a software engineer focusing in security"}
36+
],
37+
}
38+
}
39+
]
40+
},
41+
scope={"user_id": "test-user-id"},
42+
)
43+
assert (
44+
len(
45+
list(
46+
client.agent_engines.list_memories(
47+
name=agent_engine.api_resource.name,
48+
)
49+
)
50+
)
51+
>= 1
52+
)
53+
54+
55+
pytestmark = pytest_helper.setup(
56+
file=__file__,
57+
globals_for_file=globals(),
58+
test_method="agent_engines.generate_memories",
59+
)

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
19+
20+
def test_get_memory(client):
21+
agent_engine = client.agent_engines.create()
22+
operation = client.agent_engines.create_memory(
23+
name=agent_engine.api_resource.name,
24+
fact="memory_fact",
25+
scope={"user_id": "123"},
26+
)
27+
memory = client.agent_engines.get_memory(
28+
name=operation.response.name,
29+
)
30+
assert memory.name == operation.response.name
31+
32+
33+
pytestmark = pytest_helper.setup(
34+
file=__file__,
35+
globals_for_file=globals(),
36+
test_method="agent_engines.get_memory",
37+
)

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
19+
20+
def test_list_memories(client):
21+
agent_engine = client.agent_engines.create()
22+
assert not list(
23+
client.agent_engines.list_memories(
24+
name=agent_engine.api_resource.name,
25+
)
26+
)
27+
client.agent_engines.create_memory(
28+
name=agent_engine.api_resource.name,
29+
fact="memory_fact",
30+
scope={"user_id": "123"},
31+
)
32+
assert (
33+
len(
34+
list(
35+
client.agent_engines.list_memories(
36+
name=agent_engine.api_resource.name,
37+
)
38+
)
39+
)
40+
== 1
41+
)
42+
43+
44+
pytestmark = pytest_helper.setup(
45+
file=__file__,
46+
globals_for_file=globals(),
47+
test_method="agent_engines.list_memories",
48+
)

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
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+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
19+
20+
def test_retrieve_memories(client):
21+
agent_engine = client.agent_engines.create()
22+
assert not list(
23+
client.agent_engines.retrieve_memories(
24+
name=agent_engine.api_resource.name,
25+
scope={"user_id": "123"},
26+
)
27+
)
28+
client.agent_engines.create_memory(
29+
name=agent_engine.api_resource.name,
30+
fact="memory_fact_1",
31+
scope={"user_id": "123"},
32+
)
33+
assert (
34+
len(
35+
list(
36+
client.agent_engines.retrieve_memories(
37+
name=agent_engine.api_resource.name,
38+
scope={"user_id": "123"},
39+
)
40+
)
41+
)
42+
== 1
43+
)
44+
assert not list(
45+
client.agent_engines.retrieve_memories(
46+
name=agent_engine.api_resource.name,
47+
scope={"user_id": "456"},
48+
)
49+
)
50+
client.agent_engines.create_memory(
51+
name=agent_engine.api_resource.name,
52+
fact="memory_fact_2",
53+
scope={"user_id": "123"},
54+
)
55+
assert (
56+
len(
57+
list(
58+
client.agent_engines.retrieve_memories(
59+
name=agent_engine.api_resource.name,
60+
scope={"user_id": "123"},
61+
)
62+
)
63+
)
64+
== 2
65+
)
66+
67+
68+
pytestmark = pytest_helper.setup(
69+
file=__file__,
70+
globals_for_file=globals(),
71+
test_method="agent_engines.create_memory",
72+
)

0 commit comments

Comments
 (0)