Skip to content

feat: Allow custom headers in file upload requests. #1452

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Allow custom headers in file upload requests. #1452

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 51 additions & 6 deletions google/genai/_api_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -1379,6 +1410,8 @@ def upload_file(
file, upload_url, upload_size, http_options=http_options
)



def _upload_fd(
self,
file: io.IOBase,
Expand Down Expand Up @@ -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),
Expand All @@ -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)
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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),
Expand All @@ -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)
Expand Down
133 changes: 73 additions & 60 deletions google/genai/files.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
),
)

Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
Loading