google-genai 1.32.0__py3-none-any.whl → 1.34.0__py3-none-any.whl

google/genai/_api_client.py

@@ -584,13 +584,9 @@ class BaseApiClient:
     # Initialize the lock. This lock will be used to protect access to the
     # credentials. This is crucial for thread safety when multiple coroutines
     # might be accessing the credentials at the same time.
-    try:
-      self._sync_auth_lock = threading.Lock()
-      self._async_auth_lock = asyncio.Lock()
-    except RuntimeError:
-      asyncio.set_event_loop(asyncio.new_event_loop())
-      self._sync_auth_lock = threading.Lock()
-      self._async_auth_lock = asyncio.Lock()
+    self._sync_auth_lock = threading.Lock()
+    self._async_auth_lock: Optional[asyncio.Lock] = None
+    self._async_auth_lock_creation_lock: Optional[asyncio.Lock] = None
 
     # Handle when to use Vertex AI in express mode (api key).
     # Explicit initializer arguments are already validated above.
@@ -903,10 +899,36 @@ class BaseApiClient:
       else:
         raise RuntimeError('Could not resolve API token from the environment')
 
+  async def _get_async_auth_lock(self) -> asyncio.Lock:
+    """Lazily initializes and returns an asyncio.Lock for async authentication.
+
+    This method ensures that a single `asyncio.Lock` instance is created and
+    shared among all asynchronous operations that require authentication,
+    preventing race conditions when accessing or refreshing credentials.
+
+    The lock is created on the first call to this method. An internal async lock
+    is used to protect the creation of the main authentication lock to ensure
+    it's a singleton within the client instance.
+
+    Returns:
+        The asyncio.Lock instance for asynchronous authentication operations.
+    """
+    if self._async_auth_lock is None:
+      # Create async creation lock if needed
+      if self._async_auth_lock_creation_lock is None:
+        self._async_auth_lock_creation_lock = asyncio.Lock()
+
+      async with self._async_auth_lock_creation_lock:
+        if self._async_auth_lock is None:
+          self._async_auth_lock = asyncio.Lock()
+
+    return self._async_auth_lock
+
   async def _async_access_token(self) -> Union[str, Any]:
     """Retrieves the access token for the credentials asynchronously."""
     if not self._credentials:
-      async with self._async_auth_lock:
+      async_auth_lock = await self._get_async_auth_lock()
+      async with async_auth_lock:
         # This ensures that only one coroutine can execute the auth logic at a
         # time for thread safety.
         if not self._credentials:
@@ -920,7 +942,8 @@ class BaseApiClient:
     if self._credentials:
       if self._credentials.expired or not self._credentials.token:
         # Only refresh when it needs to. Default expiration is 3600 seconds.
-        async with self._async_auth_lock:
+        async_auth_lock = await self._get_async_auth_lock()
+        async with async_auth_lock:
           if self._credentials.expired or not self._credentials.token:
             # Double check that the credentials expired before refreshing.
             await asyncio.to_thread(refresh_auth, self._credentials)

google/genai/_common.py

@@ -100,7 +100,11 @@ def set_value_by_path(data: Optional[dict[Any, Any]], keys: list[str], value: An
             f' Existing value: {existing_data}; New value: {value}.'
         )
     else:
-      data[keys[-1]] = value
+      if (keys[-1] == '_self' and isinstance(data, dict)
+          and isinstance(value, dict)):
+        data.update(value)
+      else:
+        data[keys[-1]] = value
 
 
 def get_value_by_path(data: Any, keys: list[str]) -> Any:

google/genai/_extra_utils.py

@@ -90,14 +90,12 @@ def _get_bigquery_uri(
 
 
 def format_destination(
-    src: Union[str, types.BatchJobSourceOrDict],
-    config: Optional[types.CreateBatchJobConfigOrDict] = None,
+    src: Union[str, types.BatchJobSource],
+    config: Optional[types.CreateBatchJobConfig] = None,
 ) -> types.CreateBatchJobConfig:
   """Formats the destination uri based on the source uri for Vertex AI."""
-  config = (
-      types._CreateBatchJobParameters(config=config).config
-      or types.CreateBatchJobConfig()
-  )
+  if config is None:
+    config = types.CreateBatchJobConfig()
 
   unique_name = None
   if not config.display_name:
@@ -113,8 +111,7 @@ def format_destination(
     elif bigquery_source_uri:
       unique_name = unique_name or _common.timestamped_unique_name()
       config.dest = f'{bigquery_source_uri}_dest_{unique_name}'
-    else:
-      raise ValueError(f'The source {src} is not supported.')
+
   return config
 
 

google/genai/_live_converters.py

@@ -165,6 +165,23 @@ def _FileData_to_mldev(
   return to_object
 
 
+def _FunctionCall_to_mldev(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['id']) is not None:
+    setv(to_object, ['id'], getv(from_object, ['id']))
+
+  if getv(from_object, ['args']) is not None:
+    setv(to_object, ['args'], getv(from_object, ['args']))
+
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  return to_object
+
+
 def _Part_to_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -203,6 +220,13 @@ def _Part_to_mldev(
         getv(from_object, ['thought_signature']),
     )
 
+  if getv(from_object, ['function_call']) is not None:
+    setv(
+        to_object,
+        ['functionCall'],
+        _FunctionCall_to_mldev(getv(from_object, ['function_call']), to_object),
+    )
+
   if getv(from_object, ['code_execution_result']) is not None:
     setv(
         to_object,
@@ -213,9 +237,6 @@ def _Part_to_mldev(
   if getv(from_object, ['executable_code']) is not None:
     setv(to_object, ['executableCode'], getv(from_object, ['executable_code']))
 
-  if getv(from_object, ['function_call']) is not None:
-    setv(to_object, ['functionCall'], getv(from_object, ['function_call']))
-
   if getv(from_object, ['function_response']) is not None:
     setv(
         to_object,
@@ -1317,6 +1338,23 @@ def _FileData_to_vertex(
   return to_object
 
 
+def _FunctionCall_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['id']) is not None:
+    raise ValueError('id parameter is not supported in Vertex AI.')
+
+  if getv(from_object, ['args']) is not None:
+    setv(to_object, ['args'], getv(from_object, ['args']))
+
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  return to_object
+
+
 def _Part_to_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -1355,6 +1393,15 @@ def _Part_to_vertex(
         getv(from_object, ['thought_signature']),
     )
 
+  if getv(from_object, ['function_call']) is not None:
+    setv(
+        to_object,
+        ['functionCall'],
+        _FunctionCall_to_vertex(
+            getv(from_object, ['function_call']), to_object
+        ),
+    )
+
   if getv(from_object, ['code_execution_result']) is not None:
     setv(
         to_object,
@@ -1365,9 +1412,6 @@ def _Part_to_vertex(
   if getv(from_object, ['executable_code']) is not None:
     setv(to_object, ['executableCode'], getv(from_object, ['executable_code']))
 
-  if getv(from_object, ['function_call']) is not None:
-    setv(to_object, ['functionCall'], getv(from_object, ['function_call']))
-
   if getv(from_object, ['function_response']) is not None:
     setv(
         to_object,
@@ -2394,6 +2438,23 @@ def _FileData_from_mldev(
   return to_object
 
 
+def _FunctionCall_from_mldev(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['id']) is not None:
+    setv(to_object, ['id'], getv(from_object, ['id']))
+
+  if getv(from_object, ['args']) is not None:
+    setv(to_object, ['args'], getv(from_object, ['args']))
+
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  return to_object
+
+
 def _Part_from_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -2432,6 +2493,15 @@ def _Part_from_mldev(
         getv(from_object, ['thoughtSignature']),
     )
 
+  if getv(from_object, ['functionCall']) is not None:
+    setv(
+        to_object,
+        ['function_call'],
+        _FunctionCall_from_mldev(
+            getv(from_object, ['functionCall']), to_object
+        ),
+    )
+
   if getv(from_object, ['codeExecutionResult']) is not None:
     setv(
         to_object,
@@ -2442,9 +2512,6 @@ def _Part_from_mldev(
   if getv(from_object, ['executableCode']) is not None:
     setv(to_object, ['executable_code'], getv(from_object, ['executableCode']))
 
-  if getv(from_object, ['functionCall']) is not None:
-    setv(to_object, ['function_call'], getv(from_object, ['functionCall']))
-
   if getv(from_object, ['functionResponse']) is not None:
     setv(
         to_object,
@@ -2591,23 +2658,6 @@ def _LiveServerContent_from_mldev(
   return to_object
 
 
-def _FunctionCall_from_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['id']) is not None:
-    setv(to_object, ['id'], getv(from_object, ['id']))
-
-  if getv(from_object, ['args']) is not None:
-    setv(to_object, ['args'], getv(from_object, ['args']))
-
-  if getv(from_object, ['name']) is not None:
-    setv(to_object, ['name'], getv(from_object, ['name']))
-
-  return to_object
-
-
 def _LiveServerToolCall_from_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -3111,6 +3161,21 @@ def _FileData_from_vertex(
   return to_object
 
 
+def _FunctionCall_from_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+
+  if getv(from_object, ['args']) is not None:
+    setv(to_object, ['args'], getv(from_object, ['args']))
+
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  return to_object
+
+
 def _Part_from_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -3149,6 +3214,15 @@ def _Part_from_vertex(
         getv(from_object, ['thoughtSignature']),
     )
 
+  if getv(from_object, ['functionCall']) is not None:
+    setv(
+        to_object,
+        ['function_call'],
+        _FunctionCall_from_vertex(
+            getv(from_object, ['functionCall']), to_object
+        ),
+    )
+
   if getv(from_object, ['codeExecutionResult']) is not None:
     setv(
         to_object,
@@ -3159,9 +3233,6 @@ def _Part_from_vertex(
   if getv(from_object, ['executableCode']) is not None:
     setv(to_object, ['executable_code'], getv(from_object, ['executableCode']))
 
-  if getv(from_object, ['functionCall']) is not None:
-    setv(to_object, ['function_call'], getv(from_object, ['functionCall']))
-
   if getv(from_object, ['functionResponse']) is not None:
     setv(
         to_object,
@@ -3263,21 +3334,6 @@ def _LiveServerContent_from_vertex(
   return to_object
 
 
-def _FunctionCall_from_vertex(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-
-  if getv(from_object, ['args']) is not None:
-    setv(to_object, ['args'], getv(from_object, ['args']))
-
-  if getv(from_object, ['name']) is not None:
-    setv(to_object, ['name'], getv(from_object, ['name']))
-
-  return to_object
-
-
 def _LiveServerToolCall_from_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,

google/genai/_replay_api_client.py

@@ -471,6 +471,21 @@ class ReplayApiClient(BaseApiClient):
     expected = interaction.response.sdk_response_segments[
         self._sdk_response_index
     ]
+    # The sdk_http_response.body has format in the string, need to get rid of
+    # the format information before comparing.
+    if isinstance(expected, dict):
+      if 'sdk_http_response' in expected and isinstance(
+          expected['sdk_http_response'], dict
+      ):
+        if 'body' in expected['sdk_http_response']:
+          raw_body = expected['sdk_http_response']['body']
+          print('raw_body length: ', len(raw_body))
+          print('raw_body: ', raw_body)
+          if isinstance(raw_body, str) and raw_body != '':
+            raw_body = json.loads(raw_body)
+            raw_body = json.dumps(raw_body)
+            expected['sdk_http_response']['body'] = raw_body
+
     assert (
         actual == expected
     ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}'

google/genai/_tokens_converters.py

@@ -165,6 +165,23 @@ def _FileData_to_mldev(
   return to_object
 
 
+def _FunctionCall_to_mldev(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['id']) is not None:
+    setv(to_object, ['id'], getv(from_object, ['id']))
+
+  if getv(from_object, ['args']) is not None:
+    setv(to_object, ['args'], getv(from_object, ['args']))
+
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  return to_object
+
+
 def _Part_to_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -203,6 +220,13 @@ def _Part_to_mldev(
         getv(from_object, ['thought_signature']),
     )
 
+  if getv(from_object, ['function_call']) is not None:
+    setv(
+        to_object,
+        ['functionCall'],
+        _FunctionCall_to_mldev(getv(from_object, ['function_call']), to_object),
+    )
+
   if getv(from_object, ['code_execution_result']) is not None:
     setv(
         to_object,
@@ -213,9 +237,6 @@ def _Part_to_mldev(
   if getv(from_object, ['executable_code']) is not None:
     setv(to_object, ['executableCode'], getv(from_object, ['executable_code']))
 
-  if getv(from_object, ['function_call']) is not None:
-    setv(to_object, ['functionCall'], getv(from_object, ['function_call']))
-
   if getv(from_object, ['function_response']) is not None:
     setv(
         to_object,

google/genai/_transformers.py

@@ -964,30 +964,30 @@ def t_cached_content_name(client: _api_client.BaseApiClient, name: str) -> str:
 
 def t_batch_job_source(
     client: _api_client.BaseApiClient,
-    src: Union[
-        str, List[types.InlinedRequestOrDict], types.BatchJobSourceOrDict
-    ],
+    src: types.BatchJobSourceUnionDict,
 ) -> types.BatchJobSource:
   if isinstance(src, dict):
     src = types.BatchJobSource(**src)
   if isinstance(src, types.BatchJobSource):
+    vertex_sources = sum(
+        [src.gcs_uri is not None, src.bigquery_uri is not None]
+    )
+    mldev_sources = sum([
+        src.inlined_requests is not None,
+        src.file_name is not None,
+    ])
     if client.vertexai:
-      if src.gcs_uri and src.bigquery_uri:
+      if mldev_sources or vertex_sources != 1:
         raise ValueError(
-            'Only one of `gcs_uri` or `bigquery_uri` can be set.'
-        )
-      elif not src.gcs_uri and not src.bigquery_uri:
-        raise ValueError(
-            'One of `gcs_uri` or `bigquery_uri` must be set.'
+            'Exactly one of `gcs_uri` or `bigquery_uri` must be set, other '
+            'sources are not supported in Vertex AI.'
         )
     else:
-      if src.inlined_requests and src.file_name:
+      if vertex_sources or mldev_sources != 1:
         raise ValueError(
-            'Only one of `inlined_requests` or `file_name` can be set.'
-        )
-      elif not src.inlined_requests and not src.file_name:
-        raise ValueError(
-            'One of `inlined_requests` or `file_name` must be set.'
+            'Exactly one of `inlined_requests`, `file_name`, '
+            '`inlined_embed_content_requests`, or `embed_content_file_name` '
+            'must be set, other sources are not supported in Gemini API.'
         )
     return src
 
@@ -1012,6 +1012,29 @@ def t_batch_job_source(
   raise ValueError(f'Unsupported source: {src}')
 
 
+def t_embedding_batch_job_source(
+    client: _api_client.BaseApiClient,
+    src: types.EmbeddingsBatchJobSourceOrDict,
+) -> types.EmbeddingsBatchJobSource:
+  if isinstance(src, dict):
+    src = types.EmbeddingsBatchJobSource(**src)
+
+  if isinstance(src, types.EmbeddingsBatchJobSource):
+    mldev_sources = sum([
+        src.inlined_requests is not None,
+        src.file_name is not None,
+    ])
+    if mldev_sources != 1:
+      raise ValueError(
+          'Exactly one of `inlined_requests`, `file_name`, '
+          '`inlined_embed_content_requests`, or `embed_content_file_name` '
+          'must be set, other sources are not supported in Gemini API.'
+      )
+    return src
+  else:
+    raise ValueError(f'Unsupported source type: {type(src)}')
+
+
 def t_batch_job_destination(
     dest: Union[str, types.BatchJobDestinationOrDict],
 ) -> types.BatchJobDestination:
@@ -1037,6 +1060,23 @@ def t_batch_job_destination(
     raise ValueError(f'Unsupported destination: {dest}')
 
 
+def t_recv_batch_job_destination(dest: dict[str, Any]) -> dict[str, Any]:
+  # Rename inlinedResponses if it looks like an embedding response.
+  inline_responses = dest.get('inlinedResponses', {}).get(
+      'inlinedResponses', []
+  )
+  if not inline_responses:
+    return dest
+  for response in inline_responses:
+    inner_response = response.get('response', {})
+    if not inner_response:
+      continue
+    if 'embedding' in inner_response:
+      dest['inlinedEmbedContentResponses'] = dest.pop('inlinedResponses')
+      break
+  return dest
+
+
 def t_batch_job_name(client: _api_client.BaseApiClient, name: str) -> str:
   if not client.vertexai:
     mldev_pattern = r'batches/[^/]+$'