google-genai 1.9.0__py3-none-any.whl → 1.10.0__py3-none-any.whl

google/genai/_api_client.py

@@ -19,18 +19,21 @@
 The BaseApiClient is intended to be a private module and is subject to change.
 """
 
-import anyio
 import asyncio
 import copy
 from dataclasses import dataclass
 import datetime
+import http
 import io
 import json
 import logging
+import math
 import os
 import sys
+import time
 from typing import Any, AsyncIterator, Optional, Tuple, Union
 from urllib.parse import urlparse, urlunparse
+import anyio
 import google.auth
 import google.auth.credentials
 from google.auth.credentials import Credentials
@@ -95,6 +98,14 @@ def _patch_http_options(
   return copy_option
 
 
+def _populate_server_timeout_header(
+    headers: dict[str, str], timeout_in_seconds: Optional[Union[float, int]]
+) -> None:
+  """Populates the server timeout header in the headers dict."""
+  if timeout_in_seconds and 'X-Server-Timeout' not in headers:
+    headers['X-Server-Timeout'] = str(math.ceil(timeout_in_seconds))
+
+
 def _join_url_path(base_url: str, path: str) -> str:
   parsed_base = urlparse(base_url)
   base_path = (
@@ -128,6 +139,19 @@ def _refresh_auth(credentials: Credentials) -> Credentials:
   return credentials
 
 
+def _get_timeout_in_seconds(
+    timeout: Optional[Union[float, int]],
+) -> Optional[float]:
+  """Converts the timeout to seconds."""
+  if timeout:
+    # HttpOptions.timeout is in milliseconds. But httpx.Client.request()
+    # expects seconds.
+    timeout_in_seconds = timeout / 1000.0
+  else:
+    timeout_in_seconds = None
+  return timeout_in_seconds
+
+
 @dataclass
 class HttpRequest:
   headers: dict[str, str]
@@ -520,18 +544,13 @@ class BaseApiClient:
         versioned_path,
     )
 
-    timeout_in_seconds: Optional[Union[float, int]] = (
-        patched_http_options.timeout
-    )
-    if timeout_in_seconds:
-      # HttpOptions.timeout is in milliseconds. But httpx.Client.request()
-      # expects seconds.
-      timeout_in_seconds = timeout_in_seconds / 1000.0
-    else:
-      timeout_in_seconds = None
+    timeout_in_seconds = _get_timeout_in_seconds(patched_http_options.timeout)
 
     if patched_http_options.headers is None:
       raise ValueError('Request headers must be set.')
+    _populate_server_timeout_header(
+        patched_http_options.headers, timeout_in_seconds
+    )
     return HttpRequest(
         method=http_method,
         url=url,
@@ -712,7 +731,12 @@ class BaseApiClient:
     return async_generator()
 
   def upload_file(
-      self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int
+      self,
+      file_path: Union[str, io.IOBase],
+      upload_url: str,
+      upload_size: int,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ) -> HttpResponse:
     """Transfers a file to the given URL.
 
@@ -723,18 +747,28 @@ class BaseApiClient:
       upload_url: The URL to upload the file to.
       upload_size: The size of file content to be uploaded, this will have to
         match the size requested in the resumable upload request.
+      http_options: The http options to use for the request.
 
     returns:
           The HttpResponse object from the finalize request.
     """
     if isinstance(file_path, io.IOBase):
-      return self._upload_fd(file_path, upload_url, upload_size)
+      return self._upload_fd(
+          file_path, upload_url, upload_size, http_options=http_options
+      )
     else:
       with open(file_path, 'rb') as file:
-        return self._upload_fd(file, upload_url, upload_size)
+        return self._upload_fd(
+            file, upload_url, upload_size, http_options=http_options
+        )
 
   def _upload_fd(
-      self, file: io.IOBase, upload_url: str, upload_size: int
+      self,
+      file: io.IOBase,
+      upload_url: str,
+      upload_size: int,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ) -> HttpResponse:
     """Transfers a file to the given URL.
 
@@ -743,6 +777,7 @@ class BaseApiClient:
       upload_url: The URL to upload the file to.
       upload_size: The size of file content to be uploaded, this will have to
         match the size requested in the resumable upload request.
+      http_options: The http options to use for the request.
 
     returns:
           The HttpResponse object from the finalize request.
@@ -758,15 +793,32 @@ class BaseApiClient:
       # If last chunk, finalize the upload.
       if chunk_size + offset >= upload_size:
         upload_command += ', finalize'
+      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 timeout is None:
+        # Per request timeout is not configured. Check the global timeout.
+        timeout = (
+            self._http_options.timeout
+            if isinstance(self._http_options, dict)
+            else self._http_options.timeout
+        )
+      timeout_in_seconds = _get_timeout_in_seconds(timeout)
+      upload_headers = {
+          'X-Goog-Upload-Command': upload_command,
+          'X-Goog-Upload-Offset': str(offset),
+          'Content-Length': str(chunk_size),
+      }
+      _populate_server_timeout_header(upload_headers, timeout_in_seconds)
       response = self._httpx_client.request(
           method='POST',
           url=upload_url,
-          headers={
-              'X-Goog-Upload-Command': upload_command,
-              'X-Goog-Upload-Offset': str(offset),
-              'Content-Length': str(chunk_size),
-          },
+          headers=upload_headers,
           content=file_chunk,
+          timeout=timeout_in_seconds,
       )
       offset += chunk_size
       if response.headers['x-goog-upload-status'] != 'active':
@@ -783,7 +835,12 @@ class BaseApiClient:
       )
     return HttpResponse(response.headers, response_stream=[response.text])
 
-  def download_file(self, path: str, http_options):
+  def download_file(
+      self,
+      path: str,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
+  ):
     """Downloads the file data.
 
     Args:
@@ -822,6 +879,8 @@ class BaseApiClient:
       file_path: Union[str, io.IOBase],
       upload_url: str,
       upload_size: int,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ) -> HttpResponse:
     """Transfers a file asynchronously to the given URL.
 
@@ -831,23 +890,30 @@ class BaseApiClient:
       upload_url: The URL to upload the file to.
       upload_size: The size of file content to be uploaded, this will have to
         match the size requested in the resumable upload request.
+      http_options: The http options to use for the request.
 
     returns:
           The HttpResponse object from the finalize request.
     """
     if isinstance(file_path, io.IOBase):
-      return await self._async_upload_fd(file_path, upload_url, upload_size)
+      return await self._async_upload_fd(
+          file_path, upload_url, upload_size, http_options=http_options
+      )
     else:
       file = anyio.Path(file_path)
       fd = await file.open('rb')
       async with fd:
-        return await self._async_upload_fd(fd, upload_url, upload_size)
+        return await self._async_upload_fd(
+            fd, upload_url, upload_size, http_options=http_options
+        )
 
   async def _async_upload_fd(
       self,
       file: Union[io.IOBase, anyio.AsyncFile],
       upload_url: str,
       upload_size: int,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ) -> HttpResponse:
     """Transfers a file asynchronously to the given URL.
 
@@ -856,6 +922,7 @@ class BaseApiClient:
       upload_url: The URL to upload the file to.
       upload_size: The size of file content to be uploaded, this will have to
         match the size requested in the resumable upload request.
+      http_options: The http options to use for the request.
 
     returns:
           The HttpResponse object from the finalized request.
@@ -874,15 +941,32 @@ class BaseApiClient:
       # If last chunk, finalize the upload.
       if chunk_size + offset >= upload_size:
         upload_command += ', finalize'
+      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 timeout is None:
+        # Per request timeout is not configured. Check the global timeout.
+        timeout = (
+            self._http_options.timeout
+            if isinstance(self._http_options, dict)
+            else self._http_options.timeout
+        )
+      timeout_in_seconds = _get_timeout_in_seconds(timeout)
+      upload_headers = {
+          'X-Goog-Upload-Command': upload_command,
+          'X-Goog-Upload-Offset': str(offset),
+          'Content-Length': str(chunk_size),
+      }
+      _populate_server_timeout_header(upload_headers, timeout_in_seconds)
       response = await self._async_httpx_client.request(
           method='POST',
           url=upload_url,
           content=file_chunk,
-          headers={
-              'X-Goog-Upload-Command': upload_command,
-              'X-Goog-Upload-Offset': str(offset),
-              'Content-Length': str(chunk_size),
-          },
+          headers=upload_headers,
+          timeout=timeout_in_seconds,
       )
       offset += chunk_size
       if response.headers.get('x-goog-upload-status') != 'active':
@@ -899,7 +983,12 @@ class BaseApiClient:
       )
     return HttpResponse(response.headers, response_stream=[response.text])
 
-  async def async_download_file(self, path: str, http_options):
+  async def async_download_file(
+      self,
+      path: str,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
+  ):
     """Downloads the file data.
 
     Args:

google/genai/_automatic_function_calling_util.py

@@ -28,7 +28,7 @@ from . import types
 if sys.version_info >= (3, 10):
   VersionedUnionType = builtin_types.UnionType
 else:
-  VersionedUnionType = typing._UnionGenericAlias
+  VersionedUnionType = typing._UnionGenericAlias  # type: ignore[attr-defined]
 
 _py_builtin_type_to_schema_type = {
     str: types.Type.STRING,

google/genai/_extra_utils.py

@@ -30,7 +30,7 @@ from . import types
 if sys.version_info >= (3, 10):
   from types import UnionType
 else:
-  UnionType = typing._UnionGenericAlias
+  UnionType = typing._UnionGenericAlias  # type: ignore[attr-defined]
 
 _DEFAULT_MAX_REMOTE_CALLS_AFC = 10
 

google/genai/_replay_api_client.py

@@ -34,6 +34,8 @@ from ._api_client import HttpOptions
 from ._api_client import HttpRequest
 from ._api_client import HttpResponse
 from ._common import BaseModel
+from .types import HttpOptionsOrDict
+from .types import GenerateVideosOperation
 
 
 def _redact_version_numbers(version_string: str) -> str:
@@ -396,7 +398,12 @@ class ReplayApiClient(BaseApiClient):
     if isinstance(response_model, list):
       response_model = response_model[0]
     print('response_model: ', response_model.model_dump(exclude_none=True))
-    actual = response_model.model_dump(exclude_none=True, mode='json')
+    if isinstance(response_model, GenerateVideosOperation):
+      actual = response_model.model_dump(
+          exclude={'result'}, exclude_none=True, mode='json'
+      )
+    else:
+      actual = response_model.model_dump(exclude_none=True, mode='json')
     expected = interaction.response.sdk_response_segments[
         self._sdk_response_index
     ]
@@ -461,7 +468,14 @@ class ReplayApiClient(BaseApiClient):
     else:
       return self._build_response_from_replay(http_request)
 
-  def upload_file(self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int) -> HttpResponse:
+  def upload_file(
+      self,
+      file_path: Union[str, io.IOBase],
+      upload_url: str,
+      upload_size: int,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
+  ) -> HttpResponse:
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
       content = file_path.read()
@@ -479,7 +493,9 @@ class ReplayApiClient(BaseApiClient):
     if self._should_call_api():
       result: Union[str, HttpResponse]
       try:
-        result = super().upload_file(file_path, upload_url, upload_size)
+        result = super().upload_file(
+            file_path, upload_url, upload_size, http_options=http_options
+        )
       except HTTPError as e:
         result = HttpResponse(
             dict(e.response.headers), [json.dumps({'reason': e.response.reason})]
@@ -496,6 +512,8 @@ class ReplayApiClient(BaseApiClient):
       file_path: Union[str, io.IOBase],
       upload_url: str,
       upload_size: int,
+      *,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ) -> HttpResponse:
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
@@ -515,7 +533,7 @@ class ReplayApiClient(BaseApiClient):
       result: HttpResponse
       try:
         result = await super().async_upload_file(
-            file_path, upload_url, upload_size
+            file_path, upload_url, upload_size, http_options=http_options
         )
       except HTTPError as e:
         result = HttpResponse(
@@ -528,14 +546,16 @@ class ReplayApiClient(BaseApiClient):
     else:
       return self._build_response_from_replay(request)
 
-  def download_file(self, path: str, http_options: HttpOptions):
+  def download_file(
+      self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None
+  ):
     self._initialize_replay_session_if_not_loaded()
     request = self._build_request(
         'get', path=path, request_dict={}, http_options=http_options
     )
     if self._should_call_api():
       try:
-        result = super().download_file(path, http_options)
+        result = super().download_file(path, http_options=http_options)
       except HTTPError as e:
         result = HttpResponse(
             dict(e.response.headers), [json.dumps({'reason': e.response.reason})]
@@ -547,14 +567,18 @@ class ReplayApiClient(BaseApiClient):
     else:
       return self._build_response_from_replay(request).byte_stream[0]
 
-  async def async_download_file(self, path: str, http_options):
+  async def async_download_file(
+      self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None
+  ):
     self._initialize_replay_session_if_not_loaded()
     request = self._build_request(
         'get', path=path, request_dict={}, http_options=http_options
     )
     if self._should_call_api():
       try:
-        result = await super().async_download_file(path, http_options)
+        result = await super().async_download_file(
+            path, http_options=http_options
+        )
       except HTTPError as e:
         result = HttpResponse(
             dict(e.response.headers), [json.dumps({'reason': e.response.reason})]

google/genai/_transformers.py

@@ -43,7 +43,7 @@ if sys.version_info >= (3, 10):
   _UNION_TYPES = (typing.Union, builtin_types.UnionType)
   from typing import TypeGuard
 else:
-  VersionedUnionType = typing._UnionGenericAlias
+  VersionedUnionType = typing._UnionGenericAlias  # type: ignore[attr-defined]
   _UNION_TYPES = (typing.Union,)
   from typing_extensions import TypeGuard
 
@@ -657,51 +657,55 @@ def process_schema(
       )
 
   if schema.get('title') == 'PlaceholderLiteralEnum':
-    schema.pop('title', None)
-
-  # If a dict is provided directly to response_schema, it may use `any_of`
-  # instead of `anyOf`. Otherwise model_json_schema() uses `anyOf`
-  if schema.get('any_of', None) is not None:
-    schema['anyOf'] = schema.pop('any_of')
+    del schema['title']
+
+  # Standardize spelling for relevant schema fields.  For example, if a dict is
+  # provided directly to response_schema, it may use `any_of` instead of `anyOf.
+  # Otherwise, model_json_schema() uses `anyOf`.
+  for from_name, to_name in [
+      ('any_of', 'anyOf'),
+      ('property_ordering', 'propertyOrdering'),
+  ]:
+    if (value := schema.pop(from_name, None)) is not None:
+      schema[to_name] = value
 
   if defs is None:
     defs = schema.pop('$defs', {})
     for _, sub_schema in defs.items():
-      process_schema(sub_schema, client, defs)
+      # We can skip the '$ref' check, because JSON schema forbids a '$ref' from
+      # directly referencing another '$ref':
+      # https://json-schema.org/understanding-json-schema/structuring#recursion
+      process_schema(
+          sub_schema, client, defs, order_properties=order_properties
+      )
 
   handle_null_fields(schema)
 
   # After removing null fields, Optional fields with only one possible type
   # will have a $ref key that needs to be flattened
   # For example: {'default': None, 'description': 'Name of the person', 'nullable': True, '$ref': '#/$defs/TestPerson'}
-  schema_ref = schema.get('$ref', None)
-  if schema_ref is not None:
-    ref = defs[schema_ref.split('defs/')[-1]]
-    for schema_key in list(ref.keys()):
-      schema[schema_key] = ref[schema_key]
-    del schema['$ref']
-
-  any_of = schema.get('anyOf', None)
-  if any_of is not None:
-    for sub_schema in any_of:
-      # $ref is present in any_of if the schema is a union of Pydantic classes
-      ref_key = sub_schema.get('$ref', None)
-      if ref_key is None:
-        process_schema(sub_schema, client, defs)
-      else:
-        ref = defs[ref_key.split('defs/')[-1]]
-        any_of.append(ref)
-    schema['anyOf'] = [item for item in any_of if '$ref' not in item]
+  if (ref := schema.pop('$ref', None)) is not None:
+    schema.update(defs[ref.split('defs/')[-1]])
+
+  def _recurse(sub_schema: dict[str, Any]) -> dict[str, Any]:
+    """Returns the processed `sub_schema`, resolving its '$ref' if any."""
+    if (ref := sub_schema.pop('$ref', None)) is not None:
+      sub_schema = defs[ref.split('defs/')[-1]]
+    process_schema(sub_schema, client, defs, order_properties=order_properties)
+    return sub_schema
+
+  if (any_of := schema.get('anyOf')) is not None:
+    schema['anyOf'] = [_recurse(sub_schema) for sub_schema in any_of]
     return
 
-  schema_type = schema.get('type', None)
+  schema_type = schema.get('type')
   if isinstance(schema_type, Enum):
     schema_type = schema_type.value
   schema_type = schema_type.upper()
 
   # model_json_schema() returns a schema with a 'const' field when a Literal with one value is provided as a pydantic field
   # For example `genre: Literal['action']` becomes: {'const': 'action', 'title': 'Genre', 'type': 'string'}
-  const = schema.get('const', None)
+  const = schema.get('const')
   if const is not None:
     if schema_type == 'STRING':
       schema['enum'] = [const]
@@ -710,38 +714,18 @@ def process_schema(
       raise ValueError('Literal values must be strings.')
 
   if schema_type == 'OBJECT':
-    properties = schema.get('properties', None)
-    if properties is None:
-      return
-    for name, sub_schema in properties.items():
-      ref_key = sub_schema.get('$ref', None)
-      if ref_key is None:
-        process_schema(sub_schema, client, defs)
-      else:
-        ref = defs[ref_key.split('defs/')[-1]]
-        process_schema(ref, client, defs)
-        properties[name] = ref
-    if (
-        len(properties.items()) > 1
-        and order_properties
-        and all(
-            ordering_key not in schema
-            for ordering_key in ['property_ordering', 'propertyOrdering']
-        )
-    ):
-      property_names = list(properties.keys())
-      schema['property_ordering'] = property_names
+    if (properties := schema.get('properties')) is not None:
+      for name, sub_schema in list(properties.items()):
+        properties[name] = _recurse(sub_schema)
+      if (
+          len(properties.items()) > 1
+          and order_properties
+          and 'propertyOrdering' not in schema
+      ):
+        schema['property_ordering'] = list(properties.keys())
   elif schema_type == 'ARRAY':
-    sub_schema = schema.get('items', None)
-    if sub_schema is None:
-      return
-    ref_key = sub_schema.get('$ref', None)
-    if ref_key is None:
-      process_schema(sub_schema, client, defs)
-    else:
-      ref = defs[ref_key.split('defs/')[-1]]
-      process_schema(ref, client, defs)
-      schema['items'] = ref
+    if (items := schema.get('items')) is not None:
+      schema['items'] = _recurse(items)
 
 
 def _process_enum(

google/genai/files.py

@@ -658,6 +658,14 @@ class Files(_api_module.BaseModule):
     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='',
@@ -685,11 +693,11 @@ class Files(_api_module.BaseModule):
 
     if isinstance(file, io.IOBase):
       return_file = self._api_client.upload_file(
-          file, upload_url, file_obj.size_bytes
+          file, upload_url, file_obj.size_bytes, http_options=http_options
       )
     else:
       return_file = self._api_client.upload_file(
-          fs_path, upload_url, file_obj.size_bytes
+          fs_path, upload_url, file_obj.size_bytes, http_options=http_options
       )
 
     return types.File._from_response(
@@ -778,7 +786,7 @@ class Files(_api_module.BaseModule):
 
     data = self._api_client.download_file(
         path,
-        http_options,
+        http_options=http_options,
     )
 
     if isinstance(file, types.Video):
@@ -1122,6 +1130,14 @@ class AsyncFiles(_api_module.BaseModule):
     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='',
@@ -1148,11 +1164,11 @@ class AsyncFiles(_api_module.BaseModule):
 
     if isinstance(file, io.IOBase):
       return_file = await self._api_client.async_upload_file(
-          file, upload_url, file_obj.size_bytes
+          file, upload_url, file_obj.size_bytes, http_options=http_options
       )
     else:
       return_file = await self._api_client.async_upload_file(
-          fs_path, upload_url, file_obj.size_bytes
+          fs_path, upload_url, file_obj.size_bytes, http_options=http_options
       )
 
     return types.File._from_response(
@@ -1231,7 +1247,7 @@ class AsyncFiles(_api_module.BaseModule):
 
     data = await self._api_client.async_download_file(
         path,
-        http_options,
+        http_options=http_options,
     )
 
     return data