google-genai 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

google/genai/_api_client.py

@@ -29,16 +29,14 @@ import json
 import logging
 import os
 import sys
-from typing import Any, AsyncIterator, Optional, Tuple, TypedDict, Union
+from typing import Any, AsyncIterator, Optional, Tuple, Union
 from urllib.parse import urlparse, urlunparse
 import google.auth
 import google.auth.credentials
 from google.auth.credentials import Credentials
-from google.auth.transport.requests import AuthorizedSession
 from google.auth.transport.requests import Request
 import httpx
-from pydantic import BaseModel, ConfigDict, Field, ValidationError
-import requests
+from pydantic import BaseModel, Field, ValidationError
 from . import _common
 from . import errors
 from . import version
@@ -88,7 +86,8 @@ def _patch_http_options(
       copy_option[patch_key].update(patch_value)
     elif patch_value is not None:  # Accept empty values.
       copy_option[patch_key] = patch_value
-  _append_library_version_headers(copy_option['headers'])
+  if copy_option['headers']:
+    _append_library_version_headers(copy_option['headers'])
   return copy_option
 
 
@@ -103,7 +102,7 @@ def _join_url_path(base_url: str, path: str) -> str:
   return urlunparse(parsed_base._replace(path=base_path + '/' + path))
 
 
-def _load_auth(*, project: Union[str, None]) -> tuple[Credentials, str]:
+def _load_auth(*, project: Union[str, None]) -> Tuple[Credentials, str]:
   """Loads google auth credentials and project id."""
   credentials, loaded_project_id = google.auth.default(
       scopes=['https://www.googleapis.com/auth/cloud-platform'],
@@ -273,7 +272,9 @@ class BaseApiClient:
     validated_http_options: dict[str, Any]
     if isinstance(http_options, dict):
       try:
-        validated_http_options = HttpOptions.model_validate(http_options).model_dump()
+        validated_http_options = HttpOptions.model_validate(
+            http_options
+        ).model_dump()
       except ValidationError as e:
         raise ValueError(f'Invalid http_options: {e}')
     elif isinstance(http_options, HttpOptions):
@@ -359,7 +360,9 @@ class BaseApiClient:
       self._http_options['headers']['x-goog-api-key'] = self.api_key
     # Update the http options with the user provided http options.
     if http_options:
-      self._http_options = _patch_http_options(self._http_options, validated_http_options)
+      self._http_options = _patch_http_options(
+          self._http_options, validated_http_options
+      )
     else:
       _append_library_version_headers(self._http_options['headers'])
 
@@ -367,8 +370,27 @@ class BaseApiClient:
     url_parts = urlparse(self._http_options['base_url'])
     return url_parts._replace(scheme='wss').geturl()
 
-  async def _async_access_token(self) -> str:
+  def _access_token(self) -> str:
     """Retrieves the access token for the credentials."""
+    if not self._credentials:
+      self._credentials, project = _load_auth(project=self.project)
+      if not self.project:
+        self.project = project
+
+    if self._credentials:
+      if (
+          self._credentials.expired or not self._credentials.token
+      ):
+        # Only refresh when it needs to. Default expiration is 3600 seconds.
+        _refresh_auth(self._credentials)
+      if not self._credentials.token:
+        raise RuntimeError('Could not resolve API token from the environment')
+      return self._credentials.token
+    else:
+      raise RuntimeError('Could not resolve API token from the environment')
+
+  async def _async_access_token(self) -> str:
+    """Retrieves the access token for the credentials asynchronously."""
     if not self._credentials:
       async with self._auth_lock:
         # This ensures that only one coroutine can execute the auth logic at a
@@ -437,8 +459,8 @@ class BaseApiClient:
     ):
       path = f'projects/{self.project}/locations/{self.location}/' + path
     url = _join_url_path(
-        patched_http_options['base_url'],
-        patched_http_options['api_version'] + '/' + path,
+        patched_http_options.get('base_url', ''),
+        patched_http_options.get('api_version', '') + '/' + path,
     )
 
     timeout_in_seconds: Optional[Union[float, int]] = patched_http_options.get(
@@ -464,59 +486,56 @@ class BaseApiClient:
       http_request: HttpRequest,
       stream: bool = False,
   ) -> HttpResponse:
+    data: Optional[Union[str, bytes]] = None
     if self.vertexai and not self.api_key:
-      if not self._credentials:
-        self._credentials, _ = _load_auth(project=self.project)
-      if self._credentials.quota_project_id:
+      http_request.headers['Authorization'] = (
+          f'Bearer {self._access_token()}'
+      )
+      if self._credentials and self._credentials.quota_project_id:
         http_request.headers['x-goog-user-project'] = (
             self._credentials.quota_project_id
         )
-      authed_session = AuthorizedSession(self._credentials)
-      authed_session.stream = stream
-      response = authed_session.request(
-          http_request.method.upper(),
-          http_request.url,
+      data = json.dumps(http_request.data)
+    else:
+      if http_request.data:
+        if not isinstance(http_request.data, bytes):
+          data = json.dumps(http_request.data)
+        else:
+          data = http_request.data
+
+    if stream:
+      client = httpx.Client()
+      httpx_request = client.build_request(
+          method=http_request.method,
+          url=http_request.url,
+          content=data,
           headers=http_request.headers,
-          data=json.dumps(http_request.data) if http_request.data else None,
           timeout=http_request.timeout,
       )
+      response = client.send(httpx_request, stream=stream)
       errors.APIError.raise_for_response(response)
       return HttpResponse(
           response.headers, response if stream else [response.text]
       )
     else:
-      return self._request_unauthorized(http_request, stream)
-
-  def _request_unauthorized(
-      self,
-      http_request: HttpRequest,
-      stream: bool = False,
-  ) -> HttpResponse:
-    data: Optional[Union[str, bytes]] = None
-    if http_request.data:
-      if not isinstance(http_request.data, bytes):
-        data = json.dumps(http_request.data)
-      else:
-        data = http_request.data
-
-    http_session = requests.Session()
-    response = http_session.request(
-        method=http_request.method,
-        url=http_request.url,
-        headers=http_request.headers,
-        data=data,
-        timeout=http_request.timeout,
-        stream=stream,
-    )
-    errors.APIError.raise_for_response(response)
-    return HttpResponse(
-        response.headers, response if stream else [response.text]
-    )
+      with httpx.Client() as client:
+        response = client.request(
+            method=http_request.method,
+            url=http_request.url,
+            headers=http_request.headers,
+            content=data,
+            timeout=http_request.timeout,
+        )
+        errors.APIError.raise_for_response(response)
+        return HttpResponse(
+            response.headers, response if stream else [response.text]
+        )
 
   async def _async_request(
       self, http_request: HttpRequest, stream: bool = False
   ):
-    if self.vertexai:
+    data: Optional[Union[str, bytes]] = None
+    if self.vertexai and not self.api_key:
       http_request.headers['Authorization'] = (
           f'Bearer {await self._async_access_token()}'
       )
@@ -524,12 +543,20 @@ class BaseApiClient:
         http_request.headers['x-goog-user-project'] = (
             self._credentials.quota_project_id
         )
+      data = json.dumps(http_request.data)
+    else:
+      if http_request.data:
+        if not isinstance(http_request.data, bytes):
+          data = json.dumps(http_request.data)
+        else:
+          data = http_request.data
+
     if stream:
       aclient = httpx.AsyncClient()
       httpx_request = aclient.build_request(
           method=http_request.method,
           url=http_request.url,
-          content=json.dumps(http_request.data),
+          content=data,
           headers=http_request.headers,
           timeout=http_request.timeout,
       )
@@ -547,7 +574,7 @@ class BaseApiClient:
             method=http_request.method,
             url=http_request.url,
             headers=http_request.headers,
-            content=json.dumps(http_request.data) if http_request.data else None,
+            content=data,
             timeout=http_request.timeout,
         )
         errors.APIError.raise_for_response(response)
@@ -633,7 +660,7 @@ class BaseApiClient:
 
   def upload_file(
       self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int
-  ) -> str:
+  ) -> dict[str, str]:
     """Transfers a file to the given URL.
 
     Args:
@@ -655,7 +682,7 @@ class BaseApiClient:
 
   def _upload_fd(
       self, file: io.IOBase, upload_url: str, upload_size: int
-  ) -> str:
+  ) -> dict[str, str]:
     """Transfers a file to the given URL.
 
     Args:
@@ -689,7 +716,7 @@ class BaseApiClient:
           data=file_chunk,
       )
 
-      response = self._request_unauthorized(request, stream=False)
+      response = self._request(request, stream=False)
       offset += chunk_size
       if response.headers['X-Goog-Upload-Status'] != 'active':
         break  # upload is complete or it has been interrupted.
@@ -732,25 +759,24 @@ class BaseApiClient:
       else:
         data = http_request.data
 
-    http_session = requests.Session()
-    response = http_session.request(
-        method=http_request.method,
-        url=http_request.url,
-        headers=http_request.headers,
-        data=data,
-        timeout=http_request.timeout,
-        stream=False,
-    )
+    with httpx.Client(follow_redirects=True) as client:
+      response = client.request(
+          method=http_request.method,
+          url=http_request.url,
+          headers=http_request.headers,
+          content=data,
+          timeout=http_request.timeout,
+      )
 
-    errors.APIError.raise_for_response(response)
-    return HttpResponse(response.headers, byte_stream=[response.content])
+      errors.APIError.raise_for_response(response)
+      return HttpResponse(response.headers, byte_stream=[response.read()])
 
   async def async_upload_file(
       self,
       file_path: Union[str, io.IOBase],
       upload_url: str,
       upload_size: int,
-  ) -> str:
+  ) -> dict[str, str]:
     """Transfers a file asynchronously to the given URL.
 
     Args:
@@ -776,7 +802,7 @@ class BaseApiClient:
       file: Union[io.IOBase, anyio.AsyncFile],
       upload_url: str,
       upload_size: int,
-  ) -> str:
+  ) -> dict[str, str]:
     """Transfers a file asynchronously to the given URL.
 
     Args:
@@ -842,7 +868,7 @@ class BaseApiClient:
         'get', path=path, request_dict={}, http_options=http_options
     )
 
-    data: Optional[Union[str, bytes]]
+    data: Optional[Union[str, bytes]] = None
     if http_request.data:
       if not isinstance(http_request.data, bytes):
         data = json.dumps(http_request.data)

google/genai/_automatic_function_calling_util.py

@@ -17,7 +17,7 @@ import inspect
 import sys
 import types as builtin_types
 import typing
-from typing import _GenericAlias, Any, Callable, get_args, get_origin, Literal, Optional, Union
+from typing import _GenericAlias, Any, Callable, get_args, get_origin, Literal, Optional, Union  # type: ignore[attr-defined]
 
 import pydantic
 
@@ -41,19 +41,11 @@ _py_builtin_type_to_schema_type = {
 
 
 def _is_builtin_primitive_or_compound(
-    annotation: inspect.Parameter.annotation,
+    annotation: inspect.Parameter.annotation,  # type: ignore[valid-type]
 ) -> bool:
   return annotation in _py_builtin_type_to_schema_type.keys()
 
 
-def _raise_for_any_of_if_mldev(schema: types.Schema):
-  if schema.any_of:
-    raise ValueError(
-        'AnyOf is not supported in function declaration schema for'
-        ' the Gemini API.'
-    )
-
-
 def _raise_for_default_if_mldev(schema: types.Schema):
   if schema.default is not None:
     raise ValueError(
@@ -64,12 +56,11 @@ def _raise_for_default_if_mldev(schema: types.Schema):
 
 def _raise_if_schema_unsupported(api_option: Literal['VERTEX_AI', 'GEMINI_API'], schema: types.Schema):
   if api_option == 'GEMINI_API':
-    _raise_for_any_of_if_mldev(schema)
     _raise_for_default_if_mldev(schema)
 
 
 def _is_default_value_compatible(
-    default_value: Any, annotation: inspect.Parameter.annotation
+    default_value: Any, annotation: inspect.Parameter.annotation  # type: ignore[valid-type]
 ) -> bool:
   # None type is expected to be handled external to this function
   if _is_builtin_primitive_or_compound(annotation):
@@ -292,8 +283,7 @@ def _parse_schema_from_parameter(
           ),
           func_name,
       )
-    if api_option == 'VERTEX_AI':
-      schema.required = _get_required_fields(schema)
+    schema.required = _get_required_fields(schema)
     _raise_if_schema_unsupported(api_option, schema)
     return schema
   raise ValueError(

google/genai/_transformers.py

@@ -26,9 +26,7 @@ import sys
 import time
 import types as builtin_types
 import typing
-from typing import Any, GenericAlias, Optional, Union
-
-import types as builtin_types
+from typing import Any, GenericAlias, Optional, Union  # type: ignore[attr-defined]
 
 if typing.TYPE_CHECKING:
   import PIL.Image
@@ -43,10 +41,11 @@ logger = logging.getLogger('google_genai._transformers')
 if sys.version_info >= (3, 10):
   VersionedUnionType = builtin_types.UnionType
   _UNION_TYPES = (typing.Union, builtin_types.UnionType)
+  from typing import TypeGuard
 else:
   VersionedUnionType = typing._UnionGenericAlias
   _UNION_TYPES = (typing.Union,)
-
+  from typing_extensions import TypeGuard
 
 def _resource_name(
     client: _api_client.BaseApiClient,
@@ -165,7 +164,9 @@ def t_model(client: _api_client.BaseApiClient, model: str):
       return f'models/{model}'
 
 
-def t_models_url(api_client: _api_client.BaseApiClient, base_models: bool) -> str:
+def t_models_url(
+    api_client: _api_client.BaseApiClient, base_models: bool
+) -> str:
   if api_client.vertexai:
     if base_models:
       return 'publishers/google/models'
@@ -179,7 +180,8 @@ def t_models_url(api_client: _api_client.BaseApiClient, base_models: bool) -> st
 
 
 def t_extract_models(
-    api_client: _api_client.BaseApiClient, response: dict[str, list[types.ModelDict]]
+    api_client: _api_client.BaseApiClient,
+    response: dict[str, list[types.ModelDict]],
 ) -> Optional[list[types.ModelDict]]:
   if not response:
     return []
@@ -240,9 +242,7 @@ def pil_to_blob(img) -> types.Blob:
   return types.Blob(mime_type=mime_type, data=data)
 
 
-def t_part(
-    part: Optional[types.PartUnionDict]
-) -> types.Part:
+def t_part(part: Optional[types.PartUnionDict]) -> types.Part:
   try:
     import PIL.Image
 
@@ -268,7 +268,7 @@ def t_part(
 
 
 def t_parts(
-    parts: Optional[Union[list[types.PartUnionDict], types.PartUnionDict]],
+    parts: Optional[Union[list[types.PartUnionDict], types.PartUnionDict, list[types.Part]]],
 ) -> list[types.Part]:
   #
   if parts is None or (isinstance(parts, list) and not parts):
@@ -332,22 +332,35 @@ def t_content(
 def t_contents_for_embed(
     client: _api_client.BaseApiClient,
     contents: Union[list[types.Content], list[types.ContentDict], ContentType],
-):
-  if client.vertexai and isinstance(contents, list):
-    # TODO: Assert that only text is supported.
-    return [t_content(client, content).parts[0].text for content in contents]
-  elif client.vertexai:
-    return [t_content(client, contents).parts[0].text]
-  elif isinstance(contents, list):
-    return [t_content(client, content) for content in contents]
+) -> Union[list[str], list[types.Content]]:
+  if isinstance(contents, list):
+    transformed_contents = [t_content(client, content) for content in contents]
   else:
-    return [t_content(client, contents)]
+    transformed_contents = [t_content(client, contents)]
+
+  if client.vertexai:
+    text_parts = []
+    for content in transformed_contents:
+      if content is not None:
+        if isinstance(content, dict):
+          content = types.Content.model_validate(content)
+        if content.parts is not None:
+          for part in content.parts:
+            if part.text:
+              text_parts.append(part.text)
+            else:
+              logger.warning(
+                  f'Non-text part found, only returning text parts.'
+              )
+    return text_parts
+  else:
+    return transformed_contents
 
 
 def t_contents(
     client: _api_client.BaseApiClient,
     contents: Optional[
-        Union[types.ContentListUnion, types.ContentListUnionDict]
+        Union[types.ContentListUnion, types.ContentListUnionDict, types.Content]
     ],
 ) -> list[types.Content]:
   if contents is None or (isinstance(contents, list) and not contents):
@@ -365,7 +378,7 @@ def t_contents(
   result: list[types.Content] = []
   accumulated_parts: list[types.Part] = []
 
-  def _is_part(part: types.PartUnionDict) -> bool:
+  def _is_part(part: Union[types.PartUnionDict, Any]) -> TypeGuard[types.PartUnionDict]:
     if (
         isinstance(part, str)
         or isinstance(part, types.File)
@@ -429,11 +442,11 @@ def t_contents(
     ):
       _append_accumulated_parts_as_content(result, accumulated_parts)
       if isinstance(content, list):
-        result.append(types.UserContent(parts=content))
+        result.append(types.UserContent(parts=content))  # type: ignore[arg-type]
       else:
         result.append(content)
-    elif (_is_part(content)): # type: ignore
-      _handle_current_part(result, accumulated_parts, content) # type: ignore
+    elif (_is_part(content)):
+      _handle_current_part(result, accumulated_parts, content)
     elif isinstance(content, dict):
       # PactDict is already handled in _is_part
       result.append(types.Content.model_validate(content))
@@ -499,7 +512,7 @@ def handle_null_fields(schema: dict[str, Any]):
         schema['anyOf'].remove({'type': 'null'})
         if len(schema['anyOf']) == 1:
           # If there is only one type left after removing null, remove the anyOf field.
-          for key,val in schema['anyOf'][0].items():
+          for key, val in schema['anyOf'][0].items():
             schema[key] = val
           del schema['anyOf']
 
@@ -574,7 +587,8 @@ def process_schema(
 
     if schema.get('default') is not None:
       raise ValueError(
-          'Default value is not supported in the response schema for the Gemini API.'
+          'Default value is not supported in the response schema for the Gemini'
+          ' API.'
       )
 
   if schema.get('title') == 'PlaceholderLiteralEnum':
@@ -604,10 +618,6 @@ def process_schema(
 
   any_of = schema.get('anyOf', None)
   if any_of is not None:
-    if client and not client.vertexai:
-      raise ValueError(
-          'AnyOf is not supported in the response schema for the Gemini API.'
-      )
     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)
@@ -670,7 +680,7 @@ def process_schema(
 
 
 def _process_enum(
-    enum: EnumMeta, client: Optional[_api_client.BaseApiClient] = None  # type: ignore
+    enum: EnumMeta, client: _api_client.BaseApiClient
 ) -> types.Schema:
   for member in enum:  # type: ignore
     if not isinstance(member.value, str):
@@ -680,7 +690,7 @@ def _process_enum(
       )
 
   class Placeholder(pydantic.BaseModel):
-    placeholder: enum
+    placeholder: enum  # type: ignore[valid-type]
 
   enum_schema = Placeholder.model_json_schema()
   process_schema(enum_schema, client)
@@ -688,12 +698,19 @@ def _process_enum(
   return types.Schema.model_validate(enum_schema)
 
 
+def _is_type_dict_str_any(origin: Union[types.SchemaUnionDict, Any]) -> TypeGuard[dict[str, Any]]:
+  """Verifies the schema is of type dict[str, Any] for mypy type checking."""
+  return isinstance(origin, dict) and all(
+      isinstance(key, str) for key in origin
+  )
+
+
 def t_schema(
     client: _api_client.BaseApiClient, origin: Union[types.SchemaUnionDict, Any]
 ) -> Optional[types.Schema]:
   if not origin:
     return None
-  if isinstance(origin, dict):
+  if isinstance(origin, dict) and _is_type_dict_str_any(origin):
     process_schema(origin, client, order_properties=False)
     return types.Schema.model_validate(origin)
   if isinstance(origin, EnumMeta):
@@ -724,7 +741,7 @@ def t_schema(
   ):
 
     class Placeholder(pydantic.BaseModel):
-      placeholder: origin
+      placeholder: origin  # type: ignore[valid-type]
 
     schema = Placeholder.model_json_schema()
     process_schema(schema, client)
@@ -735,7 +752,8 @@ def t_schema(
 
 
 def t_speech_config(
-    _: _api_client.BaseApiClient, origin: Union[types.SpeechConfigUnionDict, Any]
+    _: _api_client.BaseApiClient,
+    origin: Union[types.SpeechConfigUnionDict, Any],
 ) -> Optional[types.SpeechConfig]:
   if not origin:
     return None
@@ -794,7 +812,10 @@ def t_tools(
     transformed_tool = t_tool(client, tool)
     # All functions should be merged into one tool.
     if transformed_tool is not None:
-      if transformed_tool.function_declarations:
+      if (
+          transformed_tool.function_declarations
+          and function_tool.function_declarations is not None
+      ):
         function_tool.function_declarations += (
             transformed_tool.function_declarations
         )
@@ -896,7 +917,10 @@ def t_file_name(
   elif isinstance(name, types.Video):
     name = name.uri
   elif isinstance(name, types.GeneratedVideo):
-    name = name.video.uri
+    if name.video is not None:
+      name = name.video.uri
+    else:
+      name = None
 
   if name is None:
     raise ValueError('File name is required.')

google/genai/batches.py

@@ -998,6 +998,8 @@ class Batches(_api_module.BaseModule):
       for batch_job in batch_jobs:
         print(f"Batch job: {batch_job.name}, state {batch_job.state}")
     """
+    if config is None:
+      config = types.ListBatchJobsConfig()
     return Pager(
         'batch_jobs',
         self._list,
@@ -1373,6 +1375,8 @@ class AsyncBatches(_api_module.BaseModule):
       await batch_jobs_pager.next_page()
       print(f"next page: {batch_jobs_pager.page}")
     """
+    if config is None:
+      config = types.ListBatchJobsConfig()
     return AsyncPager(
         'batch_jobs',
         self._list,