diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index 20ad6a1fe..8e1445fdd 100644 --- a/google/genai/_extra_utils.py +++ b/google/genai/_extra_utils.py @@ -19,11 +19,12 @@ import inspect import io import logging -import sys -import typing -from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin import mimetypes import os +import sys +import typing +from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, get_args, get_origin + import pydantic from . import _common @@ -47,6 +48,9 @@ McpClientSession: typing.Type = Any McpTool: typing.Type = Any + +C = TypeVar('C', bound='_common.BaseModel') + _DEFAULT_MAX_REMOTE_CALLS_AFC = 10 logger = logging.getLogger('google_genai.models') @@ -698,21 +702,29 @@ def has_agent_platform_mcp_servers( def get_usage_header( - config: Optional[types.GenerateContentConfigOrDict] = None, - usage: str = 'afc', -) -> types.GenerateContentConfig: - """Sets the afc version label.""" + config: Optional[Union[dict[str, Any], C]], config_cls: Type[C], usage: str +) -> C: + """Returns the usage header for the config.""" usage_header = f'google-genai-sdk/{public_version.__version__}+{usage}' if not config: - config_model = types.GenerateContentConfig() + config_model = config_cls() elif isinstance(config, dict): - config_model = types.GenerateContentConfig(**config) + config_model = config_cls(**config) else: config_model = config - if not config_model.http_options: - config_model.http_options = types.HttpOptions() - existing_headers = config_model.http_options.headers or {} + # Many configs have http_options, safely initialize it if it's missing + if not hasattr(config_model, 'http_options'): + return config_model + + http_options = getattr(config_model, 'http_options', None) + if http_options is None: + http_options = types.HttpOptions() + setattr(config_model, 'http_options', http_options) + + http_options = typing.cast(types.HttpOptions, http_options) + existing_headers = http_options.headers or {} + for header_key in ('user-agent', 'x-goog-api-client'): if header_key in existing_headers: if ( @@ -730,5 +742,6 @@ def get_usage_header( existing_headers[header_key] += f' {usage_header}' else: existing_headers[header_key] = usage_header - config_model.http_options.headers = existing_headers + + http_options.headers = existing_headers return config_model diff --git a/google/genai/chats.py b/google/genai/chats.py index 8782f47ce..8ab765db0 100644 --- a/google/genai/chats.py +++ b/google/genai/chats.py @@ -115,7 +115,9 @@ def __init__( history: list[ContentOrDict], ): self._model = model - self._config = _extra_utils.get_usage_header(config, usage="chat") + self._config = _extra_utils.get_usage_header( + config, types.GenerateContentConfig, usage="chat" # type: ignore[arg-type] + ) content_models = [] for content in history: if not isinstance(content, Content): @@ -253,7 +255,7 @@ def send_message( input_content = t.t_content(message) method_config = config if config else self._config method_config = _extra_utils.get_usage_header( - method_config, usage="chat" + method_config, types.GenerateContentConfig, usage="chat" # type: ignore[arg-type] ) response = self._modules.generate_content( model=self._model, @@ -314,7 +316,7 @@ def send_message_stream( chunk = None method_config = config if config else self._config method_config = _extra_utils.get_usage_header( - method_config, usage="chat" + method_config, types.GenerateContentConfig, usage="chat" # type: ignore[arg-type] ) if isinstance(self._modules, Models): for chunk in self._modules.generate_content_stream( @@ -423,7 +425,7 @@ async def send_message( input_content = t.t_content(message) method_config = config if config else self._config method_config = _extra_utils.get_usage_header( - method_config, usage="chat" + method_config, types.GenerateContentConfig, usage="chat" # type: ignore[arg-type] ) response = await self._modules.generate_content( model=self._model, @@ -481,7 +483,7 @@ async def send_message_stream( method_config = config if config else self._config method_config = _extra_utils.get_usage_header( - method_config, usage="chat" + method_config, types.GenerateContentConfig, usage="chat" # type: ignore[arg-type] ) async def async_generator(): # type: ignore[no-untyped-def] diff --git a/google/genai/models.py b/google/genai/models.py index 3e34a0bf0..d77346acf 100644 --- a/google/genai/models.py +++ b/google/genai/models.py @@ -6596,7 +6596,9 @@ def generate_content( function_map = _extra_utils.get_function_map(parsed_config) if function_map: parsed_config_to_call = _extra_utils.get_usage_header( - parsed_config_to_call + parsed_config_to_call, + types.GenerateContentConfig, + 'afc', ) i += 1 response = self._generate_content( @@ -6771,7 +6773,7 @@ def generate_content_stream( function_map = _extra_utils.get_function_map(parsed_config) if function_map: parsed_config_to_call = _extra_utils.get_usage_header( - parsed_config_to_call + parsed_config_to_call, types.GenerateContentConfig, 'afc' ) i += 1 response = self._generate_content_stream( @@ -7091,7 +7093,7 @@ def generate_videos( time.sleep(10) operation = client.operations.get(operation) - operation.response.generated_videos[0].video.uri + operation.result.generated_videos[0].video.uri ``` """ if prompt or image or video: @@ -7109,6 +7111,11 @@ def generate_videos( stacklevel=2, ) Models._logged_generate_videos_deprecation_warning = True + _extra_utils.get_usage_header( + config, # type: ignore[arg-type] + types.GenerateVideosConfig, + 'nonsource', + ) # Gemini Developer API does not support video bytes. video_dct: dict[str, Any] = {} if not self._api_client.vertexai and video: @@ -8804,7 +8811,9 @@ async def generate_content( ) if function_map: final_parsed_config_to_call = _extra_utils.get_usage_header( - final_parsed_config_to_call + final_parsed_config_to_call, + types.GenerateContentConfig, + usage='afc', ) response = await self._generate_content( model=model, contents=contents, config=final_parsed_config_to_call @@ -9058,7 +9067,9 @@ async def stream_generator(): # type: ignore[no-untyped-def] ) if function_map: final_parsed_config_to_call = _extra_utils.get_usage_header( - final_parsed_config_to_call + final_parsed_config_to_call, + types.GenerateContentConfig, + usage='afc', ) i += 1 @@ -9463,7 +9474,12 @@ async def generate_videos( stacklevel=2, ) AsyncModels._logged_generate_videos_deprecation_warning = True - # Gemini Developer API does not support video bytes. + # Gemini Developer API does not support video bytes. + _extra_utils.get_usage_header( + config, # type: ignore[arg-type] + types.GenerateVideosConfig, + 'nonsource', + ) video_dct: dict[str, Any] = {} if not self._api_client.vertexai and video: if isinstance(video, types.Video): diff --git a/google/genai/tests/afc/test_get_usage_header.py b/google/genai/tests/afc/test_get_usage_header.py index e53ab9965..120306b03 100644 --- a/google/genai/tests/afc/test_get_usage_header.py +++ b/google/genai/tests/afc/test_get_usage_header.py @@ -14,20 +14,22 @@ """Unit tests for _extra_utils.get_usage_header.""" +from ... import _extra_utils from ... import types from ... import version as public_version -from ..._extra_utils import get_usage_header + +get_usage_header = _extra_utils.get_usage_header def test_get_usage_header_none_config(): - config = get_usage_header(None) + config = get_usage_header(None, types.GenerateContentConfig, 'afc') expected_header = f'google-genai-sdk/{public_version.__version__}+afc' assert config.http_options.headers['user-agent'] == expected_header assert config.http_options.headers['x-goog-api-client'] == expected_header def test_get_usage_header_dict_config(): - config = get_usage_header({'temperature': 0.5}, usage='afc') + config = get_usage_header({'temperature': 0.5}, types.GenerateContentConfig, 'afc') expected_header = f'google-genai-sdk/{public_version.__version__}+afc' assert config.temperature == 0.5 assert config.http_options.headers['user-agent'] == expected_header @@ -35,7 +37,7 @@ def test_get_usage_header_dict_config(): def test_get_usage_header_custom_usage(): - config = get_usage_header(None, usage='chat') + config = get_usage_header(None, types.GenerateContentConfig, 'chat') expected_header = f'google-genai-sdk/{public_version.__version__}+chat' assert config.http_options.headers['user-agent'] == expected_header assert config.http_options.headers['x-goog-api-client'] == expected_header @@ -53,7 +55,7 @@ def test_get_usage_header_with_existing_sdk_headers(): config = types.GenerateContentConfig( http_options=types.HttpOptions(headers=initial_headers) ) - config = get_usage_header(config, usage='afc') + config = get_usage_header(config, types.GenerateContentConfig, 'afc') expected_header = ( f'google-genai-sdk/{public_version.__version__}+afc gl-python/3.12' ) @@ -62,19 +64,19 @@ def test_get_usage_header_with_existing_sdk_headers(): def test_get_usage_header_idempotent_no_duplicate_usage(): - config = get_usage_header(None, usage='afc') - config = get_usage_header(config, usage='afc') + config = get_usage_header(None, types.GenerateContentConfig, 'afc') + config = get_usage_header(config, types.GenerateContentConfig, 'afc') expected_header = f'google-genai-sdk/{public_version.__version__}+afc' assert config.http_options.headers['user-agent'] == expected_header assert config.http_options.headers['x-goog-api-client'] == expected_header def test_get_usage_header_multiple_usages_no_duplicate(): - config = get_usage_header(None, usage='chat') - config = get_usage_header(config, usage='afc') + config = get_usage_header(None, types.GenerateContentConfig, 'chat') + config = get_usage_header(config, types.GenerateContentConfig, 'afc') # Call again with chat and afc to ensure no duplicate entries are appended - config = get_usage_header(config, usage='chat') - config = get_usage_header(config, usage='afc') + config = get_usage_header(config, types.GenerateContentConfig, 'chat') + config = get_usage_header(config, types.GenerateContentConfig, 'afc') user_agent = config.http_options.headers['user-agent'] x_goog_api_client = config.http_options.headers['x-goog-api-client'] @@ -98,9 +100,9 @@ def test_get_usage_header_with_custom_user_agent(): config = types.GenerateContentConfig( http_options=types.HttpOptions(headers=initial_headers) ) - config = get_usage_header(config, usage='afc') + config = get_usage_header(config, types.GenerateContentConfig, 'afc') # Call again to ensure idempotency - config = get_usage_header(config, usage='afc') + config = get_usage_header(config, types.GenerateContentConfig, 'afc') expected_usage = f'google-genai-sdk/{public_version.__version__}+afc' assert 'custom-agent/1.0' in config.http_options.headers['user-agent'] @@ -115,8 +117,24 @@ def test_get_usage_header_preserves_other_http_options(): headers={'custom-header': 'value'}, ) ) - config = get_usage_header(config, usage='afc') + config = get_usage_header(config, types.GenerateContentConfig, 'afc') assert config.http_options.api_version == 'v1alpha' assert config.http_options.headers['custom-header'] == 'value' assert '+afc' in config.http_options.headers['user-agent'] assert '+afc' in config.http_options.headers['x-goog-api-client'] + + +def test_get_usage_header_with_generate_images_config(): + config = types.GenerateImagesConfig() + config = get_usage_header(config, types.GenerateImagesConfig, 'image') + expected_header = f'google-genai-sdk/{public_version.__version__}+image' + assert config.http_options.headers['user-agent'] == expected_header + assert config.http_options.headers['x-goog-api-client'] == expected_header + + +def test_get_usage_header_dict_generate_images_config(): + config = get_usage_header({'negative_prompt': 'blue'}, types.GenerateImagesConfig, 'image') + expected_header = f'google-genai-sdk/{public_version.__version__}+image' + assert config.negative_prompt == 'blue' + assert config.http_options.headers['user-agent'] == expected_header + assert config.http_options.headers['x-goog-api-client'] == expected_header