encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

chore: Make VertexAiSessionService fully asynchronous

Browse Source 
This CL refactors VertexAiSessionService to use the asynchronous aio client for all Vertex AI API calls. This ensures that the service methods are non-blocking and can be used effectively in an asyncio environment.

PiperOrigin-RevId: 824573356
This commit is contained in:
Shangjie Chen
2025-10-27 10:18:58 -07:00
committed by Copybara-Service
parent 48ddd07894
commit f7e2a7a40e
2 changed files with 47 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/sessions/vertex_ai_session_service.py
+30 -26
View File
@@ -13,6 +13,7 @@
# limitations under the License.
from __future__ import annotations
import asyncio
import datetime
import json
import logging
@@ -104,7 +105,7 @@ class VertexAiSessionService(BaseSessionService):
if _is_vertex_express_mode(self._project, self._location):
config['wait_for_completion'] = False
api_response = api_client.agent_engines.sessions.create(
api_response = await api_client.aio.agent_engines.sessions.create(
name=f'reasoningEngines/{reasoning_engine_id}',
user_id=user_id,
config=config,
@@ -123,7 +124,7 @@ class VertexAiSessionService(BaseSessionService):
)
async def _poll_session_resource():
try:
return api_client.agent_engines.sessions.get(
return await api_client.aio.agent_engines.sessions.get(
name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
)
except ClientError:
@@ -135,11 +136,11 @@ class VertexAiSessionService(BaseSessionService):
except Exception as exc:
raise ValueError('Failed to create session.') from exc
get_session_response = api_client.agent_engines.sessions.get(
get_session_response = await api_client.aio.agent_engines.sessions.get(
name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
)
else:
api_response = api_client.agent_engines.sessions.create(
api_response = await api_client.aio.agent_engines.sessions.create(
name=f'reasoningEngines/{reasoning_engine_id}',
user_id=user_id,
config=config,
@@ -168,10 +169,28 @@ class VertexAiSessionService(BaseSessionService):
) -> Optional[Session]:
reasoning_engine_id = self._get_reasoning_engine_id(app_name)
api_client = self._get_api_client()
session_resource_name = (
f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
)
# Get session resource
get_session_response = api_client.agent_engines.sessions.get(
name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
# Get session resource and events in parallel.
list_events_kwargs = {}
if config and not config.num_recent_events and config.after_timestamp:
# Filter events based on timestamp.
list_events_kwargs['config'] = {
'filter': 'timestamp>="{}"'.format(
datetime.datetime.fromtimestamp(
config.after_timestamp, tz=datetime.timezone.utc
).isoformat()
)
}
get_session_response, events_iterator = await asyncio.gather(
api_client.aio.agent_engines.sessions.get(name=session_resource_name),
api_client.aio.agent_engines.sessions.events.list(
name=session_resource_name,
**list_events_kwargs,
),
)
if get_session_response.user_id != user_id:
@@ -187,29 +206,14 @@ class VertexAiSessionService(BaseSessionService):
state=getattr(get_session_response, 'session_state', None) or {},
last_update_time=update_timestamp,
)
list_events_kwargs = {}
if config and not config.num_recent_events and config.after_timestamp:
list_events_kwargs['config'] = {
'filter': 'timestamp>="{}"'.format(
datetime.datetime.fromtimestamp(
config.after_timestamp, tz=datetime.timezone.utc
).isoformat()
)
}
events_iterator = api_client.agent_engines.sessions.events.list(
name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
**list_events_kwargs,
)
session.events += [
_from_api_event(event)
for event in events_iterator
if event.timestamp.timestamp() <= update_timestamp
]
# Filter events based on config
if config:
# Filter events based on num_recent_events.
if config.num_recent_events:
session.events = session.events[-config.num_recent_events :]
@@ -226,7 +230,7 @@ class VertexAiSessionService(BaseSessionService):
config = {}
if user_id is not None:
config['filter'] = f'user_id="{user_id}"'
sessions_iterator = api_client.agent_engines.sessions.list(
sessions_iterator = await api_client.aio.agent_engines.sessions.list(
name=f'reasoningEngines/{reasoning_engine_id}',
config=config,
)
@@ -251,7 +255,7 @@ class VertexAiSessionService(BaseSessionService):
api_client = self._get_api_client()
try:
api_client.agent_engines.sessions.delete(
await api_client.aio.agent_engines.sessions.delete(
name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
)
except Exception as e:
@@ -308,7 +312,7 @@ class VertexAiSessionService(BaseSessionService):
)
config['event_metadata'] = metadata_dict
api_client.agent_engines.sessions.events.append(
await api_client.aio.agent_engines.sessions.events.append(
name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}',
author=event.author,
invocation_id=event.invocation_id,
tests/unittests/sessions/test_vertex_ai_session_service.py
+17 -13
View File
@@ -235,22 +235,24 @@ class MockApiClient:
"""Initializes MockClient."""
self.session_dict: dict[str, Any] = {}
self.event_dict: dict[str, Tuple[List[Any], Optional[str]]] = {}
self.agent_engines = mock.Mock()
self.agent_engines.sessions.get.side_effect = self._get_session
self.agent_engines.sessions.list.side_effect = self._list_sessions
self.agent_engines.sessions.delete.side_effect = self._delete_session
self.agent_engines.sessions.create.side_effect = self._create_session
self.agent_engines.sessions.events.list.side_effect = self._list_events
self.agent_engines.sessions.events.append.side_effect = self._append_event
self.aio = mock.Mock()
self.aio.agent_engines.sessions.get.side_effect = self._get_session
self.aio.agent_engines.sessions.list.side_effect = self._list_sessions
self.aio.agent_engines.sessions.delete.side_effect = self._delete_session
self.aio.agent_engines.sessions.create.side_effect = self._create_session
self.aio.agent_engines.sessions.events.list.side_effect = self._list_events
self.aio.agent_engines.sessions.events.append.side_effect = (
self._append_event
)
self.last_create_session_config: dict[str, Any] = {}
def _get_session(self, name: str):
async def _get_session(self, name: str):
session_id = name.split('/')[-1]
if session_id in self.session_dict:
return _convert_to_object(self.session_dict[session_id])
raise api_core_exceptions.NotFound(f'Session not found: {session_id}')
def _list_sessions(self, name: str, config: dict[str, Any]):
async def _list_sessions(self, name: str, config: dict[str, Any]):
filter_val = config.get('filter', '')
user_id_match = re.search(r'user_id="([^"]+)"', filter_val)
if user_id_match:
@@ -271,11 +273,13 @@ class MockApiClient:
_convert_to_object(session) for session in self.session_dict.values()
]
def _delete_session(self, name: str):
async def _delete_session(self, name: str):
session_id = name.split('/')[-1]
self.session_dict.pop(session_id)
def _create_session(self, name: str, user_id: str, config: dict[str, Any]):
async def _create_session(
self, name: str, user_id: str, config: dict[str, Any]
):
self.last_create_session_config = config
new_session_id = '4'
self.session_dict[new_session_id] = {
@@ -299,7 +303,7 @@ class MockApiClient:
'response': self.session_dict['4'],
})
def _list_events(self, name: str, **kwargs):
async def _list_events(self, name: str, **kwargs):
session_id = name.split('/')[-1]
events = []
if session_id in self.event_dict:
@@ -322,7 +326,7 @@ class MockApiClient:
]
return [_convert_to_object(event) for event in events]
def _append_event(
async def _append_event(
self,
name: str,
author: str,