Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions agentplatform/agent_engines/templates/adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -1568,6 +1568,10 @@ async def async_create_session(
will be be generated for the session.
state (dict[str, Any]):
Optional. The initial state of the session.
ttl (str):
Optional. The time-to-live for the session.
expire_time (str):
Optional. The expiration time for the session.
**kwargs (dict[str, Any]):
Optional. Additional keyword arguments to pass to the
session service.
Expand Down
46 changes: 46 additions & 0 deletions tests/unit/vertex_adk/test_agent_engine_templates_adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
from vertexai.agent_engines.templates import adk as adk_template
from google.genai import types
import pytest
from google.adk.sessions.base_session_service import BaseSessionService


try:
Expand Down Expand Up @@ -554,6 +555,51 @@ async def test_async_create_session(self, get_project_id_mock: mock.Mock):
assert session2["user_id"] == _TEST_USER_ID
assert session2["id"] == "test_session_id"

@pytest.mark.asyncio
async def test_async_create_session_with_ttl_kwargs(
self, get_project_id_mock: mock.Mock
):
mock_session_service = mock.Mock(spec=BaseSessionService)
mock_session_service.create_session = mock.AsyncMock()
app = agent_engines.AdkApp(
agent=_TEST_AGENT,
session_service_builder=lambda: mock_session_service,
)
await app.async_create_session(
user_id=_TEST_USER_ID,
ttl="7200s",
)
mock_session_service.create_session.assert_called_once_with(
app_name=app._app_name(),
user_id=_TEST_USER_ID,
session_id=None,
state=None,
ttl="7200s",
)

@pytest.mark.asyncio
async def test_async_create_session_with_expire_time_kwargs(
self, get_project_id_mock: mock.Mock
):
mock_session_service = mock.Mock(spec=BaseSessionService)
mock_session_service.create_session = mock.AsyncMock()
app = agent_engines.AdkApp(
agent=_TEST_AGENT,
session_service_builder=lambda: mock_session_service,
)
await app.async_create_session(
user_id=_TEST_USER_ID,
expire_time="2026-03-01T00:00:00Z",
)
mock_session_service.create_session.assert_called_once_with(
app_name=app._app_name(),
user_id=_TEST_USER_ID,
session_id=None,
state=None,
expire_time="2026-03-01T00:00:00Z",
)


@pytest.mark.asyncio
async def test_async_get_session(self, get_project_id_mock: mock.Mock):
app = agent_engines.AdkApp(agent=_TEST_AGENT)
Expand Down
4 changes: 4 additions & 0 deletions vertexai/agent_engines/templates/adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -1568,6 +1568,10 @@ async def async_create_session(
will be be generated for the session.
state (dict[str, Any]):
Optional. The initial state of the session.
ttl (str):
Optional. The time-to-live for the session.
expire_time (str):
Optional. The expiration time for the session.
**kwargs (dict[str, Any]):
Optional. Additional keyword arguments to pass to the
session service.
Expand Down
Loading