14
14
from __future__ import annotations
15
15
16
16
import asyncio
17
+ import json
17
18
import logging
18
19
import re
19
20
from typing import Any
@@ -87,6 +88,7 @@ async def create_session(
87
88
path = f'reasoningEngines/{ reasoning_engine_id } /sessions' ,
88
89
request_dict = session_json_dict ,
89
90
)
91
+ api_response = _convert_api_response (api_response )
90
92
logger .info (f'Create Session response { api_response } ' )
91
93
92
94
session_id = api_response ['name' ].split ('/' )[- 3 ]
@@ -100,6 +102,7 @@ async def create_session(
100
102
path = f'operations/{ operation_id } ' ,
101
103
request_dict = {},
102
104
)
105
+ lro_response = _convert_api_response (lro_response )
103
106
104
107
if lro_response .get ('done' , None ):
105
108
break
@@ -118,6 +121,7 @@ async def create_session(
118
121
path = f'reasoningEngines/{ reasoning_engine_id } /sessions/{ session_id } ' ,
119
122
request_dict = {},
120
123
)
124
+ get_session_api_response = _convert_api_response (get_session_api_response )
121
125
122
126
update_timestamp = isoparse (
123
127
get_session_api_response ['updateTime' ]
@@ -149,6 +153,7 @@ async def get_session(
149
153
path = f'reasoningEngines/{ reasoning_engine_id } /sessions/{ session_id } ' ,
150
154
request_dict = {},
151
155
)
156
+ get_session_api_response = _convert_api_response (get_session_api_response )
152
157
153
158
session_id = get_session_api_response ['name' ].split ('/' )[- 1 ]
154
159
update_timestamp = isoparse (
@@ -167,9 +172,12 @@ async def get_session(
167
172
path = f'reasoningEngines/{ reasoning_engine_id } /sessions/{ session_id } /events' ,
168
173
request_dict = {},
169
174
)
175
+ list_events_api_response = _convert_api_response (list_events_api_response )
170
176
171
177
# Handles empty response case
172
- if list_events_api_response .get ('httpHeaders' , None ):
178
+ if not list_events_api_response or list_events_api_response .get (
179
+ 'httpHeaders' , None
180
+ ):
173
181
return session
174
182
175
183
session .events += [
@@ -226,9 +234,10 @@ async def list_sessions(
226
234
path = path ,
227
235
request_dict = {},
228
236
)
237
+ api_response = _convert_api_response (api_response )
229
238
230
239
# Handles empty response case
231
- if api_response .get ('httpHeaders' , None ):
240
+ if not api_response or api_response .get ('httpHeaders' , None ):
232
241
return ListSessionsResponse ()
233
242
234
243
sessions = []
@@ -303,6 +312,13 @@ def _get_api_client(self):
303
312
return client ._api_client
304
313
305
314
315
+ def _convert_api_response (api_response ):
316
+ """Converts the API response to a JSON object based on the type."""
317
+ if hasattr (api_response , 'body' ):
318
+ return json .loads (api_response .body )
319
+ return api_response
320
+
321
+
306
322
def _convert_event_to_json (event : Event ) -> Dict [str , Any ]:
307
323
metadata_json = {
308
324
'partial' : event .partial ,
0 commit comments