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:19:25 -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
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,