feat: GenAI SDK client: Add async Memory and Memory Revisions methods

PiperOrigin-RevId: 821824583

Author: vertex-sdk-bot

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

@@ -23,15 +23,16 @@
 
 def test_delete_memory(client):
  agent_engine = client.agent_engines.create()
- operation = client.agent_engines.create_memory(
+ operation = client.agent_engines.memories.create(
  name=agent_engine.api_resource.name,
  fact="memory_fact",
  scope={"user_id": "123"},
  )
  memory = operation.response
- operation = client.agent_engines.delete_memory(name=memory.name)
+ operation = client.agent_engines.memories.delete(name=memory.name)
  assert isinstance(operation, types.DeleteAgentEngineMemoryOperation)
  assert operation.name.startswith(memory.name + "/operations/")
+ client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
 
 
 pytestmark = pytest_helper.setup(
@@ -46,14 +47,16 @@ def test_delete_memory(client):
 
 @pytest.mark.asyncio
 async def test_delete_memory_async(client):
- # TODO(b/431785750): use async methods for create() and create_memory() when available
  agent_engine = client.agent_engines.create()
- operation = client.agent_engines.create_memory(
+ operation = await client.aio.agent_engines.memories.create(
  name=agent_engine.api_resource.name,
  fact="memory_fact",
  scope={"user_id": "123"},
  )
  memory = operation.response
- operation = await client.aio.agent_engines.delete_memory(name=memory.name)
+ operation = await client.aio.agent_engines.memories.delete(name=memory.name)
  assert isinstance(operation, types.DeleteAgentEngineMemoryOperation)
  assert operation.name.startswith(memory.name + "/operations/")
+ await client.aio.agent_engines.delete(
+ name=agent_engine.api_resource.name, force=True
+ )

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

@@ -14,12 +14,16 @@
 #
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
+import pytest
+
+
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
 from google.genai import types as genai_types
 
 
 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/"
  )
@@ -146,3 +150,64 @@ def test_generate_memories_direct_memories_source(client):
  globals_for_file=globals(),
  test_method="agent_engines.generate_memories",
 )
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@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,
+ 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),
+ )
+ memories_pager = await client.aio.agent_engines.memories.list(
+ name=agent_engine.api_resource.name
+ )
+ memory_list = [item async for item in memories_pager]
+ assert len(memory_list) >= 1
+
+ revisions_pager = await client.aio.agent_engines.memories.revisions.list(
+ name=memory_list[0].name
+ )
+ memory_revisions = [item async for item in revisions_pager]
+ assert len(memory_revisions) >= 1
+ revision_name = memory_revisions[0].name
+
+ # Update the memory.
+ client.agent_engines.memories._update(
+ name=memory_list[0].name,
+ fact="This is temporary",
+ scope={"user_id": "test-user-id"},
+ )
+ memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name)
+ assert memory.fact == "This is temporary"
+
+ # Rollback to the revision with the original fact that was created by the
+ # generation request.
+ await client.aio.agent_engines.memories.rollback(
+ name=memory_list[0].name,
+ target_revision_id=revision_name.split("/")[-1],
+ )
+ memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name)
+ assert memory.fact == memory_revisions[0].fact
+
+ await client.aio.agent_engines.delete(
+ name=agent_engine.api_resource.name, force=True
+ )

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

@@ -22,23 +22,24 @@
 
 def test_get_memory(client):
  agent_engine = client.agent_engines.create()
- operation = client.agent_engines.create_memory(
+ operation = client.agent_engines.memories.create(
  name=agent_engine.api_resource.name,
  fact="memory_fact",
  scope={"user_id": "123"},
  )
  assert isinstance(operation, types.AgentEngineMemoryOperation)
- memory = client.agent_engines.get_memory(
+ memory = client.agent_engines.memories.get(
  name=operation.response.name,
  )
  assert isinstance(memory, types.Memory)
  assert memory.name == operation.response.name
+ 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.get_memory",
+ test_method="agent_engines.memories.get",
 )
 
 
@@ -47,16 +48,18 @@ def test_get_memory(client):
 
 @pytest.mark.asyncio
 async def test_get_memory_async(client):
- # TODO(b/431785750): use async methods for create() and create_memory() when available
  agent_engine = client.agent_engines.create()
- operation = client.agent_engines.create_memory(
+ operation = await client.aio.agent_engines.memories.create(
  name=agent_engine.api_resource.name,
  fact="memory_fact",
  scope={"user_id": "123"},
  )
  assert isinstance(operation, types.AgentEngineMemoryOperation)
- memory = await client.aio.agent_engines.get_memory(
+ memory = await client.aio.agent_engines.memories.get(
  name=operation.response.name,
  )
  assert isinstance(memory, types.Memory)
  assert memory.name == operation.response.name
+ await client.aio.agent_engines.delete(
+ name=agent_engine.api_resource.name, force=True
+ )

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

@@ -14,6 +14,8 @@
 #
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
+import pytest
+
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
 
@@ -59,5 +61,36 @@ def test_list_memories(client):
 pytestmark = pytest_helper.setup(
  file=__file__,
  globals_for_file=globals(),
- test_method="agent_engines.list_memories",
+ test_method="agent_engines.memories.list",
 )
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_async_list_memories(client):
+ agent_engine = client.agent_engines.create()
+ pager = await client.aio.agent_engines.memories.list(
+ name=agent_engine.api_resource.name
+ )
+ assert not [item async for item in pager]
+
+ await client.aio.agent_engines.memories.create(
+ name=agent_engine.api_resource.name,
+ fact="memory_fact_2",
+ scope={"user_id": "456"},
+ config={
+ "wait_for_completion": True,
+ },
+ )
+ pager = await client.aio.agent_engines.memories.list(
+ name=agent_engine.api_resource.name
+ )
+ memory_list = [item async for item in pager]
+ assert len(memory_list) == 1
+ assert isinstance(memory_list[0], types.Memory)
+
+ await client.aio.agent_engines.delete(
+ name=agent_engine.api_resource.name, force=True
+ )

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

@@ -14,6 +14,9 @@
 #
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
+import pytest
+
+
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
 from google.genai import pagers
@@ -22,23 +25,23 @@
 def test_retrieve_memories_with_similarity_search_params(client):
  agent_engine = client.agent_engines.create()
  assert not list(
- client.agent_engines.retrieve_memories(
+ client.agent_engines.memories.retrieve(
  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(
+ client.agent_engines.memories.create(
  name=agent_engine.api_resource.name,
  fact="memory_fact_1",
  scope={"user_id": "123"},
  )
  assert (
  len(
  list(
- client.agent_engines.retrieve_memories(
+ client.agent_engines.memories.retrieve(
  name=agent_engine.api_resource.name,
  scope={"user_id": "123"},
  )
@@ -47,20 +50,20 @@ def test_retrieve_memories_with_similarity_search_params(client):
  == 1
  )
  assert not list(
- client.agent_engines.retrieve_memories(
+ client.agent_engines.memories.retrieve(
  name=agent_engine.api_resource.name,
  scope={"user_id": "456"},
  )
  )
- client.agent_engines.create_memory(
+ client.agent_engines.memories.create(
  name=agent_engine.api_resource.name,
  fact="memory_fact_2",
  scope={"user_id": "123"},
  )
  assert (
  len(
  list(
- client.agent_engines.retrieve_memories(
+ client.agent_engines.memories.retrieve(
  name=agent_engine.api_resource.name,
  scope={"user_id": "123"},
  )
@@ -74,12 +77,12 @@ def test_retrieve_memories_with_similarity_search_params(client):
 
 def test_retrieve_memories_with_simple_retrieval_params(client):
  agent_engine = client.agent_engines.create()
- client.agent_engines.create_memory(
+ client.agent_engines.memories.create(
  name=agent_engine.api_resource.name,
  fact="memory_fact_1",
  scope={"user_id": "123"},
  )
- memories = client.agent_engines.retrieve_memories(
+ memories = client.agent_engines.memories.retrieve(
  name=agent_engine.api_resource.name,
  scope={"user_id": "123"},
  simple_retrieval_params=types.RetrieveMemoriesRequestSimpleRetrievalParams(
@@ -98,3 +101,27 @@ def test_retrieve_memories_with_simple_retrieval_params(client):
  globals_for_file=globals(),
  test_method="agent_engines.create_memory",
 )
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_retrieve_memories_async(client):
+ agent_engine = client.agent_engines.create()
+ operation = await client.aio.agent_engines.memories.create(
+ name=agent_engine.api_resource.name,
+ fact="memory_fact",
+ scope={"user_id": "123"},
+ )
+ assert isinstance(operation, types.AgentEngineMemoryOperation)
+ pager = await client.aio.agent_engines.memories.retrieve(
+ name=agent_engine.api_resource.name,
+ scope={"user_id": "123"},
+ )
+ memories = [item async for item in pager]
+ assert len(memories) == 1
+ assert isinstance(memories[0], types.RetrieveMemoriesResponseRetrievedMemory)
+ await client.aio.agent_engines.delete(
+ name=agent_engine.api_resource.name, force=True
+ )