1616import asyncio
1717import logging
1818import re
19- import time
2019from typing import Any
20+ from typing import Dict
2121from typing import Optional
2222import urllib .parse
2323
@@ -50,9 +50,6 @@ def __init__(
5050 self .project = project
5151 self .location = location
5252
53- client = genai .Client (vertexai = True , project = project , location = location )
54- self .api_client = client ._api_client
55-
5653 @override
5754 async def create_session (
5855 self ,
@@ -86,6 +83,7 @@ async def create_session(
8683 operation_id = api_response ['name' ].split ('/' )[- 1 ]
8784
8885 max_retry_attempt = 5
86+ lro_response = None
8987 while max_retry_attempt >= 0 :
9088 lro_response = await api_client .async_request (
9189 http_method = 'GET' ,
@@ -99,6 +97,11 @@ async def create_session(
9997 await asyncio .sleep (1 )
10098 max_retry_attempt -= 1
10199
100+ if lro_response is None or not lro_response .get ('done' , None ):
101+ raise TimeoutError (
102+ f'Timeout waiting for operation { operation_id } to complete.'
103+ )
104+
102105 # Get session resource
103106 get_session_api_response = await api_client .async_request (
104107 http_method = 'GET' ,
@@ -235,11 +238,15 @@ async def delete_session(
235238 ) -> None :
236239 reasoning_engine_id = _parse_reasoning_engine_id (app_name )
237240 api_client = _get_api_client (self .project , self .location )
238- await api_client .async_request (
239- http_method = 'DELETE' ,
240- path = f'reasoningEngines/{ reasoning_engine_id } /sessions/{ session_id } ' ,
241- request_dict = {},
242- )
241+ try :
242+ await api_client .async_request (
243+ http_method = 'DELETE' ,
244+ path = f'reasoningEngines/{ reasoning_engine_id } /sessions/{ session_id } ' ,
245+ request_dict = {},
246+ )
247+ except Exception as e :
248+ logger .error (f'Error deleting session { session_id } : { e } ' )
249+ raise e
243250
244251 @override
245252 async def append_event (self , session : Session , event : Event ) -> Event :
@@ -266,7 +273,7 @@ def _get_api_client(project: str, location: str):
266273 return client ._api_client
267274
268275
269- def _convert_event_to_json (event : Event ):
276+ def _convert_event_to_json (event : Event ) -> Dict [ str , Any ] :
270277 metadata_json = {
271278 'partial' : event .partial ,
272279 'turn_complete' : event .turn_complete ,
@@ -318,7 +325,7 @@ def _convert_event_to_json(event: Event):
318325 return event_json
319326
320327
321- def _from_api_event (api_event : dict ) -> Event :
328+ def _from_api_event (api_event : Dict [ str , Any ] ) -> Event :
322329 event_actions = EventActions ()
323330 if api_event .get ('actions' , None ):
324331 event_actions = EventActions (
0 commit comments