diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 12101cd01..9f2e55181 100644 --- a/google/genai/_api_client.py +++ b/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) diff --git a/google/genai/files.py b/google/genai/files.py index 4d02e8f79..45028224e 100644 --- a/google/genai/files.py +++ b/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( diff --git a/google/genai/tests/client/test_http_options.py b/google/genai/tests/client/test_http_options.py index 9b89892aa..cc606911d 100644 --- a/google/genai/tests/client/test_http_options.py +++ b/google/genai/tests/client/test_http_options.py @@ -86,6 +86,101 @@ def test_patch_http_options_appends_version_headers(): assert 'x-goog-api-client' in patched.headers +def test_update_headers_appends_headers_for_reserved_keys(): + client_options = types.HttpOptions( + headers={ + 'user-agent': 'client_user_agent', + 'x-goog-api-client': 'client_x_goog_api_client', + } + ) + config_options = types.HttpOptions( + headers={ + 'user-agent': 'config_user_agent', + 'x-goog-api-client': 'config_x_goog_api_client', + } + ) + updated_headers = _api_client._update_headers_with_append_keys( + client_options, + config_options, + ) + assert updated_headers['user-agent'] == 'client_user_agent, config_user_agent' + assert ( + updated_headers['x-goog-api-client'] + == 'client_x_goog_api_client, config_x_goog_api_client' + ) + + +def test_update_headers_behavior_for_non_reserved_keys(): + client_options = types.HttpOptions( + headers={ + 'user-agent': 'client_user_agent', + 'x-goog-api-client': 'client_x_goog_api_client', + 'custom-header': 'client_custom_header', + 'custom-header-2': 'client_custom_header_2', + } + ) + config_options = types.HttpOptions( + headers={ + 'user-agent': 'config_user_agent', + 'x-goog-api-client': 'config_x_goog_api_client', + 'custom-header': 'config_custom_header', + 'custom-header-3': 'config_custom_header_3', + } + ) + updated_headers = _api_client._update_headers_with_append_keys( + client_options, + config_options, + ) + assert updated_headers['user-agent'] == 'client_user_agent, config_user_agent' + assert ( + updated_headers['x-goog-api-client'] + == 'client_x_goog_api_client, config_x_goog_api_client' + ) + assert updated_headers['custom-header'] == 'config_custom_header' + assert updated_headers['custom-header-2'] == 'client_custom_header_2' + assert updated_headers['custom-header-3'] == 'config_custom_header_3' + + +def test_update_headers_behavior_for_missing_input(): + client_options = types.HttpOptions( + headers={ + 'user-agent': 'client_user_agent', + 'x-goog-api-client': 'client_x_goog_api_client', + 'custom-header': 'client_custom_header', + 'custom-header-2': 'client_custom_header_2', + } + ) + config_options = types.HttpOptions( + headers={ + 'user-agent': 'config_user_agent', + 'x-goog-api-client': 'config_x_goog_api_client', + 'custom-header': 'config_custom_header', + 'custom-header-3': 'config_custom_header_3', + } + ) + # No original headers, no config options, result should be + # empty + update_headers_1 = _api_client._update_headers_with_append_keys( + None, None + ) + assert update_headers_1 == {} + + # No original headers, config options,result should be + # config headers + update_headers_2 = _api_client._update_headers_with_append_keys( + None, config_options + ) + assert update_headers_2 == config_options.headers + + # Original headers, no config options, no reserved keys, result should be + # original headers + update_headers_3 = _api_client._update_headers_with_append_keys( + client_options, None + ) + assert update_headers_3 == client_options.headers + + + def test_setting_timeout_populates_server_timeout_header(): api_client = _api_client.BaseApiClient( vertexai=False, diff --git a/google/genai/tests/files/test_upload.py b/google/genai/tests/files/test_upload.py index d12918877..11fd01969 100644 --- a/google/genai/tests/files/test_upload.py +++ b/google/genai/tests/files/test_upload.py @@ -32,6 +32,23 @@ test_table=test_table, ) +def test_image_png_upload_with_config_headers(client): + with pytest_helper.exception_if_vertex(client, ValueError): + new_headers = { + **client._api_client._http_options.headers, + 'custom-header': 'test_value_set_by_client', + 'client-header-2': 'test_value_set_by_client_2', + 'Content-Type': 'will_be_overridden', + } + client._api_client._http_options.headers = new_headers + file = client.files.upload( + file='tests/data/google.png', + config=types.UploadFileConfig( + http_options={ + 'headers': {'custom-header': 'test_value_set_by_config'} + }), + ) + assert file.name.startswith('files/') def test_image_png_upload(client): with pytest_helper.exception_if_vertex(client, ValueError): @@ -183,28 +200,50 @@ def test_audio_m4a_upload_with_config_dict(client): ) assert file.name.startswith('files/') +@pytest.mark.asyncio +async def test_image_png_upload_with_config_headers_async(client): + with pytest_helper.exception_if_vertex(client, ValueError): + async with client.aio as aio_client: + new_headers = { + **aio_client._api_client._http_options.headers, + 'custom-header': 'test_value_set_by_client', + 'client-header-2': 'test_value_set_by_client_2', + 'Content-Type': 'will_be_overridden', + } + aio_client._api_client._http_options.headers = new_headers + file = await aio_client.files.upload( + file='tests/data/google.png', + config=types.UploadFileConfig( + http_options={ + 'headers': {'custom-header': 'test_value_set_by_config'} + }), + ) + assert file.name.startswith('files/') @pytest.mark.asyncio async def test_image_upload_async(client): with pytest_helper.exception_if_vertex(client, ValueError): - file = await client.aio.files.upload(file='tests/data/google.png') - assert file.name.startswith('files/') + async with client.aio as aio_client: + file = await aio_client.files.upload(file='tests/data/google.png') + assert file.name.startswith('files/') @pytest.mark.asyncio async def test_image_upload_with_config_async(client): with pytest_helper.exception_if_vertex(client, ValueError): - file = await client.aio.files.upload( - file='tests/data/google.png', - config=types.UploadFileConfig(display_name='test_image'), - ) - assert file.name.startswith('files/') + async with client.aio as aio_client: + file = await aio_client.files.upload( + file='tests/data/google.png', + config=types.UploadFileConfig(display_name='test_image'), + ) + assert file.name.startswith('files/') @pytest.mark.asyncio async def test_image_upload_with_config_dict_async(client): with pytest_helper.exception_if_vertex(client, ValueError): - file = await client.aio.files.upload( + async with client.aio as aio_client: + file = await aio_client.files.upload( file='tests/data/google.png', config={ 'display_name': 'test_image', @@ -217,14 +256,15 @@ async def test_image_upload_with_config_dict_async(client): @pytest.mark.asyncio async def test_image_upload_with_bytesio_async(client): with pytest_helper.exception_if_vertex(client, ValueError): - with open('tests/data/google.png', 'rb') as f: - buffer = io.BytesIO(f.read()) - file = await client.aio.files.upload( - file=buffer, - config=types.UploadFileConfig( - mime_type='image/png'), - ) - assert file.name.startswith('files/') + async with client.aio as aio_client: + with open('tests/data/google.png', 'rb') as f: + buffer = io.BytesIO(f.read()) + file = await aio_client.files.upload( + file=buffer, + config=types.UploadFileConfig( + mime_type='image/png'), + ) + assert file.name.startswith('files/') @pytest.mark.asyncio