google-genai 1.19.0__py3-none-any.whl → 1.21.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -25,6 +25,7 @@ import copy
25
25
  from dataclasses import dataclass
26
26
  import datetime
27
27
  import http
28
+ import inspect
28
29
  import io
29
30
  import json
30
31
  import logging
@@ -34,7 +35,7 @@ import ssl
34
35
  import sys
35
36
  import threading
36
37
  import time
37
- from typing import Any, AsyncIterator, Optional, Tuple, Union
38
+ from typing import Any, AsyncIterator, Optional, TYPE_CHECKING, Tuple, Union
38
39
  from urllib.parse import urlparse
39
40
  from urllib.parse import urlunparse
40
41
 
@@ -48,6 +49,7 @@ import httpx
48
49
  from pydantic import BaseModel
49
50
  from pydantic import Field
50
51
  from pydantic import ValidationError
52
+ import tenacity
51
53
 
52
54
  from . import _common
53
55
  from . import errors
@@ -55,6 +57,23 @@ from . import version
55
57
  from .types import HttpOptions
56
58
  from .types import HttpOptionsDict
57
59
  from .types import HttpOptionsOrDict
60
+ from .types import HttpResponse as SdkHttpResponse
61
+ from .types import HttpRetryOptions
62
+
63
+
64
+ has_aiohttp = False
65
+ try:
66
+ import aiohttp
67
+
68
+ has_aiohttp = True
69
+ except ImportError:
70
+ pass
71
+
72
+ # internal comment
73
+
74
+
75
+ if TYPE_CHECKING:
76
+ from multidict import CIMultiDictProxy
58
77
 
59
78
 
60
79
  logger = logging.getLogger('google_genai._api_client')
@@ -81,8 +100,7 @@ def _get_env_api_key() -> Optional[str]:
81
100
  env_gemini_api_key = os.environ.get('GEMINI_API_KEY', None)
82
101
  if env_google_api_key and env_gemini_api_key:
83
102
  logger.warning(
84
- 'Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using'
85
- ' GOOGLE_API_KEY.'
103
+ 'Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using GOOGLE_API_KEY.'
86
104
  )
87
105
 
88
106
  return env_google_api_key or env_gemini_api_key or None
@@ -104,7 +122,9 @@ def _append_library_version_headers(headers: dict[str, str]) -> None:
104
122
  'x-goog-api-client' in headers
105
123
  and version_header_value not in headers['x-goog-api-client']
106
124
  ):
107
- headers['x-goog-api-client'] = f'{version_header_value} ' + headers['x-goog-api-client']
125
+ headers['x-goog-api-client'] = (
126
+ f'{version_header_value} ' + headers['x-goog-api-client']
127
+ )
108
128
  elif 'x-goog-api-client' not in headers:
109
129
  headers['x-goog-api-client'] = version_header_value
110
130
 
@@ -200,23 +220,11 @@ class HttpRequest:
200
220
  timeout: Optional[float] = None
201
221
 
202
222
 
203
- # TODO(b/394358912): Update this class to use a SDKResponse class that can be
204
- # generated and used for all languages.
205
- class BaseResponse(_common.BaseModel):
206
- http_headers: Optional[dict[str, str]] = Field(
207
- default=None, description='The http headers of the response.'
208
- )
209
-
210
- json_payload: Optional[Any] = Field(
211
- default=None, description='The json payload of the response.'
212
- )
213
-
214
-
215
223
  class HttpResponse:
216
224
 
217
225
  def __init__(
218
226
  self,
219
- headers: Union[dict[str, str], httpx.Headers],
227
+ headers: Union[dict[str, str], httpx.Headers, 'CIMultiDictProxy[str]'],
220
228
  response_stream: Union[Any, str] = None,
221
229
  byte_stream: Union[Any, bytes] = None,
222
230
  ):
@@ -282,6 +290,17 @@ class HttpResponse:
282
290
  if chunk.startswith('data: '):
283
291
  chunk = chunk[len('data: ') :]
284
292
  yield json.loads(chunk)
293
+ elif hasattr(self.response_stream, 'content'):
294
+ async for chunk in self.response_stream.content.iter_any():
295
+ # This is aiohttp.ClientResponse.
296
+ if chunk:
297
+ # In async streaming mode, the chunk of JSON is prefixed with
298
+ # "data:" which we must strip before parsing.
299
+ if not isinstance(chunk, str):
300
+ chunk = chunk.decode('utf-8')
301
+ if chunk.startswith('data: '):
302
+ chunk = chunk[len('data: ') :]
303
+ yield json.loads(chunk)
285
304
  else:
286
305
  raise ValueError('Error parsing streaming response.')
287
306
 
@@ -303,6 +322,56 @@ class HttpResponse:
303
322
  response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
304
323
 
305
324
 
325
+ # Default retry options.
326
+ # The config is based on https://cloud.google.com/storage/docs/retry-strategy.
327
+ _RETRY_ATTEMPTS = 3
328
+ _RETRY_INITIAL_DELAY = 1.0 # seconds
329
+ _RETRY_MAX_DELAY = 120.0 # seconds
330
+ _RETRY_EXP_BASE = 2
331
+ _RETRY_JITTER = 1
332
+ _RETRY_HTTP_STATUS_CODES = (
333
+ 408, # Request timeout.
334
+ 429, # Too many requests.
335
+ 500, # Internal server error.
336
+ 502, # Bad gateway.
337
+ 503, # Service unavailable.
338
+ 504, # Gateway timeout
339
+ )
340
+
341
+
342
+ def _retry_args(options: Optional[HttpRetryOptions]) -> dict[str, Any]:
343
+ """Returns the retry args for the given http retry options.
344
+
345
+ Args:
346
+ options: The http retry options to use for the retry configuration. If None,
347
+ the 'never retry' stop strategy will be used.
348
+
349
+ Returns:
350
+ The arguments passed to the tenacity.(Async)Retrying constructor.
351
+ """
352
+ if options is None:
353
+ return {'stop': tenacity.stop_after_attempt(1)}
354
+
355
+ stop = tenacity.stop_after_attempt(options.attempts or _RETRY_ATTEMPTS)
356
+ retriable_codes = options.http_status_codes or _RETRY_HTTP_STATUS_CODES
357
+ retry = tenacity.retry_if_result(
358
+ lambda response: response.status_code in retriable_codes,
359
+ )
360
+ retry_error_callback = lambda retry_state: retry_state.outcome.result()
361
+ wait = tenacity.wait_exponential_jitter(
362
+ initial=options.initial_delay or _RETRY_INITIAL_DELAY,
363
+ max=options.max_delay or _RETRY_MAX_DELAY,
364
+ exp_base=options.exp_base or _RETRY_EXP_BASE,
365
+ jitter=options.jitter or _RETRY_JITTER,
366
+ )
367
+ return {
368
+ 'stop': stop,
369
+ 'retry': retry,
370
+ 'retry_error_callback': retry_error_callback,
371
+ 'wait': wait,
372
+ }
373
+
374
+
306
375
  class SyncHttpxClient(httpx.Client):
307
376
  """Sync httpx client."""
308
377
 
@@ -384,7 +453,7 @@ class BaseApiClient:
384
453
  try:
385
454
  validated_http_options = HttpOptions.model_validate(http_options)
386
455
  except ValidationError as e:
387
- raise ValueError(f'Invalid http_options: {e}')
456
+ raise ValueError('Invalid http_options') from e
388
457
  elif isinstance(http_options, HttpOptions):
389
458
  validated_http_options = http_options
390
459
 
@@ -480,14 +549,26 @@ class BaseApiClient:
480
549
  if self._http_options.headers is not None:
481
550
  _append_library_version_headers(self._http_options.headers)
482
551
 
483
- client_args, async_client_args = self._ensure_ssl_ctx(self._http_options)
552
+ client_args, async_client_args = self._ensure_httpx_ssl_ctx(
553
+ self._http_options
554
+ )
484
555
  self._httpx_client = SyncHttpxClient(**client_args)
485
556
  self._async_httpx_client = AsyncHttpxClient(**async_client_args)
557
+ if has_aiohttp:
558
+ # Do it once at the genai.Client level. Share among all requests.
559
+ self._async_client_session_request_args = self._ensure_aiohttp_ssl_ctx(
560
+ self._http_options
561
+ )
562
+
563
+ retry_kwargs = _retry_args(self._http_options.retry_options)
564
+ self._retry = tenacity.Retrying(**retry_kwargs)
565
+ self._async_retry = tenacity.AsyncRetrying(**retry_kwargs)
486
566
 
487
567
  @staticmethod
488
- def _ensure_ssl_ctx(options: HttpOptions) -> (
489
- Tuple[dict[str, Any], dict[str, Any]]):
490
- """Ensures the SSL context is present in the client args.
568
+ def _ensure_httpx_ssl_ctx(
569
+ options: HttpOptions,
570
+ ) -> Tuple[dict[str, Any], dict[str, Any]]:
571
+ """Ensures the SSL context is present in the HTTPX client args.
491
572
 
492
573
  Creates a default SSL context if one is not provided.
493
574
 
@@ -502,8 +583,11 @@ class BaseApiClient:
502
583
  args = options.client_args
503
584
  async_args = options.async_client_args
504
585
  ctx = (
505
- args.get(verify) if args else None
506
- or async_args.get(verify) if async_args else None
586
+ args.get(verify)
587
+ if args
588
+ else None or async_args.get(verify)
589
+ if async_args
590
+ else None
507
591
  )
508
592
 
509
593
  if not ctx:
@@ -534,13 +618,76 @@ class BaseApiClient:
534
618
  if not args or not args.get(verify):
535
619
  args = (args or {}).copy()
536
620
  args[verify] = ctx
537
- return args
621
+ # Drop the args that isn't used by the httpx client.
622
+ copied_args = args.copy()
623
+ for key in copied_args.copy():
624
+ if key not in inspect.signature(httpx.Client.__init__).parameters:
625
+ del copied_args[key]
626
+ return copied_args
538
627
 
539
628
  return (
540
629
  _maybe_set(args, ctx),
541
630
  _maybe_set(async_args, ctx),
542
631
  )
543
632
 
633
+ @staticmethod
634
+ def _ensure_aiohttp_ssl_ctx(options: HttpOptions) -> dict[str, Any]:
635
+ """Ensures the SSL context is present in the async client args.
636
+
637
+ Creates a default SSL context if one is not provided.
638
+
639
+ Args:
640
+ options: The http options to check for SSL context.
641
+
642
+ Returns:
643
+ An async aiohttp ClientSession._request args.
644
+ """
645
+
646
+ verify = 'ssl' # keep it consistent with httpx.
647
+ async_args = options.async_client_args
648
+ ctx = async_args.get(verify) if async_args else None
649
+
650
+ if not ctx:
651
+ # Initialize the SSL context for the httpx client.
652
+ # Unlike requests, the aiohttp package does not automatically pull in the
653
+ # environment variables SSL_CERT_FILE or SSL_CERT_DIR. They need to be
654
+ # enabled explicitly. Instead of 'verify' at client level in httpx,
655
+ # aiohttp uses 'ssl' at request level.
656
+ ctx = ssl.create_default_context(
657
+ cafile=os.environ.get('SSL_CERT_FILE', certifi.where()),
658
+ capath=os.environ.get('SSL_CERT_DIR'),
659
+ )
660
+
661
+ def _maybe_set(
662
+ args: Optional[dict[str, Any]],
663
+ ctx: ssl.SSLContext,
664
+ ) -> dict[str, Any]:
665
+ """Sets the SSL context in the client args if not set.
666
+
667
+ Does not override the SSL context if it is already set.
668
+
669
+ Args:
670
+ args: The client args to to check for SSL context.
671
+ ctx: The SSL context to set.
672
+
673
+ Returns:
674
+ The client args with the SSL context included.
675
+ """
676
+ if not args or not args.get(verify):
677
+ args = (args or {}).copy()
678
+ args[verify] = ctx
679
+ # Drop the args that isn't in the aiohttp RequestOptions.
680
+ copied_args = args.copy()
681
+ for key in copied_args.copy():
682
+ if (
683
+ key
684
+ not in inspect.signature(aiohttp.ClientSession._request).parameters
685
+ ):
686
+ del copied_args[key]
687
+ return copied_args
688
+
689
+ return _maybe_set(async_args, ctx)
690
+
544
691
  def _websocket_base_url(self) -> str:
545
692
  url_parts = urlparse(self._http_options.base_url)
546
693
  return url_parts._replace(scheme='wss').geturl() # type: ignore[arg-type, return-value]
@@ -645,6 +792,14 @@ class BaseApiClient:
645
792
  else:
646
793
  base_url = patched_http_options.base_url
647
794
 
795
+ if (
796
+ hasattr(patched_http_options, 'extra_body')
797
+ and patched_http_options.extra_body
798
+ ):
799
+ _common.recursive_dict_update(
800
+ request_dict, patched_http_options.extra_body
801
+ )
802
+
648
803
  url = _join_url_path(
649
804
  base_url,
650
805
  versioned_path,
@@ -670,7 +825,7 @@ class BaseApiClient:
670
825
  timeout=timeout_in_seconds,
671
826
  )
672
827
 
673
- def _request(
828
+ def _request_once(
674
829
  self,
675
830
  http_request: HttpRequest,
676
831
  stream: bool = False,
@@ -716,7 +871,14 @@ class BaseApiClient:
716
871
  response.headers, response if stream else [response.text]
717
872
  )
718
873
 
719
- async def _async_request(
874
+ def _request(
875
+ self,
876
+ http_request: HttpRequest,
877
+ stream: bool = False,
878
+ ) -> HttpResponse:
879
+ return self._retry(self._request_once, http_request, stream) # type: ignore[no-any-return]
880
+
881
+ async def _async_request_once(
720
882
  self, http_request: HttpRequest, stream: bool = False
721
883
  ) -> HttpResponse:
722
884
  data: Optional[Union[str, bytes]] = None
@@ -737,33 +899,72 @@ class BaseApiClient:
737
899
  data = http_request.data
738
900
 
739
901
  if stream:
740
- httpx_request = self._async_httpx_client.build_request(
741
- method=http_request.method,
742
- url=http_request.url,
743
- content=data,
744
- headers=http_request.headers,
745
- timeout=http_request.timeout,
746
- )
747
- response = await self._async_httpx_client.send(
748
- httpx_request,
749
- stream=stream,
750
- )
751
- await errors.APIError.raise_for_async_response(response)
752
- return HttpResponse(
753
- response.headers, response if stream else [response.text]
754
- )
902
+ if has_aiohttp:
903
+ session = aiohttp.ClientSession(
904
+ headers=http_request.headers,
905
+ trust_env=True,
906
+ )
907
+ response = await session.request(
908
+ method=http_request.method,
909
+ url=http_request.url,
910
+ headers=http_request.headers,
911
+ data=data,
912
+ timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
913
+ **self._async_client_session_request_args,
914
+ )
915
+ await errors.APIError.raise_for_async_response(response)
916
+ return HttpResponse(response.headers, response)
917
+ else:
918
+ # aiohttp is not available. Fall back to httpx.
919
+ httpx_request = self._async_httpx_client.build_request(
920
+ method=http_request.method,
921
+ url=http_request.url,
922
+ content=data,
923
+ headers=http_request.headers,
924
+ timeout=http_request.timeout,
925
+ )
926
+ client_response = await self._async_httpx_client.send(
927
+ httpx_request,
928
+ stream=stream,
929
+ )
930
+ await errors.APIError.raise_for_async_response(client_response)
931
+ return HttpResponse(client_response.headers, client_response)
755
932
  else:
756
- response = await self._async_httpx_client.request(
757
- method=http_request.method,
758
- url=http_request.url,
759
- headers=http_request.headers,
760
- content=data,
761
- timeout=http_request.timeout,
762
- )
763
- await errors.APIError.raise_for_async_response(response)
764
- return HttpResponse(
765
- response.headers, response if stream else [response.text]
766
- )
933
+ if has_aiohttp:
934
+ async with aiohttp.ClientSession(
935
+ headers=http_request.headers,
936
+ trust_env=True,
937
+ ) as session:
938
+ response = await session.request(
939
+ method=http_request.method,
940
+ url=http_request.url,
941
+ headers=http_request.headers,
942
+ data=data,
943
+ timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
944
+ **self._async_client_session_request_args,
945
+ )
946
+ await errors.APIError.raise_for_async_response(response)
947
+ return HttpResponse(response.headers, [await response.text()])
948
+ else:
949
+ # aiohttp is not available. Fall back to httpx.
950
+ client_response = await self._async_httpx_client.request(
951
+ method=http_request.method,
952
+ url=http_request.url,
953
+ headers=http_request.headers,
954
+ content=data,
955
+ timeout=http_request.timeout,
956
+ )
957
+ await errors.APIError.raise_for_async_response(client_response)
958
+ return HttpResponse(client_response.headers, [client_response.text])
959
+
960
+ async def _async_request(
961
+ self,
962
+ http_request: HttpRequest,
963
+ stream: bool = False,
964
+ ) -> HttpResponse:
965
+ return await self._async_retry( # type: ignore[no-any-return]
966
+ self._async_request_once, http_request, stream
967
+ )
767
968
 
768
969
  def get_read_only_http_options(self) -> dict[str, Any]:
769
970
  if isinstance(self._http_options, BaseModel):
@@ -778,17 +979,16 @@ class BaseApiClient:
778
979
  path: str,
779
980
  request_dict: dict[str, object],
780
981
  http_options: Optional[HttpOptionsOrDict] = None,
781
- ) -> Union[BaseResponse, Any]:
982
+ ) -> SdkHttpResponse:
782
983
  http_request = self._build_request(
783
984
  http_method, path, request_dict, http_options
784
985
  )
785
986
  response = self._request(http_request, stream=False)
786
- json_response = response.json
787
- if not json_response:
788
- return BaseResponse(http_headers=response.headers).model_dump(
789
- by_alias=True
790
- )
791
- return json_response
987
+ response_body = response.response_stream[0] if response.response_stream else ''
988
+ return SdkHttpResponse(
989
+ headers=response.headers, body=response_body
990
+ )
991
+
792
992
 
793
993
  def request_streamed(
794
994
  self,
@@ -811,16 +1011,17 @@ class BaseApiClient:
811
1011
  path: str,
812
1012
  request_dict: dict[str, object],
813
1013
  http_options: Optional[HttpOptionsOrDict] = None,
814
- ) -> Union[BaseResponse, Any]:
1014
+ ) -> SdkHttpResponse:
815
1015
  http_request = self._build_request(
816
1016
  http_method, path, request_dict, http_options
817
1017
  )
818
1018
 
819
1019
  result = await self._async_request(http_request=http_request, stream=False)
820
- json_response = result.json
821
- if not json_response:
822
- return BaseResponse(http_headers=result.headers).model_dump(by_alias=True)
823
- return json_response
1020
+ response_body = result.response_stream[0] if result.response_stream else ''
1021
+ return SdkHttpResponse(
1022
+ headers=result.headers, body=response_body
1023
+ )
1024
+
824
1025
 
825
1026
  async def async_request_streamed(
826
1027
  self,
@@ -949,9 +1150,7 @@ class BaseApiClient:
949
1150
  )
950
1151
 
951
1152
  if response.headers.get('x-goog-upload-status') != 'final':
952
- raise ValueError(
953
- 'Failed to upload file: Upload status is not finalized.'
954
- )
1153
+ raise ValueError('Failed to upload file: Upload status is not finalized.')
955
1154
  return HttpResponse(response.headers, response_stream=[response.text])
956
1155
 
957
1156
  def download_file(
@@ -959,7 +1158,7 @@ class BaseApiClient:
959
1158
  path: str,
960
1159
  *,
961
1160
  http_options: Optional[HttpOptionsOrDict] = None,
962
- ) -> Union[Any,bytes]:
1161
+ ) -> Union[Any, bytes]:
963
1162
  """Downloads the file data.
964
1163
 
965
1164
  Args:
@@ -1048,68 +1247,162 @@ class BaseApiClient:
1048
1247
  """
1049
1248
  offset = 0
1050
1249
  # Upload the file in chunks
1051
- while True:
1052
- if isinstance(file, io.IOBase):
1053
- file_chunk = file.read(CHUNK_SIZE)
1054
- else:
1055
- file_chunk = await file.read(CHUNK_SIZE)
1056
- chunk_size = 0
1057
- if file_chunk:
1058
- chunk_size = len(file_chunk)
1059
- upload_command = 'upload'
1060
- # If last chunk, finalize the upload.
1061
- if chunk_size + offset >= upload_size:
1062
- upload_command += ', finalize'
1063
- http_options = http_options if http_options else self._http_options
1064
- timeout = (
1065
- http_options.get('timeout')
1066
- if isinstance(http_options, dict)
1067
- else http_options.timeout
1068
- )
1069
- if timeout is None:
1070
- # Per request timeout is not configured. Check the global timeout.
1071
- timeout = (
1072
- self._http_options.timeout
1073
- if isinstance(self._http_options, dict)
1074
- else self._http_options.timeout
1250
+ if has_aiohttp: # pylint: disable=g-import-not-at-top
1251
+ async with aiohttp.ClientSession(
1252
+ headers=self._http_options.headers,
1253
+ trust_env=True,
1254
+ ) as session:
1255
+ while True:
1256
+ if isinstance(file, io.IOBase):
1257
+ file_chunk = file.read(CHUNK_SIZE)
1258
+ else:
1259
+ file_chunk = await file.read(CHUNK_SIZE)
1260
+ chunk_size = 0
1261
+ if file_chunk:
1262
+ chunk_size = len(file_chunk)
1263
+ upload_command = 'upload'
1264
+ # If last chunk, finalize the upload.
1265
+ if chunk_size + offset >= upload_size:
1266
+ upload_command += ', finalize'
1267
+ http_options = http_options if http_options else self._http_options
1268
+ timeout = (
1269
+ http_options.get('timeout')
1270
+ if isinstance(http_options, dict)
1271
+ else http_options.timeout
1272
+ )
1273
+ if timeout is None:
1274
+ # Per request timeout is not configured. Check the global timeout.
1275
+ timeout = (
1276
+ self._http_options.timeout
1277
+ if isinstance(self._http_options, dict)
1278
+ else self._http_options.timeout
1279
+ )
1280
+ timeout_in_seconds = _get_timeout_in_seconds(timeout)
1281
+ upload_headers = {
1282
+ 'X-Goog-Upload-Command': upload_command,
1283
+ 'X-Goog-Upload-Offset': str(offset),
1284
+ 'Content-Length': str(chunk_size),
1285
+ }
1286
+ _populate_server_timeout_header(upload_headers, timeout_in_seconds)
1287
+
1288
+ retry_count = 0
1289
+ response = None
1290
+ while retry_count < MAX_RETRY_COUNT:
1291
+ response = await session.request(
1292
+ method='POST',
1293
+ url=upload_url,
1294
+ data=file_chunk,
1295
+ headers=upload_headers,
1296
+ timeout=aiohttp.ClientTimeout(connect=timeout_in_seconds),
1297
+ )
1298
+
1299
+ if response.headers.get('X-Goog-Upload-Status'):
1300
+ break
1301
+ delay_seconds = INITIAL_RETRY_DELAY * (
1302
+ DELAY_MULTIPLIER**retry_count
1303
+ )
1304
+ retry_count += 1
1305
+ time.sleep(delay_seconds)
1306
+
1307
+ offset += chunk_size
1308
+ if (
1309
+ response is not None
1310
+ and response.headers.get('X-Goog-Upload-Status') != 'active'
1311
+ ):
1312
+ break # upload is complete or it has been interrupted.
1313
+
1314
+ if upload_size <= offset: # Status is not finalized.
1315
+ raise ValueError(
1316
+ f'All content has been uploaded, but the upload status is not'
1317
+ f' finalized.'
1318
+ )
1319
+ if (
1320
+ response is not None
1321
+ and response.headers.get('X-Goog-Upload-Status') != 'final'
1322
+ ):
1323
+ raise ValueError(
1324
+ 'Failed to upload file: Upload status is not finalized.'
1325
+ )
1326
+ return HttpResponse(
1327
+ response.headers, response_stream=[await response.text()]
1075
1328
  )
1076
- timeout_in_seconds = _get_timeout_in_seconds(timeout)
1077
- upload_headers = {
1078
- 'X-Goog-Upload-Command': upload_command,
1079
- 'X-Goog-Upload-Offset': str(offset),
1080
- 'Content-Length': str(chunk_size),
1081
- }
1082
- _populate_server_timeout_header(upload_headers, timeout_in_seconds)
1083
-
1084
- retry_count = 0
1085
- while retry_count < MAX_RETRY_COUNT:
1086
- response = await self._async_httpx_client.request(
1087
- method='POST',
1088
- url=upload_url,
1089
- content=file_chunk,
1090
- headers=upload_headers,
1091
- timeout=timeout_in_seconds,
1329
+ else:
1330
+ # aiohttp is not available. Fall back to httpx.
1331
+ while True:
1332
+ if isinstance(file, io.IOBase):
1333
+ file_chunk = file.read(CHUNK_SIZE)
1334
+ else:
1335
+ file_chunk = await file.read(CHUNK_SIZE)
1336
+ chunk_size = 0
1337
+ if file_chunk:
1338
+ chunk_size = len(file_chunk)
1339
+ upload_command = 'upload'
1340
+ # If last chunk, finalize the upload.
1341
+ if chunk_size + offset >= upload_size:
1342
+ upload_command += ', finalize'
1343
+ http_options = http_options if http_options else self._http_options
1344
+ timeout = (
1345
+ http_options.get('timeout')
1346
+ if isinstance(http_options, dict)
1347
+ else http_options.timeout
1092
1348
  )
1093
- if response.headers.get('x-goog-upload-status'):
1094
- break
1095
- delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
1096
- retry_count += 1
1097
- time.sleep(delay_seconds)
1098
-
1099
- offset += chunk_size
1100
- if response.headers.get('x-goog-upload-status') != 'active':
1101
- break # upload is complete or it has been interrupted.
1102
-
1103
- if upload_size <= offset: # Status is not finalized.
1349
+ if timeout is None:
1350
+ # Per request timeout is not configured. Check the global timeout.
1351
+ timeout = (
1352
+ self._http_options.timeout
1353
+ if isinstance(self._http_options, dict)
1354
+ else self._http_options.timeout
1355
+ )
1356
+ timeout_in_seconds = _get_timeout_in_seconds(timeout)
1357
+ upload_headers = {
1358
+ 'X-Goog-Upload-Command': upload_command,
1359
+ 'X-Goog-Upload-Offset': str(offset),
1360
+ 'Content-Length': str(chunk_size),
1361
+ }
1362
+ _populate_server_timeout_header(upload_headers, timeout_in_seconds)
1363
+
1364
+ retry_count = 0
1365
+ client_response = None
1366
+ while retry_count < MAX_RETRY_COUNT:
1367
+ client_response = await self._async_httpx_client.request(
1368
+ method='POST',
1369
+ url=upload_url,
1370
+ content=file_chunk,
1371
+ headers=upload_headers,
1372
+ timeout=timeout_in_seconds,
1373
+ )
1374
+ if (
1375
+ client_response is not None
1376
+ and client_response.headers
1377
+ and client_response.headers.get('x-goog-upload-status')
1378
+ ):
1379
+ break
1380
+ delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
1381
+ retry_count += 1
1382
+ time.sleep(delay_seconds)
1383
+
1384
+ offset += chunk_size
1385
+ if (
1386
+ client_response is not None
1387
+ and client_response.headers.get('x-goog-upload-status') != 'active'
1388
+ ):
1389
+ break # upload is complete or it has been interrupted.
1390
+
1391
+ if upload_size <= offset: # Status is not finalized.
1392
+ raise ValueError(
1393
+ 'All content has been uploaded, but the upload status is not'
1394
+ ' finalized.'
1395
+ )
1396
+ if (
1397
+ client_response is not None
1398
+ and client_response.headers.get('x-goog-upload-status') != 'final'
1399
+ ):
1104
1400
  raise ValueError(
1105
- 'All content has been uploaded, but the upload status is not'
1106
- f' finalized.'
1401
+ 'Failed to upload file: Upload status is not finalized.'
1107
1402
  )
1108
- if response.headers.get('x-goog-upload-status') != 'final':
1109
- raise ValueError(
1110
- 'Failed to upload file: Upload status is not finalized.'
1403
+ return HttpResponse(
1404
+ client_response.headers, response_stream=[client_response.text]
1111
1405
  )
1112
- return HttpResponse(response.headers, response_stream=[response.text])
1113
1406
 
1114
1407
  async def async_download_file(
1115
1408
  self,
@@ -1137,18 +1430,37 @@ class BaseApiClient:
1137
1430
  else:
1138
1431
  data = http_request.data
1139
1432
 
1140
- response = await self._async_httpx_client.request(
1141
- method=http_request.method,
1142
- url=http_request.url,
1143
- headers=http_request.headers,
1144
- content=data,
1145
- timeout=http_request.timeout,
1146
- )
1147
- await errors.APIError.raise_for_async_response(response)
1433
+ if has_aiohttp:
1434
+ async with aiohttp.ClientSession(
1435
+ headers=http_request.headers,
1436
+ trust_env=True,
1437
+ ) as session:
1438
+ response = await session.request(
1439
+ method=http_request.method,
1440
+ url=http_request.url,
1441
+ headers=http_request.headers,
1442
+ data=data,
1443
+ timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
1444
+ )
1445
+ await errors.APIError.raise_for_async_response(response)
1148
1446
 
1149
- return HttpResponse(
1150
- response.headers, byte_stream=[response.read()]
1151
- ).byte_stream[0]
1447
+ return HttpResponse(
1448
+ response.headers, byte_stream=[await response.read()]
1449
+ ).byte_stream[0]
1450
+ else:
1451
+ # aiohttp is not available. Fall back to httpx.
1452
+ client_response = await self._async_httpx_client.request(
1453
+ method=http_request.method,
1454
+ url=http_request.url,
1455
+ headers=http_request.headers,
1456
+ content=data,
1457
+ timeout=http_request.timeout,
1458
+ )
1459
+ await errors.APIError.raise_for_async_response(client_response)
1460
+
1461
+ return HttpResponse(
1462
+ client_response.headers, byte_stream=[client_response.read()]
1463
+ ).byte_stream[0]
1152
1464
 
1153
1465
  # This method does nothing in the real api client. It is used in the
1154
1466
  # replay_api_client to verify the response from the SDK method matches the