google-genai 1.22.0__tar.gz → 1.24.0__tar.gz

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.
Files changed (41) hide show
  1. {google_genai-1.22.0/google_genai.egg-info → google_genai-1.24.0}/PKG-INFO +120 -6
  2. {google_genai-1.22.0 → google_genai-1.24.0}/README.md +119 -5
  3. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_api_client.py +111 -32
  4. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_live_converters.py +14 -6
  5. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_tokens_converters.py +6 -0
  6. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/batches.py +84 -12
  7. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/caches.py +6 -0
  8. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/live.py +4 -1
  9. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/models.py +6 -0
  10. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/tunings.py +36 -0
  11. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/types.py +311 -36
  12. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/version.py +1 -1
  13. {google_genai-1.22.0 → google_genai-1.24.0/google_genai.egg-info}/PKG-INFO +120 -6
  14. {google_genai-1.22.0 → google_genai-1.24.0}/pyproject.toml +1 -1
  15. {google_genai-1.22.0 → google_genai-1.24.0}/LICENSE +0 -0
  16. {google_genai-1.22.0 → google_genai-1.24.0}/MANIFEST.in +0 -0
  17. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/__init__.py +0 -0
  18. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_adapters.py +0 -0
  19. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_api_module.py +0 -0
  20. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_automatic_function_calling_util.py +0 -0
  21. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_base_url.py +0 -0
  22. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_common.py +0 -0
  23. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_extra_utils.py +0 -0
  24. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_mcp_utils.py +0 -0
  25. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_replay_api_client.py +0 -0
  26. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_test_api_client.py +0 -0
  27. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/_transformers.py +0 -0
  28. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/chats.py +0 -0
  29. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/client.py +0 -0
  30. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/errors.py +0 -0
  31. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/files.py +0 -0
  32. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/live_music.py +0 -0
  33. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/operations.py +0 -0
  34. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/pagers.py +0 -0
  35. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/py.typed +0 -0
  36. {google_genai-1.22.0 → google_genai-1.24.0}/google/genai/tokens.py +0 -0
  37. {google_genai-1.22.0 → google_genai-1.24.0}/google_genai.egg-info/SOURCES.txt +0 -0
  38. {google_genai-1.22.0 → google_genai-1.24.0}/google_genai.egg-info/dependency_links.txt +0 -0
  39. {google_genai-1.22.0 → google_genai-1.24.0}/google_genai.egg-info/requires.txt +0 -0
  40. {google_genai-1.22.0 → google_genai-1.24.0}/google_genai.egg-info/top_level.txt +0 -0
  41. {google_genai-1.22.0 → google_genai-1.24.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: google-genai
3
- Version: 1.22.0
3
+ Version: 1.24.0
4
4
  Summary: GenAI Python SDK
5
5
  Author-email: Google LLC <googleapis-packages@google.com>
6
6
  License: Apache-2.0
@@ -743,6 +743,53 @@ response = client.models.generate_content(
743
743
  ),
744
744
  )
745
745
  ```
746
+
747
+ #### Model Context Protocol (MCP) support (experimental)
748
+
749
+ Built-in [MCP](https://modelcontextprotocol.io/introduction) support is an
750
+ experimental feature. You can pass a local MCP server as a tool directly.
751
+
752
+ ```python
753
+ import os
754
+ import asyncio
755
+ from datetime import datetime
756
+ from mcp import ClientSession, StdioServerParameters
757
+ from mcp.client.stdio import stdio_client
758
+ from google import genai
759
+
760
+ client = genai.Client()
761
+
762
+ # Create server parameters for stdio connection
763
+ server_params = StdioServerParameters(
764
+ command="npx", # Executable
765
+ args=["-y", "@philschmid/weather-mcp"], # MCP Server
766
+ env=None, # Optional environment variables
767
+ )
768
+
769
+ async def run():
770
+ async with stdio_client(server_params) as (read, write):
771
+ async with ClientSession(read, write) as session:
772
+ # Prompt to get the weather for the current day in London.
773
+ prompt = f"What is the weather in London in {datetime.now().strftime('%Y-%m-%d')}?"
774
+
775
+ # Initialize the connection between client and server
776
+ await session.initialize()
777
+
778
+ # Send request to the model with MCP function declarations
779
+ response = await client.aio.models.generate_content(
780
+ model="gemini-2.5-flash",
781
+ contents=prompt,
782
+ config=genai.types.GenerateContentConfig(
783
+ temperature=0,
784
+ tools=[session], # uses the session, will automatically call the tool using automatic function calling
785
+ ),
786
+ )
787
+ print(response.text)
788
+
789
+ # Start the asyncio event loop and run the main function
790
+ asyncio.run(run())
791
+ ```
792
+
746
793
  ### JSON Response Schema
747
794
 
748
795
  However you define your schema, don't duplicate it in your input prompt,
@@ -1085,9 +1132,9 @@ response3.generated_images[0].image.show()
1085
1132
 
1086
1133
  ### Veo
1087
1134
 
1088
- #### Generate Videos
1135
+ Support for generating videos is considered public preview
1089
1136
 
1090
- Support for generate videos in Vertex and Gemini Developer API is behind an allowlist
1137
+ #### Generate Videos (Text to Video)
1091
1138
 
1092
1139
  ```python
1093
1140
  from google.genai import types
@@ -1098,7 +1145,6 @@ operation = client.models.generate_videos(
1098
1145
  prompt='A neon hologram of a cat driving at top speed',
1099
1146
  config=types.GenerateVideosConfig(
1100
1147
  number_of_videos=1,
1101
- fps=24,
1102
1148
  duration_seconds=5,
1103
1149
  enhance_prompt=True,
1104
1150
  ),
@@ -1109,7 +1155,73 @@ while not operation.done:
1109
1155
  time.sleep(20)
1110
1156
  operation = client.operations.get(operation)
1111
1157
 
1112
- video = operation.result.generated_videos[0].video
1158
+ video = operation.response.generated_videos[0].video
1159
+ video.show()
1160
+ ```
1161
+
1162
+ #### Generate Videos (Image to Video)
1163
+
1164
+ ```python
1165
+ from google.genai import types
1166
+
1167
+ # Read local image (uses mimetypes.guess_type to infer mime type)
1168
+ image = types.Image.from_file("local/path/file.png")
1169
+
1170
+ # Create operation
1171
+ operation = client.models.generate_videos(
1172
+ model='veo-2.0-generate-001',
1173
+ # Prompt is optional if image is provided
1174
+ prompt='Night sky',
1175
+ image=image,
1176
+ config=types.GenerateVideosConfig(
1177
+ number_of_videos=1,
1178
+ duration_seconds=5,
1179
+ enhance_prompt=True,
1180
+ # Can also pass an Image into last_frame for frame interpolation
1181
+ ),
1182
+ )
1183
+
1184
+ # Poll operation
1185
+ while not operation.done:
1186
+ time.sleep(20)
1187
+ operation = client.operations.get(operation)
1188
+
1189
+ video = operation.response.generated_videos[0].video
1190
+ video.show()
1191
+ ```
1192
+
1193
+ #### Generate Videos (Video to Video)
1194
+
1195
+ Currently, only Vertex supports Video to Video generation (Video extension).
1196
+
1197
+ ```python
1198
+ from google.genai import types
1199
+
1200
+ # Read local video (uses mimetypes.guess_type to infer mime type)
1201
+ video = types.Video.from_file("local/path/video.mp4")
1202
+
1203
+ # Create operation
1204
+ operation = client.models.generate_videos(
1205
+ model='veo-2.0-generate-001',
1206
+ # Prompt is optional if Video is provided
1207
+ prompt='Night sky',
1208
+ # Input video must be in GCS
1209
+ video=types.Video(
1210
+ uri="gs://bucket-name/inputs/videos/cat_driving.mp4",
1211
+ ),
1212
+ config=types.GenerateVideosConfig(
1213
+ number_of_videos=1,
1214
+ duration_seconds=5,
1215
+ enhance_prompt=True,
1216
+ ),
1217
+ )
1218
+
1219
+ # Poll operation
1220
+ while not operation.done:
1221
+ time.sleep(20)
1222
+ operation = client.operations.get(operation)
1223
+
1224
+ video = operation.response.generated_videos[0].video
1113
1225
  video.show()
1114
1226
  ```
1115
1227
 
@@ -1260,7 +1372,7 @@ client.
1260
1372
 
1261
1373
  ### Tune
1262
1374
 
1263
- - Vertex AI supports tuning from GCS source
1375
+ - Vertex AI supports tuning from GCS source or from a Vertex Multimodal Dataset
1264
1376
  - Gemini Developer API supports tuning from inline examples
1265
1377
 
1266
1378
  ```python
@@ -1269,10 +1381,12 @@ from google.genai import types
1269
1381
  if client.vertexai:
1270
1382
  model = 'gemini-2.0-flash-001'
1271
1383
  training_dataset = types.TuningDataset(
1384
+ # or gcs_uri=my_vertex_multimodal_dataset
1272
1385
  gcs_uri='gs://cloud-samples-data/ai-platform/generative_ai/gemini-1_5/text/sft_train_data.jsonl',
1273
1386
  )
1274
1387
  else:
1275
1388
  model = 'models/gemini-2.0-flash-001'
1389
+ # or gcs_uri=my_vertex_multimodal_dataset.resource_name
1276
1390
  training_dataset = types.TuningDataset(
1277
1391
  examples=[
1278
1392
  types.TuningExample(
@@ -709,6 +709,53 @@ response = client.models.generate_content(
709
709
  ),
710
710
  )
711
711
  ```
712
+
713
+ #### Model Context Protocol (MCP) support (experimental)
714
+
715
+ Built-in [MCP](https://modelcontextprotocol.io/introduction) support is an
716
+ experimental feature. You can pass a local MCP server as a tool directly.
717
+
718
+ ```python
719
+ import os
720
+ import asyncio
721
+ from datetime import datetime
722
+ from mcp import ClientSession, StdioServerParameters
723
+ from mcp.client.stdio import stdio_client
724
+ from google import genai
725
+
726
+ client = genai.Client()
727
+
728
+ # Create server parameters for stdio connection
729
+ server_params = StdioServerParameters(
730
+ command="npx", # Executable
731
+ args=["-y", "@philschmid/weather-mcp"], # MCP Server
732
+ env=None, # Optional environment variables
733
+ )
734
+
735
+ async def run():
736
+ async with stdio_client(server_params) as (read, write):
737
+ async with ClientSession(read, write) as session:
738
+ # Prompt to get the weather for the current day in London.
739
+ prompt = f"What is the weather in London in {datetime.now().strftime('%Y-%m-%d')}?"
740
+
741
+ # Initialize the connection between client and server
742
+ await session.initialize()
743
+
744
+ # Send request to the model with MCP function declarations
745
+ response = await client.aio.models.generate_content(
746
+ model="gemini-2.5-flash",
747
+ contents=prompt,
748
+ config=genai.types.GenerateContentConfig(
749
+ temperature=0,
750
+ tools=[session], # uses the session, will automatically call the tool using automatic function calling
751
+ ),
752
+ )
753
+ print(response.text)
754
+
755
+ # Start the asyncio event loop and run the main function
756
+ asyncio.run(run())
757
+ ```
758
+
712
759
  ### JSON Response Schema
713
760
 
714
761
  However you define your schema, don't duplicate it in your input prompt,
@@ -1051,9 +1098,9 @@ response3.generated_images[0].image.show()
1051
1098
 
1052
1099
  ### Veo
1053
1100
 
1054
- #### Generate Videos
1101
+ Support for generating videos is considered public preview
1055
1102
 
1056
- Support for generate videos in Vertex and Gemini Developer API is behind an allowlist
1103
+ #### Generate Videos (Text to Video)
1057
1104
 
1058
1105
  ```python
1059
1106
  from google.genai import types
@@ -1064,7 +1111,6 @@ operation = client.models.generate_videos(
1064
1111
  prompt='A neon hologram of a cat driving at top speed',
1065
1112
  config=types.GenerateVideosConfig(
1066
1113
  number_of_videos=1,
1067
- fps=24,
1068
1114
  duration_seconds=5,
1069
1115
  enhance_prompt=True,
1070
1116
  ),
@@ -1075,7 +1121,73 @@ while not operation.done:
1075
1121
  time.sleep(20)
1076
1122
  operation = client.operations.get(operation)
1077
1123
 
1078
- video = operation.result.generated_videos[0].video
1124
+ video = operation.response.generated_videos[0].video
1125
+ video.show()
1126
+ ```
1127
+
1128
+ #### Generate Videos (Image to Video)
1129
+
1130
+ ```python
1131
+ from google.genai import types
1132
+
1133
+ # Read local image (uses mimetypes.guess_type to infer mime type)
1134
+ image = types.Image.from_file("local/path/file.png")
1135
+
1136
+ # Create operation
1137
+ operation = client.models.generate_videos(
1138
+ model='veo-2.0-generate-001',
1139
+ # Prompt is optional if image is provided
1140
+ prompt='Night sky',
1141
+ image=image,
1142
+ config=types.GenerateVideosConfig(
1143
+ number_of_videos=1,
1144
+ duration_seconds=5,
1145
+ enhance_prompt=True,
1146
+ # Can also pass an Image into last_frame for frame interpolation
1147
+ ),
1148
+ )
1149
+
1150
+ # Poll operation
1151
+ while not operation.done:
1152
+ time.sleep(20)
1153
+ operation = client.operations.get(operation)
1154
+
1155
+ video = operation.response.generated_videos[0].video
1156
+ video.show()
1157
+ ```
1158
+
1159
+ #### Generate Videos (Video to Video)
1160
+
1161
+ Currently, only Vertex supports Video to Video generation (Video extension).
1162
+
1163
+ ```python
1164
+ from google.genai import types
1165
+
1166
+ # Read local video (uses mimetypes.guess_type to infer mime type)
1167
+ video = types.Video.from_file("local/path/video.mp4")
1168
+
1169
+ # Create operation
1170
+ operation = client.models.generate_videos(
1171
+ model='veo-2.0-generate-001',
1172
+ # Prompt is optional if Video is provided
1173
+ prompt='Night sky',
1174
+ # Input video must be in GCS
1175
+ video=types.Video(
1176
+ uri="gs://bucket-name/inputs/videos/cat_driving.mp4",
1177
+ ),
1178
+ config=types.GenerateVideosConfig(
1179
+ number_of_videos=1,
1180
+ duration_seconds=5,
1181
+ enhance_prompt=True,
1182
+ ),
1183
+ )
1184
+
1185
+ # Poll operation
1186
+ while not operation.done:
1187
+ time.sleep(20)
1188
+ operation = client.operations.get(operation)
1189
+
1190
+ video = operation.response.generated_videos[0].video
1079
1191
  video.show()
1080
1192
  ```
1081
1193
 
@@ -1226,7 +1338,7 @@ client.
1226
1338
 
1227
1339
  ### Tune
1228
1340
 
1229
- - Vertex AI supports tuning from GCS source
1341
+ - Vertex AI supports tuning from GCS source or from a Vertex Multimodal Dataset
1230
1342
  - Gemini Developer API supports tuning from inline examples
1231
1343
 
1232
1344
  ```python
@@ -1235,10 +1347,12 @@ from google.genai import types
1235
1347
  if client.vertexai:
1236
1348
  model = 'gemini-2.0-flash-001'
1237
1349
  training_dataset = types.TuningDataset(
1350
+ # or gcs_uri=my_vertex_multimodal_dataset
1238
1351
  gcs_uri='gs://cloud-samples-data/ai-platform/generative_ai/gemini-1_5/text/sft_train_data.jsonl',
1239
1352
  )
1240
1353
  else:
1241
1354
  model = 'models/gemini-2.0-flash-001'
1355
+ # or gcs_uri=my_vertex_multimodal_dataset.resource_name
1242
1356
  training_dataset = types.TuningDataset(
1243
1357
  examples=[
1244
1358
  types.TuningExample(
@@ -60,6 +60,11 @@ from .types import HttpOptionsOrDict
60
60
  from .types import HttpResponse as SdkHttpResponse
61
61
  from .types import HttpRetryOptions
62
62
 
63
+ try:
64
+ from websockets.asyncio.client import connect as ws_connect
65
+ except ModuleNotFoundError:
66
+ # This try/except is for TAP, mypy complains about it which is why we have the type: ignore
67
+ from websockets.client import connect as ws_connect # type: ignore
63
68
 
64
69
  has_aiohttp = False
65
70
  try:
@@ -69,8 +74,6 @@ try:
69
74
  except ImportError:
70
75
  pass
71
76
 
72
- # internal comment
73
-
74
77
 
75
78
  if TYPE_CHECKING:
76
79
  from multidict import CIMultiDictProxy
@@ -227,11 +230,13 @@ class HttpResponse:
227
230
  headers: Union[dict[str, str], httpx.Headers, 'CIMultiDictProxy[str]'],
228
231
  response_stream: Union[Any, str] = None,
229
232
  byte_stream: Union[Any, bytes] = None,
233
+ session: Optional['aiohttp.ClientSession'] = None,
230
234
  ):
231
235
  self.status_code: int = 200
232
236
  self.headers = headers
233
237
  self.response_stream = response_stream
234
238
  self.byte_stream = byte_stream
239
+ self._session = session
235
240
 
236
241
  # Async iterator for async streaming.
237
242
  def __aiter__(self) -> 'HttpResponse':
@@ -291,16 +296,23 @@ class HttpResponse:
291
296
  chunk = chunk[len('data: ') :]
292
297
  yield json.loads(chunk)
293
298
  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:
299
+ # This is aiohttp.ClientResponse.
300
+ try:
301
+ while True:
302
+ chunk = await self.response_stream.content.readline()
303
+ if not chunk:
304
+ break
297
305
  # In async streaming mode, the chunk of JSON is prefixed with
298
306
  # "data:" which we must strip before parsing.
299
- if not isinstance(chunk, str):
300
- chunk = chunk.decode('utf-8')
307
+ chunk = chunk.decode('utf-8')
301
308
  if chunk.startswith('data: '):
302
309
  chunk = chunk[len('data: ') :]
303
- yield json.loads(chunk)
310
+ chunk = chunk.strip()
311
+ if chunk:
312
+ yield json.loads(chunk)
313
+ finally:
314
+ if hasattr(self, '_session') and self._session:
315
+ await self._session.close()
304
316
  else:
305
317
  raise ValueError('Error parsing streaming response.')
306
318
 
@@ -324,9 +336,11 @@ class HttpResponse:
324
336
 
325
337
  # Default retry options.
326
338
  # The config is based on https://cloud.google.com/storage/docs/retry-strategy.
327
- _RETRY_ATTEMPTS = 3
339
+ # By default, the client will retry 4 times with approximately 1.0, 2.0, 4.0,
340
+ # 8.0 seconds between each attempt.
341
+ _RETRY_ATTEMPTS = 5 # including the initial call.
328
342
  _RETRY_INITIAL_DELAY = 1.0 # seconds
329
- _RETRY_MAX_DELAY = 120.0 # seconds
343
+ _RETRY_MAX_DELAY = 60.0 # seconds
330
344
  _RETRY_EXP_BASE = 2
331
345
  _RETRY_JITTER = 1
332
346
  _RETRY_HTTP_STATUS_CODES = (
@@ -350,14 +364,13 @@ def _retry_args(options: Optional[HttpRetryOptions]) -> dict[str, Any]:
350
364
  The arguments passed to the tenacity.(Async)Retrying constructor.
351
365
  """
352
366
  if options is None:
353
- return {'stop': tenacity.stop_after_attempt(1)}
367
+ return {'stop': tenacity.stop_after_attempt(1), 'reraise': True}
354
368
 
355
369
  stop = tenacity.stop_after_attempt(options.attempts or _RETRY_ATTEMPTS)
356
370
  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,
371
+ retry = tenacity.retry_if_exception(
372
+ lambda e: isinstance(e, errors.APIError) and e.code in retriable_codes,
359
373
  )
360
- retry_error_callback = lambda retry_state: retry_state.outcome.result()
361
374
  wait = tenacity.wait_exponential_jitter(
362
375
  initial=options.initial_delay or _RETRY_INITIAL_DELAY,
363
376
  max=options.max_delay or _RETRY_MAX_DELAY,
@@ -367,7 +380,7 @@ def _retry_args(options: Optional[HttpRetryOptions]) -> dict[str, Any]:
367
380
  return {
368
381
  'stop': stop,
369
382
  'retry': retry,
370
- 'retry_error_callback': retry_error_callback,
383
+ 'reraise': True,
371
384
  'wait': wait,
372
385
  }
373
386
 
@@ -538,6 +551,7 @@ class BaseApiClient:
538
551
  # Default options for both clients.
539
552
  self._http_options.headers = {'Content-Type': 'application/json'}
540
553
  if self.api_key:
554
+ self.api_key = self.api_key.strip()
541
555
  if self._http_options.headers is not None:
542
556
  self._http_options.headers['x-goog-api-key'] = self.api_key
543
557
  # Update the http options with the user provided http options.
@@ -554,15 +568,16 @@ class BaseApiClient:
554
568
  )
555
569
  self._httpx_client = SyncHttpxClient(**client_args)
556
570
  self._async_httpx_client = AsyncHttpxClient(**async_client_args)
557
- if has_aiohttp:
571
+ if self._use_aiohttp():
558
572
  # Do it once at the genai.Client level. Share among all requests.
559
573
  self._async_client_session_request_args = self._ensure_aiohttp_ssl_ctx(
560
574
  self._http_options
561
575
  )
576
+ self._websocket_ssl_ctx = self._ensure_websocket_ssl_ctx(self._http_options)
562
577
 
563
578
  retry_kwargs = _retry_args(self._http_options.retry_options)
564
- self._retry = tenacity.Retrying(**retry_kwargs, reraise=True)
565
- self._async_retry = tenacity.AsyncRetrying(**retry_kwargs, reraise=True)
579
+ self._retry = tenacity.Retrying(**retry_kwargs)
580
+ self._async_retry = tenacity.AsyncRetrying(**retry_kwargs)
566
581
 
567
582
  @staticmethod
568
583
  def _ensure_httpx_ssl_ctx(
@@ -688,6 +703,70 @@ class BaseApiClient:
688
703
 
689
704
  return _maybe_set(async_args, ctx)
690
705
 
706
+ @staticmethod
707
+ def _ensure_websocket_ssl_ctx(options: HttpOptions) -> dict[str, Any]:
708
+ """Ensures the SSL context is present in the async client args.
709
+
710
+ Creates a default SSL context if one is not provided.
711
+
712
+ Args:
713
+ options: The http options to check for SSL context.
714
+
715
+ Returns:
716
+ An async aiohttp ClientSession._request args.
717
+ """
718
+
719
+ verify = 'ssl' # keep it consistent with httpx.
720
+ async_args = options.async_client_args
721
+ ctx = async_args.get(verify) if async_args else None
722
+
723
+ if not ctx:
724
+ # Initialize the SSL context for the httpx client.
725
+ # Unlike requests, the aiohttp package does not automatically pull in the
726
+ # environment variables SSL_CERT_FILE or SSL_CERT_DIR. They need to be
727
+ # enabled explicitly. Instead of 'verify' at client level in httpx,
728
+ # aiohttp uses 'ssl' at request level.
729
+ ctx = ssl.create_default_context(
730
+ cafile=os.environ.get('SSL_CERT_FILE', certifi.where()),
731
+ capath=os.environ.get('SSL_CERT_DIR'),
732
+ )
733
+
734
+ def _maybe_set(
735
+ args: Optional[dict[str, Any]],
736
+ ctx: ssl.SSLContext,
737
+ ) -> dict[str, Any]:
738
+ """Sets the SSL context in the client args if not set.
739
+
740
+ Does not override the SSL context if it is already set.
741
+
742
+ Args:
743
+ args: The client args to to check for SSL context.
744
+ ctx: The SSL context to set.
745
+
746
+ Returns:
747
+ The client args with the SSL context included.
748
+ """
749
+ if not args or not args.get(verify):
750
+ args = (args or {}).copy()
751
+ args[verify] = ctx
752
+ # Drop the args that isn't in the aiohttp RequestOptions.
753
+ copied_args = args.copy()
754
+ for key in copied_args.copy():
755
+ if key not in inspect.signature(ws_connect).parameters and key != 'ssl':
756
+ del copied_args[key]
757
+ return copied_args
758
+
759
+ return _maybe_set(async_args, ctx)
760
+
761
+ def _use_aiohttp(self) -> bool:
762
+ # If the instantiator has passed a custom transport, they want httpx not
763
+ # aiohttp.
764
+ return (
765
+ has_aiohttp
766
+ and (self._http_options.async_client_args or {}).get('transport')
767
+ is None
768
+ )
769
+
691
770
  def _websocket_base_url(self) -> str:
692
771
  url_parts = urlparse(self._http_options.base_url)
693
772
  return url_parts._replace(scheme='wss').geturl() # type: ignore[arg-type, return-value]
@@ -882,6 +961,7 @@ class BaseApiClient:
882
961
  self, http_request: HttpRequest, stream: bool = False
883
962
  ) -> HttpResponse:
884
963
  data: Optional[Union[str, bytes]] = None
964
+
885
965
  if self.vertexai and not self.api_key:
886
966
  http_request.headers['Authorization'] = (
887
967
  f'Bearer {await self._async_access_token()}'
@@ -899,7 +979,7 @@ class BaseApiClient:
899
979
  data = http_request.data
900
980
 
901
981
  if stream:
902
- if has_aiohttp:
982
+ if self._use_aiohttp():
903
983
  session = aiohttp.ClientSession(
904
984
  headers=http_request.headers,
905
985
  trust_env=True,
@@ -912,8 +992,9 @@ class BaseApiClient:
912
992
  timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
913
993
  **self._async_client_session_request_args,
914
994
  )
995
+
915
996
  await errors.APIError.raise_for_async_response(response)
916
- return HttpResponse(response.headers, response)
997
+ return HttpResponse(response.headers, response, session=session)
917
998
  else:
918
999
  # aiohttp is not available. Fall back to httpx.
919
1000
  httpx_request = self._async_httpx_client.build_request(
@@ -930,7 +1011,7 @@ class BaseApiClient:
930
1011
  await errors.APIError.raise_for_async_response(client_response)
931
1012
  return HttpResponse(client_response.headers, client_response)
932
1013
  else:
933
- if has_aiohttp:
1014
+ if self._use_aiohttp():
934
1015
  async with aiohttp.ClientSession(
935
1016
  headers=http_request.headers,
936
1017
  trust_env=True,
@@ -984,11 +1065,10 @@ class BaseApiClient:
984
1065
  http_method, path, request_dict, http_options
985
1066
  )
986
1067
  response = self._request(http_request, stream=False)
987
- response_body = response.response_stream[0] if response.response_stream else ''
988
- return SdkHttpResponse(
989
- headers=response.headers, body=response_body
1068
+ response_body = (
1069
+ response.response_stream[0] if response.response_stream else ''
990
1070
  )
991
-
1071
+ return SdkHttpResponse(headers=response.headers, body=response_body)
992
1072
 
993
1073
  def request_streamed(
994
1074
  self,
@@ -1003,7 +1083,9 @@ class BaseApiClient:
1003
1083
 
1004
1084
  session_response = self._request(http_request, stream=True)
1005
1085
  for chunk in session_response.segments():
1006
- yield SdkHttpResponse(headers=session_response.headers, body=json.dumps(chunk))
1086
+ yield SdkHttpResponse(
1087
+ headers=session_response.headers, body=json.dumps(chunk)
1088
+ )
1007
1089
 
1008
1090
  async def async_request(
1009
1091
  self,
@@ -1018,10 +1100,7 @@ class BaseApiClient:
1018
1100
 
1019
1101
  result = await self._async_request(http_request=http_request, stream=False)
1020
1102
  response_body = result.response_stream[0] if result.response_stream else ''
1021
- return SdkHttpResponse(
1022
- headers=result.headers, body=response_body
1023
- )
1024
-
1103
+ return SdkHttpResponse(headers=result.headers, body=response_body)
1025
1104
 
1026
1105
  async def async_request_streamed(
1027
1106
  self,
@@ -1247,7 +1326,7 @@ class BaseApiClient:
1247
1326
  """
1248
1327
  offset = 0
1249
1328
  # Upload the file in chunks
1250
- if has_aiohttp: # pylint: disable=g-import-not-at-top
1329
+ if self._use_aiohttp(): # pylint: disable=g-import-not-at-top
1251
1330
  async with aiohttp.ClientSession(
1252
1331
  headers=self._http_options.headers,
1253
1332
  trust_env=True,
@@ -1430,7 +1509,7 @@ class BaseApiClient:
1430
1509
  else:
1431
1510
  data = http_request.data
1432
1511
 
1433
- if has_aiohttp:
1512
+ if self._use_aiohttp():
1434
1513
  async with aiohttp.ClientSession(
1435
1514
  headers=http_request.headers,
1436
1515
  trust_env=True,