16
16
import asyncio
17
17
import logging
18
18
import re
19
+ import time
19
20
from typing import Any
20
- from typing import Dict
21
21
from typing import Optional
22
22
import urllib .parse
23
23
@@ -50,6 +50,9 @@ def __init__(
50
50
self .project = project
51
51
self .location = location
52
52
53
+ client = genai .Client (vertexai = True , project = project , location = location )
54
+ self .api_client = client ._api_client
55
+
53
56
@override
54
57
async def create_session (
55
58
self ,
@@ -83,7 +86,6 @@ async def create_session(
83
86
operation_id = api_response ['name' ].split ('/' )[- 1 ]
84
87
85
88
max_retry_attempt = 5
86
- lro_response = None
87
89
while max_retry_attempt >= 0 :
88
90
lro_response = await api_client .async_request (
89
91
http_method = 'GET' ,
@@ -97,11 +99,6 @@ async def create_session(
97
99
await asyncio .sleep (1 )
98
100
max_retry_attempt -= 1
99
101
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
-
105
102
# Get session resource
106
103
get_session_api_response = await api_client .async_request (
107
104
http_method = 'GET' ,
@@ -238,15 +235,11 @@ async def delete_session(
238
235
) -> None :
239
236
reasoning_engine_id = _parse_reasoning_engine_id (app_name )
240
237
api_client = _get_api_client (self .project , self .location )
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
238
+ await api_client .async_request (
239
+ http_method = 'DELETE' ,
240
+ path = f'reasoningEngines/{ reasoning_engine_id } /sessions/{ session_id } ' ,
241
+ request_dict = {},
242
+ )
250
243
251
244
@override
252
245
async def append_event (self , session : Session , event : Event ) -> Event :
@@ -273,7 +266,7 @@ def _get_api_client(project: str, location: str):
273
266
return client ._api_client
274
267
275
268
276
- def _convert_event_to_json (event : Event ) -> Dict [ str , Any ] :
269
+ def _convert_event_to_json (event : Event ):
277
270
metadata_json = {
278
271
'partial' : event .partial ,
279
272
'turn_complete' : event .turn_complete ,
@@ -325,7 +318,7 @@ def _convert_event_to_json(event: Event) -> Dict[str, Any]:
325
318
return event_json
326
319
327
320
328
- def _from_api_event (api_event : Dict [ str , Any ] ) -> Event :
321
+ def _from_api_event (api_event : dict ) -> Event :
329
322
event_actions = EventActions ()
330
323
if api_event .get ('actions' , None ):
331
324
event_actions = EventActions (
0 commit comments