feat: Allow custom headers in file upload requests.

PiperOrigin-RevId: 811009698

Author: wanlin31

google/genai/_api_client.py

@@ -34,7 +34,7 @@
 import sys
 import threading
 import time
-from typing import Any, AsyncIterator, Iterator, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Any, AsyncIterator, Dict, Set, Iterator, Optional, Tuple, TYPE_CHECKING, Union
 from urllib.parse import urlparse
 from urllib.parse import urlunparse
 import warnings
@@ -160,6 +160,37 @@ def patch_http_options(
   return copy_option
 
 
+def _update_headers_with_append_keys(
+    original_http_options: Optional[HttpOptions],
+    patching_http_options: Optional[HttpOptionsOrDict],
+) -> Dict[str, Any]:
+  """updating headers with append logic for user-agent and x-goog-api-client, and overriding logic for the rest keys."""
+  original_headers: Dict[str, Any] = {}
+  if original_http_options:
+    original_headers = original_http_options.headers or {}
+
+  patching_headers: Dict[str, Any] = {}
+  if patching_http_options:
+    if isinstance(patching_http_options, HttpOptions):
+      patching_headers = patching_http_options.headers or {}
+    else:
+      patching_headers = patching_http_options.get('headers', {}) or {}
+
+  updated_headers = original_headers.copy()
+  append_keys = {'user-agent', 'x-goog-api-client'}
+  for key, value in patching_headers.items():
+    key = key.lower()
+    if (
+        key in append_keys
+        and key in updated_headers
+        and updated_headers[key]
+    ):
+      updated_headers[key] = f'{updated_headers[key]}, {value}'
+    else:
+      updated_headers[key] = value
+  return updated_headers
+
+
 def populate_server_timeout_header(
     headers: dict[str, str], timeout_in_seconds: Optional[Union[float, int]]
 ) -> None:
@@ -1379,6 +1410,8 @@ def upload_file(
             file, upload_url, upload_size, http_options=http_options
         )
 
+
+
   def _upload_fd(
       self,
       file: io.IOBase,
@@ -1424,7 +1457,10 @@ def _upload_fd(
             else self._http_options.timeout
         )
       timeout_in_seconds = get_timeout_in_seconds(timeout)
+
+      updated_headers = _update_headers_with_append_keys(self._http_options, http_options)
       upload_headers = {
+          **updated_headers,
           'X-Goog-Upload-Command': upload_command,
           'X-Goog-Upload-Offset': str(offset),
           'Content-Length': str(chunk_size),
@@ -1439,6 +1475,7 @@ def _upload_fd(
             content=file_chunk,
             timeout=timeout_in_seconds,
         )
+        errors.APIError.raise_for_response(response)
         if response.headers.get('x-goog-upload-status'):
           break
         delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
@@ -1455,7 +1492,7 @@ def _upload_fd(
         )
 
     if response.headers.get('x-goog-upload-status') != 'final':
-      raise ValueError('Failed to upload file: Upload status is not finalized.')
+      raise ValueError('Failed to upload file, upload status is not finalized.')
     return HttpResponse(response.headers, response_stream=[response.text])
 
   def download_file(
@@ -1550,6 +1587,7 @@ async def _async_upload_fd(
     returns:
           The HttpResponse object from the finalized request.
     """
+    updated_headers = _update_headers_with_append_keys(self._http_options, http_options)
     offset = 0
     # Upload the file in chunks
     if self._use_aiohttp():  # pylint: disable=g-import-not-at-top
@@ -1569,8 +1607,8 @@ async def _async_upload_fd(
         http_options = http_options if http_options else self._http_options
         timeout = (
             http_options.get('timeout')
-             if isinstance(http_options, dict)
-             else http_options.timeout
+            if isinstance(http_options, dict)
+            else http_options.timeout
         )
         if timeout is None:
           # Per request timeout is not configured. Check the global timeout.
@@ -1580,9 +1618,12 @@ async def _async_upload_fd(
               else self._http_options.timeout
           )
         timeout_in_seconds = get_timeout_in_seconds(timeout)
+
+        # Define and merge headers, with upload headers taking priority
         upload_headers = {
+            **updated_headers,
             'X-Goog-Upload-Command': upload_command,
-             'X-Goog-Upload-Offset': str(offset),
+            'X-Goog-Upload-Offset': str(offset),
             'Content-Length': str(chunk_size),
         }
         populate_server_timeout_header(upload_headers, timeout_in_seconds)
@@ -1597,7 +1638,7 @@ async def _async_upload_fd(
               headers=upload_headers,
               timeout=aiohttp.ClientTimeout(connect=timeout_in_seconds),
           )
-
+          await errors.APIError.raise_for_async_response(response)
           if response.headers.get('X-Goog-Upload-Status'):
             break
           delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
@@ -1654,7 +1695,10 @@ async def _async_upload_fd(
               else self._http_options.timeout
           )
         timeout_in_seconds = get_timeout_in_seconds(timeout)
+
+        # Define and merge headers, with upload headers taking priority
         upload_headers = {
+            **updated_headers,
             'X-Goog-Upload-Command': upload_command,
             'X-Goog-Upload-Offset': str(offset),
             'Content-Length': str(chunk_size),
@@ -1677,6 +1721,7 @@ async def _async_upload_fd(
               and client_response.headers.get('x-goog-upload-status')
           ):
             break
+          await errors.APIError.raise_for_async_response(client_response)
           delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
           retry_count += 1
           time.sleep(delay_seconds)

google/genai/files.py

@@ -636,33 +636,31 @@ def upload(
             'Unknown mime type: Could not determine the mimetype for your'
             ' file\n    please set the `mime_type` argument'
         )
+    request_specific_headers = (
+        (config_model.http_options.headers or {})
+        if config_model.http_options
+        else {}
+    )
+
+    # Define and merge headers, with upload headers taking priority
+    create_final_headers = {
+        **request_specific_headers,
+        'Content-Type': 'application/json',
+        'X-Goog-Upload-Protocol': 'resumable',
+        'X-Goog-Upload-Command': 'start',
+        'X-Goog-Upload-Header-Content-Length': f'{file_obj.size_bytes}',
+        'X-Goog-Upload-Header-Content-Type': f'{file_obj.mime_type}',
+    }
+
+    create_http_options = types.HttpOptions(
+        api_version='',
+        headers=create_final_headers,
+    )
 
-    http_options: types.HttpOptions
-    if config_model and config_model.http_options:
-      http_options = config_model.http_options
-      http_options.api_version = ''
-      http_options.headers = {
-          'Content-Type': 'application/json',
-          'X-Goog-Upload-Protocol': 'resumable',
-          'X-Goog-Upload-Command': 'start',
-          'X-Goog-Upload-Header-Content-Length': f'{file_obj.size_bytes}',
-          'X-Goog-Upload-Header-Content-Type': f'{file_obj.mime_type}',
-      }
-    else:
-      http_options = types.HttpOptions(
-          api_version='',
-          headers={
-              'Content-Type': 'application/json',
-              'X-Goog-Upload-Protocol': 'resumable',
-              'X-Goog-Upload-Command': 'start',
-              'X-Goog-Upload-Header-Content-Length': f'{file_obj.size_bytes}',
-              'X-Goog-Upload-Header-Content-Type': f'{file_obj.mime_type}',
-          },
-      )
     response = self._create(
         file=file_obj,
         config=types.CreateFileConfig(
-            http_options=http_options, should_return_http_response=True
+            http_options=create_http_options, should_return_http_response=True
         ),
     )
 
@@ -679,11 +677,17 @@ def upload(
 
     if isinstance(file, io.IOBase):
       return_file = self._api_client.upload_file(
-          file, upload_url, file_obj.size_bytes, http_options=http_options
+          file,
+          upload_url,
+          file_obj.size_bytes,
+          http_options=config_model.http_options,
       )
     else:
       return_file = self._api_client.upload_file(
-          fs_path, upload_url, file_obj.size_bytes, http_options=http_options
+          fs_path,
+          upload_url,
+          file_obj.size_bytes,
+          http_options=config_model.http_options,
       )
 
     return types.File._from_response(
@@ -1122,58 +1126,67 @@ async def upload(
             ' file\n    please set the `mime_type` argument'
         )
 
-    http_options: types.HttpOptions
-    if config_model and config_model.http_options:
-      http_options = config_model.http_options
-      http_options.api_version = ''
-      http_options.headers = {
-          'Content-Type': 'application/json',
-          'X-Goog-Upload-Protocol': 'resumable',
-          'X-Goog-Upload-Command': 'start',
-          'X-Goog-Upload-Header-Content-Length': f'{file_obj.size_bytes}',
-          'X-Goog-Upload-Header-Content-Type': f'{file_obj.mime_type}',
-      }
-    else:
-      http_options = types.HttpOptions(
-          api_version='',
-          headers={
-              'Content-Type': 'application/json',
-              'X-Goog-Upload-Protocol': 'resumable',
-              'X-Goog-Upload-Command': 'start',
-              'X-Goog-Upload-Header-Content-Length': f'{file_obj.size_bytes}',
-              'X-Goog-Upload-Header-Content-Type': f'{file_obj.mime_type}',
-          },
-      )
+    request_specific_headers = (
+        (config_model.http_options.headers or {})
+        if config_model.http_options
+        else {}
+    )
+
+    # Define and merge headers, with upload headers taking priority
+    create_final_headers = {
+        **request_specific_headers,
+        'Content-Type': 'application/json',
+        'X-Goog-Upload-Protocol': 'resumable',
+        'X-Goog-Upload-Command': 'start',
+        'X-Goog-Upload-Header-Content-Length': f'{file_obj.size_bytes}',
+        'X-Goog-Upload-Header-Content-Type': f'{file_obj.mime_type}',
+    }
+
+    create_http_options = types.HttpOptions(
+        api_version='',
+        headers=create_final_headers,
+    )
+
     response = await self._create(
         file=file_obj,
         config=types.CreateFileConfig(
-            http_options=http_options, should_return_http_response=True
+            http_options=create_http_options, should_return_http_response=True
         ),
     )
     if (
         response.sdk_http_response is None
         or response.sdk_http_response.headers is None
-        or (
-            'x-goog-upload-url' not in response.sdk_http_response.headers
-            and 'X-Goog-Upload-URL' not in response.sdk_http_response.headers
-        )
     ):
       raise KeyError(
-          'Failed to create file. Upload URL did not returned from the create'
-          ' file request.'
+          'Failed to create file. The SDK HTTP response or its headers were'
+          ' missing.'
+      )
+
+    upload_url = None
+    for key, value in response.sdk_http_response.headers.items():
+      if key.lower() == 'x-goog-upload-url':
+        upload_url = value
+        break  # Stop as soon as we find the first match
+
+    if upload_url is None:
+      raise KeyError(
+          'Failed to create file. Upload URL was not returned in the response'
+          ' headers.'
       )
-    elif 'x-goog-upload-url' in response.sdk_http_response.headers:
-      upload_url = response.sdk_http_response.headers['x-goog-upload-url']
-    else:
-      upload_url = response.sdk_http_response.headers['X-Goog-Upload-URL']
 
     if isinstance(file, io.IOBase):
       return_file = await self._api_client.async_upload_file(
-          file, upload_url, file_obj.size_bytes, http_options=http_options
+          file,
+          upload_url,
+          file_obj.size_bytes,
+          http_options=config_model.http_options,
       )
     else:
       return_file = await self._api_client.async_upload_file(
-          fs_path, upload_url, file_obj.size_bytes, http_options=http_options
+          fs_path,
+          upload_url,
+          file_obj.size_bytes,
+          http_options=config_model.http_options,
       )
 
     return types.File._from_response(