chore: GenAI SDK client - Add more types coverage for Agent Engines replay tests

PiperOrigin-RevId: 781984631

Author: sararob

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

@@ -15,15 +15,21 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
 
 
 def test_create_memory(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.create_memory(
         name=agent_engine.api_resource.name,
         fact="memory_fact",
         scope={"user_id": "123"},
+        config=types.AgentEngineMemoryConfig(display_name="my_memory_fact"),
     )
+    assert isinstance(operation, types.AgentEngineMemoryOperation)
     assert operation.response.fact == "memory_fact"
     assert operation.response.scope == {"user_id": "123"}
     assert operation.response.name.startswith(agent_engine.api_resource.name)

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

@@ -15,6 +15,7 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
 
 
 def test_delete_memory(client):
@@ -26,6 +27,7 @@ def test_delete_memory(client):
     )
     memory = operation.response
     operation = client.agent_engines.delete_memory(name=memory.name)
+    assert isinstance(operation, types.DeleteAgentEngineMemoryOperation)
     assert operation.name.startswith(memory.name + "/operations/")
 
 

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

@@ -15,6 +15,7 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
 
 
 def test_generate_memories(client):
@@ -26,19 +27,21 @@ def test_generate_memories(client):
     )
     client.agent_engines.generate_memories(
         name=agent_engine.api_resource.name,
-        direct_contents_source={
-            "events": [
-                {
-                    "content": {
-                        "role": "model",
-                        "parts": [
-                            {"text": "I am a software engineer focusing in security"}
+        scope={"user_id": "test-user-id"},
+        direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource(
+            events=[
+                types.GenerateMemoriesRequestDirectContentsSourceEvent(
+                    content=types.Content(
+                        role="model",
+                        parts=[
+                            types.Part(
+                                text="I am a software engineer focusing in security"
+                            )
                         ],
-                    }
-                }
+                    )
+                )
             ]
-        },
-        scope={"user_id": "test-user-id"},
+        ),
     )
     assert (
         len(

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

@@ -15,6 +15,7 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
 
 
 def test_get_memory(client):
@@ -24,9 +25,11 @@ def test_get_memory(client):
         fact="memory_fact",
         scope={"user_id": "123"},
     )
+    assert isinstance(operation, types.AgentEngineMemoryOperation)
     memory = client.agent_engines.get_memory(
         name=operation.response.name,
     )
+    assert isinstance(memory, types.Memory)
     assert memory.name == operation.response.name
 
 

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

@@ -15,6 +15,7 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
 
 
 def test_list_memories(client):
@@ -29,16 +30,11 @@ def test_list_memories(client):
         fact="memory_fact",
         scope={"user_id": "123"},
     )
-    assert (
-        len(
-            list(
-                client.agent_engines.list_memories(
-                    name=agent_engine.api_resource.name,
-                )
-            )
-        )
-        == 1
+    memory_list = client.agent_engines.list_memories(
+        name=agent_engine.api_resource.name,
     )
+    assert len(memory_list) == 1
+    assert isinstance(memory_list[0], types.Memory)
 
 
 pytestmark = pytest_helper.setup(

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

@@ -15,14 +15,19 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+from google.genai import pagers
 
 
-def test_retrieve_memories(client):
+def test_retrieve_memories_with_similarity_search_params(client):
     agent_engine = client.agent_engines.create()
     assert not list(
         client.agent_engines.retrieve_memories(
             name=agent_engine.api_resource.name,
             scope={"user_id": "123"},
+            similarity_search_params=types.RetrieveMemoriesRequestSimilaritySearchParams(
+                search_query="memory_fact_1",
+            ),
         )
     )
     client.agent_engines.create_memory(
@@ -65,6 +70,25 @@ def test_retrieve_memories(client):
     )
 
 
+def test_retrieve_memories_with_simple_retrieval_params(client):
+    agent_engine = client.agent_engines.create()
+    client.agent_engines.create_memory(
+        name=agent_engine.api_resource.name,
+        fact="memory_fact_1",
+        scope={"user_id": "123"},
+    )
+    memories = client.agent_engines.retrieve_memories(
+        name=agent_engine.api_resource.name,
+        scope={"user_id": "123"},
+        simple_retrieval_params=types.RetrieveMemoriesRequestSimpleRetrievalParams(
+            page_size=1,
+        ),
+    )
+    assert isinstance(memories, pagers.Pager)
+    assert isinstance(memories.page[0], types.RetrieveMemoriesResponseRetrievedMemory)
+    assert memories.page_size == 1
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),

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

@@ -0,0 +1,44 @@
+# 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
+
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_agent_engines_update(client):
+    agent_engine = client.agent_engines.create()
+    assert agent_engine.api_resource.display_name is None
+
+    updated_agent_engine = client.agent_engines.update(
+        name=agent_engine.api_resource.name,
+        config=types.AgentEngineConfig(
+            display_name="updated_display_name",
+            description="updated description",
+        ),
+    )
+    assert isinstance(updated_agent_engine, types.AgentEngine)
+    assert updated_agent_engine.api_resource.name == agent_engine.api_resource.name
+
+    assert updated_agent_engine.api_resource.display_name == "updated_display_name"
+    assert updated_agent_engine.api_resource.description == "updated description"
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.update",
+)