chore: internal

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

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]

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):

google/genai/tests/afc/test_get_usage_header.py

@@ -14,28 +14,30 @@
 
 """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
   assert config.http_options.headers['x-goog-api-client'] == expected_header
 
 
 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