feat: Add async methods for Sessions and Session Events

PiperOrigin-RevId: 818847205

Author: vertex-sdk-bot

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

@@ -15,30 +15,31 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 import datetime
+import pytest
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
 
 
 def test_list_session_events(client):
     agent_engine = client.agent_engines.create()
-    operation = client.agent_engines.create_session(
+    operation = client.agent_engines.sessions.create(
         name=agent_engine.api_resource.name,
         user_id="test-user-123",
     )
     session = operation.response
     assert not list(
-        client.agent_engines.list_session_events(
+        client.agent_engines.sessions.events.list(
             name=session.name,
         )
     )
-    client.agent_engines.append_session_event(
+    client.agent_engines.sessions.events.append(
         name=session.name,
         author="test-user-123",
         invocation_id="test-invocation-id",
         timestamp=datetime.datetime.fromtimestamp(1234567890, tz=datetime.timezone.utc),
     )
-    session_event_list = client.agent_engines.list_session_events(
+    session_event_list = client.agent_engines.sessions.events.list(
         name=session.name,
     )
     assert len(session_event_list) == 1
@@ -48,5 +49,33 @@ def test_list_session_events(client):
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
-    test_method="agent_engines.list_session_events",
+    test_method="agent_engines.sessions.events.list",
 )
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_async_list_session_events(client):
+    agent_engine = client.agent_engines.create()
+    operation = await client.aio.agent_engines.sessions.create(
+        name=agent_engine.api_resource.name,
+        user_id="test-user-123",
+    )
+    session = operation.response
+    pager = await client.aio.agent_engines.sessions.events.list(name=session.name)
+    assert not [item async for item in pager]
+
+    await client.aio.agent_engines.sessions.events.append(
+        name=session.name,
+        author="test-user-123",
+        invocation_id="test-invocation-id",
+        timestamp=datetime.datetime.fromtimestamp(1234567890, tz=datetime.timezone.utc),
+    )
+    pager = await client.aio.agent_engines.sessions.events.list(name=session.name)
+    session_event_list = [item async for item in pager]
+    assert len(session_event_list) == 1
+    assert isinstance(session_event_list[0], types.SessionEvent)
+
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)

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

@@ -14,30 +14,59 @@
 #
 # 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
 
 
 def test_list_sessions(client):
     agent_engine = client.agent_engines.create()
     assert not list(
-        client.agent_engines.list_sessions(
+        client.agent_engines.sessions.list(
             name=agent_engine.api_resource.name,
         )
     )
-    client.agent_engines.create_session(
+    client.agent_engines.sessions.create(
         name=agent_engine.api_resource.name,
         user_id="test-user-123",
     )
-    session_list = client.agent_engines.list_sessions(
+    session_list = client.agent_engines.sessions.list(
         name=agent_engine.api_resource.name,
     )
     assert len(session_list) == 1
     assert isinstance(session_list[0], types.Session)
 
+    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.list_sessions",
+    test_method="agent_engines.sessions.list",
 )
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_async_list_sessions(client):
+    agent_engine = client.agent_engines.create()
+    pager = await client.aio.agent_engines.sessions.list(
+        name=agent_engine.api_resource.name
+    )
+    assert not [item async for item in pager]
+
+    await client.aio.agent_engines.sessions.create(
+        name=agent_engine.api_resource.name,
+        user_id="test-user-123",
+    )
+    pager = await client.aio.agent_engines.sessions.list(
+        name=agent_engine.api_resource.name,
+    )
+    session_list = [item async for item in pager]
+    assert len(session_list) == 1
+    assert isinstance(session_list[0], types.Session)
+
+    client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)

vertexai/_genai/_agent_engines_utils.py

@@ -30,6 +30,7 @@
 from typing import (
     Any,
     AsyncIterator,
+    Awaitable,
     Callable,
     Coroutine,
     Dict,
@@ -391,6 +392,41 @@ def __call__(self, *, operation_name: str, **kwargs) -> AgentEngineOperationUnio
         pass
 
 
+class GetAsyncOperationFunction(Protocol):
+    async def __call__(
+        self, *, operation_name: str, **kwargs
+    ) -> Awaitable[AgentEngineOperationUnion]:
+        pass
+
+
+async def _await_async_operation(
+    *,
+    operation_name: str,
+    get_operation_fn: GetAsyncOperationFunction,
+    poll_interval_seconds: float = 10,
+) -> Any:
+    """Waits for the operation for creating an agent engine to complete.
+
+    Args:
+        operation_name (str):
+            Required. The name of the operation for creating the Agent Engine.
+        poll_interval_seconds (float):
+            The number of seconds to wait between each poll.
+        get_operation_fn (Callable[[str], Awaitable[Any]]):
+            Optional. The async function to use for getting the operation. If not
+            provided, `self._get_agent_operation` will be used.
+
+    Returns:
+        The operation that has completed (i.e. `operation.done==True`).
+    """
+    operation = await get_operation_fn(operation_name=operation_name)
+    while not operation.done:
+        await asyncio.sleep(poll_interval_seconds)
+        operation = await get_operation_fn(operation_name=operation.name)
+
+    return operation
+
+
 def _await_operation(
     *,
     operation_name: str,

vertexai/_genai/session_events.py

@@ -26,7 +26,7 @@
 from google.genai import _common
 from google.genai._common import get_value_by_path as getv
 from google.genai._common import set_value_by_path as setv
-from google.genai.pagers import Pager
+from google.genai.pagers import AsyncPager, Pager
 
 from . import types
 
@@ -442,3 +442,32 @@ async def _list(
 
         self._api_client._verify_response(return_value)
         return return_value
+
+    async def list(
+        self,
+        *,
+        name: str,
+        config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None,
+    ) -> AsyncPager[types.SessionEvent]:
+        """Lists Agent Engine session events.
+
+        Args:
+            name (str): Required. The name of the agent engine to list session
+                events for.
+            config (ListAgentEngineSessionEventsConfig): Optional. The configuration
+                for the session events to list. Currently, the `filter` field in
+                `config` only supports filtering by `timestamp`. The timestamp
+                value must be enclosed in double quotes and include the time zone
+                information. For example:
+                `config={'filter': 'timestamp>="2025-08-07T19:44:38.4Z"'}`.
+
+        Returns:
+            AsyncPager[SessionEvent]: An async pager of session events.
+        """
+
+        return AsyncPager(
+            "session_events",
+            functools.partial(self._list, name=name),
+            await self._list(name=name, config=config),
+            config,
+        )

vertexai/_genai/sessions.py

@@ -26,7 +26,7 @@
 from google.genai import _common
 from google.genai._common import get_value_by_path as getv
 from google.genai._common import set_value_by_path as setv
-from google.genai.pagers import Pager
+from google.genai.pagers import AsyncPager, Pager
 
 from . import _agent_engines_utils
 from . import types
@@ -1078,3 +1078,95 @@ async def _update(
 
         self._api_client._verify_response(return_value)
         return return_value
+
+    _events = None
+
+    @property
+    @_common.experimental_warning(
+        "The Vertex SDK GenAI agent_engines.sessions.events module is "
+        "experimental, and may change in future versions."
+    )
+    def events(self):
+        if self._events is None:
+            try:
+                # We need to lazy load the sessions.events module to handle the
+                # possibility of ImportError when dependencies are not installed.
+                self._events = importlib.import_module(".session_events", __package__)
+            except ImportError as e:
+                raise ImportError(
+                    "The 'agent_engines.sessions.events' module requires"
+                    "additional packages. Please install them using pip install "
+                    "google-cloud-aiplatform[agent_engines]"
+                ) from e
+        return self._events.AsyncSessionEvents(self._api_client)
+
+    async def create(
+        self,
+        *,
+        name: str,
+        user_id: str,
+        config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None,
+    ) -> types.AgentEngineSessionOperation:
+        """Creates a new session in the Agent Engine.
+
+        Args:
+            name (str):
+                Required. The name of the agent engine to create the session for.
+            user_id (str):
+                Required. The user ID of the session.
+            config (CreateAgentEngineSessionConfig):
+                Optional. The configuration for the session to create.
+
+        Returns:
+            AgentEngineSessionOperation: The operation for creating the session.
+        """
+        if config is None:
+            config = types.CreateAgentEngineSessionConfig()
+        elif isinstance(config, dict):
+            config = types.CreateAgentEngineSessionConfig.model_validate(config)
+        operation = await self._create(
+            name=name,
+            user_id=user_id,
+            config=config,
+        )
+        if config.wait_for_completion and not operation.done:
+            operation = await _agent_engines_utils._await_async_operation(
+                operation_name=operation.name,
+                get_operation_fn=self._get_session_operation,
+                poll_interval_seconds=0.5,
+            )
+            if operation.response:
+                operation.response = await self.get(name=operation.response.name)
+            elif operation.error:
+                raise RuntimeError(f"Failed to create session: {operation.error}")
+            else:
+                raise RuntimeError(
+                    "Error retrieving session from the operation response. "
+                    f"Operation name: {operation.name}"
+                )
+        return operation
+
+    async def list(
+        self,
+        *,
+        name: str,
+        config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None,
+    ) -> AsyncPager[types.Session]:
+        """Lists Agent Engine sessions.
+
+        Args:
+            name (str): Required. The name of the agent engine to list sessions
+                for.
+            config (ListAgentEngineSessionConfig): Optional. The configuration
+                for the sessions to list.
+
+        Returns:
+            AsyncPager[Session]: An async pager of sessions.
+        """
+
+        return AsyncPager(
+            "sessions",
+            functools.partial(self._list, name=name),
+            await self._list(name=name, config=config),
+            config,
+        )