chore: add headers

google/genai/_extra_utils.py

@@ -690,3 +690,31 @@ def has_agent_platform_mcp_servers(
       if getattr(tool, 'mcp_servers', None):
         return True
     return False
+
+
+def get_usage_header(
+    config: Optional[types.GenerateContentConfigOrDict] = None,
+    usage: str = 'afc',
+) -> types.GenerateContentConfig:
+  """Sets the afc version label."""
+  usage_header = f'google-genai-sdk/{usage}'
+  if not config:
+    config_model = types.GenerateContentConfig()
+  elif isinstance(config, dict):
+    config_model = types.GenerateContentConfig(**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 {}
+  if 'user-agent' in existing_headers:
+    existing_headers['user-agent'] += usage_header
+  else:
+    existing_headers['user-agent'] = usage_header
+  if 'x-goog-api-client' in existing_headers:
+    existing_headers['x-goog-api-client'] += usage_header
+  else:
+    existing_headers['x-goog-api-client'] = usage_header
+  config_model.http_options.headers = existing_headers
+  return config_model

google/genai/_replay_api_client.py

@@ -100,20 +100,32 @@ def _redact_language_label(language_label: str) -> str:
   return re.sub(r'gl-python/', '{LANGUAGE_LABEL}/', language_label)
 
 
+def _redact_sdk_usage_label(header_value: str) -> str:
+  return header_value.replace('google-genai-sdk/afc', '').replace(
+      'google-genai-sdk/chat', ''
+  ).strip()
+
+
 def _redact_request_headers(headers: dict[str, str]) -> dict[str, str]:
   """Redacts headers that should not be recorded."""
   redacted_headers = {}
   for header_name, header_value in headers.items():
     if header_name.lower() == 'x-goog-api-key':
       redacted_headers[header_name] = '{REDACTED}'
     elif header_name.lower() == 'user-agent':
-      redacted_headers[header_name] = _redact_language_label(
-          _redact_version_numbers(header_value)
+      redacted_headers[header_name] = _redact_sdk_usage_label(
+          _redact_language_label(
+              _redact_version_numbers(header_value)
+          )
       ).replace('agentplatform-genai-modules', 'vertex-genai-modules')
     elif header_name.lower() == 'x-goog-api-client':
-      redacted_headers[header_name] = _redact_language_label(
-          _redact_version_numbers(header_value)
-      ).replace('agentplatform-genai-modules', 'vertex-genai-modules')
+      redacted_headers[header_name] = _redact_sdk_usage_label(
+          _redact_language_label(
+              _redact_version_numbers(header_value)
+          )
+      ).replace(
+          'agentplatform-genai-modules', 'vertex-genai-modules'
+      )
     elif header_name.lower() == 'x-goog-user-project':
       continue
     elif header_name.lower() == 'authorization':