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

google/genai/_api_client.py

@@ -20,6 +20,7 @@ The BaseApiClient is intended to be a private module and is subject to change.
 """
 
 import asyncio
+from collections.abc import Awaitable, Generator
 import copy
 from dataclasses import dataclass
 import datetime
@@ -29,21 +30,31 @@ import json
 import logging
 import math
 import os
+import ssl
 import sys
 import time
 from typing import Any, AsyncIterator, Optional, Tuple, Union
-from urllib.parse import urlparse, urlunparse
+from urllib.parse import urlparse
+from urllib.parse import urlunparse
+
 import anyio
+import certifi
 import google.auth
 import google.auth.credentials
 from google.auth.credentials import Credentials
 from google.auth.transport.requests import Request
 import httpx
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import ValidationError
+
 from . import _common
 from . import errors
 from . import version
-from .types import HttpOptions, HttpOptionsDict, HttpOptionsOrDict
+from .types import HttpOptions
+from .types import HttpOptionsDict
+from .types import HttpOptionsOrDict
+
 
 logger = logging.getLogger('google_genai._api_client')
 CHUNK_SIZE = 8 * 1024 * 1024  # 8 MB chunk size
@@ -119,7 +130,7 @@ def _join_url_path(base_url: str, path: str) -> 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(
+  credentials, loaded_project_id = google.auth.default(  # type: ignore[no-untyped-call]
       scopes=['https://www.googleapis.com/auth/cloud-platform'],
   )
 
@@ -135,7 +146,7 @@ def _load_auth(*, project: Union[str, None]) -> Tuple[Credentials, str]:
 
 
 def _refresh_auth(credentials: Credentials) -> Credentials:
-  credentials.refresh(Request())
+  credentials.refresh(Request())  # type: ignore[no-untyped-call]
   return credentials
 
 
@@ -181,17 +192,17 @@ class HttpResponse:
       response_stream: Union[Any, str] = None,
       byte_stream: Union[Any, bytes] = None,
   ):
-    self.status_code = 200
+    self.status_code: int = 200
     self.headers = headers
     self.response_stream = response_stream
     self.byte_stream = byte_stream
 
   # Async iterator for async streaming.
-  def __aiter__(self):
+  def __aiter__(self) -> 'HttpResponse':
     self.segment_iterator = self.async_segments()
     return self
 
-  async def __anext__(self):
+  async def __anext__(self) -> Any:
     try:
       return await self.segment_iterator.__anext__()
     except StopIteration:
@@ -203,7 +214,7 @@ class HttpResponse:
       return ''
     return json.loads(self.response_stream[0])
 
-  def segments(self):
+  def segments(self) -> Generator[Any, None, None]:
     if isinstance(self.response_stream, list):
       # list of objects retrieved from replay or from non-streaming API.
       for chunk in self.response_stream:
@@ -212,7 +223,7 @@ class HttpResponse:
       yield from []
     else:
       # Iterator of objects retrieved from the API.
-      for chunk in self.response_stream.iter_lines():
+      for chunk in self.response_stream.iter_lines():  # type: ignore[union-attr]
         if chunk:
           # In streaming mode, the chunk of JSON is prefixed with "data:" which
           # we must strip before parsing.
@@ -246,7 +257,7 @@ class HttpResponse:
       else:
         raise ValueError('Error parsing streaming response.')
 
-  def byte_segments(self):
+  def byte_segments(self) -> Generator[Union[bytes, Any], None, None]:
     if isinstance(self.byte_stream, list):
       # list of objects retrieved from replay or from non-streaming API.
       yield from self.byte_stream
@@ -257,7 +268,7 @@ class HttpResponse:
           'Byte segments are not supported for streaming responses.'
       )
 
-  def _copy_to_dict(self, response_payload: dict[str, object]):
+  def _copy_to_dict(self, response_payload: dict[str, object]) -> None:
     # Cannot pickle 'generator' object.
     delattr(self, 'segment_iterator')
     for attribute in dir(self):
@@ -414,7 +425,7 @@ class BaseApiClient:
       if not self.api_key:
         raise ValueError(
             'Missing key inputs argument! To use the Google AI API,'
-            'provide (`api_key`) arguments. To use the Google Cloud API,'
+            ' provide (`api_key`) arguments. To use the Google Cloud API,'
             ' provide (`vertexai`, `project` & `location`) arguments.'
         )
       self._http_options.base_url = 'https://generativelanguage.googleapis.com/'
@@ -432,13 +443,71 @@ class BaseApiClient:
     else:
       if self._http_options.headers is not None:
         _append_library_version_headers(self._http_options.headers)
-    # Initialize the httpx client.
-    self._httpx_client = SyncHttpxClient()
-    self._async_httpx_client = AsyncHttpxClient()
 
-  def _websocket_base_url(self):
+    client_args, async_client_args = self._ensure_ssl_ctx(self._http_options)
+    self._httpx_client = SyncHttpxClient(**client_args)
+    self._async_httpx_client = AsyncHttpxClient(**async_client_args)
+
+  @staticmethod
+  def _ensure_ssl_ctx(options: HttpOptions) -> (
+      Tuple[dict[str, Any], dict[str, Any]]):
+    """Ensures the SSL context is present in the client args.
+
+    Creates a default SSL context if one is not provided.
+
+    Args:
+      options: The http options to check for SSL context.
+
+    Returns:
+      A tuple of sync/async httpx client args.
+    """
+
+    verify = 'verify'
+    args = options.client_args
+    async_args = options.async_client_args
+    ctx = (
+        args.get(verify) if args else None
+        or async_args.get(verify) if async_args else None
+    )
+
+    if not ctx:
+      # Initialize the SSL context for the httpx client.
+      # Unlike requests, the httpx package does not automatically pull in the
+      # environment variables SSL_CERT_FILE or SSL_CERT_DIR. They need to be
+      # enabled explicitly.
+      ctx = ssl.create_default_context(
+          cafile=os.environ.get('SSL_CERT_FILE', certifi.where()),
+          capath=os.environ.get('SSL_CERT_DIR'),
+      )
+
+    def _maybe_set(
+        args: Optional[dict[str, Any]],
+        ctx: ssl.SSLContext,
+    ) -> dict[str, Any]:
+      """Sets the SSL context in the client args if not set.
+
+      Does not override the SSL context if it is already set.
+
+      Args:
+        args: The client args to to check for SSL context.
+        ctx: The SSL context to set.
+
+      Returns:
+        The client args with the SSL context included.
+      """
+      if not args or not args.get(verify):
+        args = (args or {}).copy()
+        args[verify] = ctx
+      return args
+
+    return (
+        _maybe_set(args, ctx),
+        _maybe_set(async_args, ctx),
+    )
+
+  def _websocket_base_url(self) -> str:
     url_parts = urlparse(self._http_options.base_url)
-    return url_parts._replace(scheme='wss').geturl()
+    return url_parts._replace(scheme='wss').geturl()  # type: ignore[arg-type, return-value]
 
   def _access_token(self) -> str:
     """Retrieves the access token for the credentials."""
@@ -453,11 +522,11 @@ class BaseApiClient:
         _refresh_auth(self._credentials)
       if not self._credentials.token:
         raise RuntimeError('Could not resolve API token from the environment')
-      return self._credentials.token
+      return self._credentials.token  # type: ignore[no-any-return]
     else:
       raise RuntimeError('Could not resolve API token from the environment')
 
-  async def _async_access_token(self) -> str:
+  async def _async_access_token(self) -> Union[str, Any]:
     """Retrieves the access token for the credentials asynchronously."""
     if not self._credentials:
       async with self._auth_lock:
@@ -607,7 +676,7 @@ class BaseApiClient:
 
   async def _async_request(
       self, http_request: HttpRequest, stream: bool = False
-  ):
+  ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
     if self.vertexai and not self.api_key:
       http_request.headers['Authorization'] = (
@@ -667,7 +736,7 @@ class BaseApiClient:
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Union[BaseResponse, Any]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
@@ -685,7 +754,7 @@ class BaseApiClient:
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Generator[Any, None, None]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
@@ -700,7 +769,7 @@ class BaseApiClient:
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ) -> dict[str, object]:
+  ) -> Union[BaseResponse, Any]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
@@ -717,18 +786,18 @@ class BaseApiClient:
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Any:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
 
     response = await self._async_request(http_request=http_request, stream=True)
 
-    async def async_generator():
+    async def async_generator():  # type: ignore[no-untyped-def]
       async for chunk in response:
         yield chunk
 
-    return async_generator()
+    return async_generator()  # type: ignore[no-untyped-call]
 
   def upload_file(
       self,
@@ -840,7 +909,7 @@ class BaseApiClient:
       path: str,
       *,
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Union[Any,bytes]:
     """Downloads the file data.
 
     Args:
@@ -909,7 +978,7 @@ class BaseApiClient:
 
   async def _async_upload_fd(
       self,
-      file: Union[io.IOBase, anyio.AsyncFile],
+      file: Union[io.IOBase, anyio.AsyncFile[Any]],
       upload_url: str,
       upload_size: int,
       *,
@@ -988,7 +1057,7 @@ class BaseApiClient:
       path: str,
       *,
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Union[Any, bytes]:
     """Downloads the file data.
 
     Args:
@@ -1025,5 +1094,5 @@ class BaseApiClient:
   # This method does nothing in the real api client. It is used in the
   # replay_api_client to verify the response from the SDK method matches the
   # recorded response.
-  def _verify_response(self, response_model: _common.BaseModel):
+  def _verify_response(self, response_model: _common.BaseModel) -> None:
     pass

google/genai/_automatic_function_calling_util.py

@@ -46,19 +46,6 @@ def _is_builtin_primitive_or_compound(
   return annotation in _py_builtin_type_to_schema_type.keys()
 
 
-def _raise_for_default_if_mldev(schema: types.Schema):
-  if schema.default is not None:
-    raise ValueError(
-        'Default value is not supported in function declaration schema for'
-        ' the Gemini API.'
-    )
-
-
-def _raise_if_schema_unsupported(api_option: Literal['VERTEX_AI', 'GEMINI_API'], schema: types.Schema):
-  if api_option == 'GEMINI_API':
-    _raise_for_default_if_mldev(schema)
-
-
 def _is_default_value_compatible(
     default_value: Any, annotation: inspect.Parameter.annotation  # type: ignore[valid-type]
 ) -> bool:
@@ -72,16 +59,16 @@ def _is_default_value_compatible(
       or isinstance(annotation, VersionedUnionType)
   ):
     origin = get_origin(annotation)
-    if origin in (Union, VersionedUnionType):
+    if origin in (Union, VersionedUnionType):  # type: ignore[comparison-overlap]
       return any(
           _is_default_value_compatible(default_value, arg)
           for arg in get_args(annotation)
       )
 
-    if origin is dict:
+    if origin is dict:  # type: ignore[comparison-overlap]
       return isinstance(default_value, dict)
 
-    if origin is list:
+    if origin is list:  # type: ignore[comparison-overlap]
       if not isinstance(default_value, list):
         return False
       # most tricky case, element in list is union type
@@ -97,7 +84,7 @@ def _is_default_value_compatible(
           for item in default_value
       )
 
-    if origin is Literal:
+    if origin is Literal:  # type: ignore[comparison-overlap]
       return default_value in get_args(annotation)
 
   # return False for any other unrecognized annotation
@@ -125,7 +112,6 @@ def _parse_schema_from_parameter(
         raise ValueError(default_value_error_msg)
       schema.default = param.default
     schema.type = _py_builtin_type_to_schema_type[param.annotation]
-    _raise_if_schema_unsupported(api_option, schema)
     return schema
   if (
       isinstance(param.annotation, VersionedUnionType)
@@ -166,7 +152,6 @@ def _parse_schema_from_parameter(
       if not _is_default_value_compatible(param.default, param.annotation):
         raise ValueError(default_value_error_msg)
       schema.default = param.default
-    _raise_if_schema_unsupported(api_option, schema)
     return schema
   if isinstance(param.annotation, _GenericAlias) or isinstance(
       param.annotation, builtin_types.GenericAlias
@@ -179,7 +164,6 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(api_option, schema)
       return schema
     if origin is Literal:
       if not all(isinstance(arg, str) for arg in args):
@@ -192,7 +176,6 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(api_option, schema)
       return schema
     if origin is list:
       schema.type = _py_builtin_type_to_schema_type[list]
@@ -209,7 +192,6 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(api_option, schema)
       return schema
     if origin is Union:
       schema.any_of = []
@@ -259,7 +241,6 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(api_option, schema)
       return schema
       # all other generic alias will be invoked in raise branch
   if (
@@ -284,7 +265,6 @@ def _parse_schema_from_parameter(
           func_name,
       )
     schema.required = _get_required_fields(schema)
-    _raise_if_schema_unsupported(api_option, schema)
     return schema
   raise ValueError(
       f'Failed to parse the parameter {param} of function {func_name} for'

google/genai/_common.py

@@ -20,7 +20,7 @@ import datetime
 import enum
 import functools
 import typing
-from typing import Any, Union
+from typing import Any, Callable, Optional, Union
 import uuid
 import warnings
 
@@ -31,7 +31,7 @@ from . import _api_client
 from . import errors
 
 
-def set_value_by_path(data, keys, value):
+def set_value_by_path(data: Optional[dict[Any, Any]], keys: list[str], value: Any) -> None:
   """Examples:
 
   set_value_by_path({}, ['a', 'b'], v)
@@ -46,54 +46,57 @@ def set_value_by_path(data, keys, value):
   for i, key in enumerate(keys[:-1]):
     if key.endswith('[]'):
       key_name = key[:-2]
-      if key_name not in data:
+      if data is not None and key_name not in data:
         if isinstance(value, list):
           data[key_name] = [{} for _ in range(len(value))]
         else:
           raise ValueError(
               f'value {value} must be a list given an array path {key}'
           )
-      if isinstance(value, list):
+      if isinstance(value, list) and data is not None:
         for j, d in enumerate(data[key_name]):
           set_value_by_path(d, keys[i + 1 :], value[j])
       else:
-        for d in data[key_name]:
-          set_value_by_path(d, keys[i + 1 :], value)
+        if data is not None:
+          for d in data[key_name]:
+            set_value_by_path(d, keys[i + 1 :], value)
       return
     elif key.endswith('[0]'):
       key_name = key[:-3]
-      if key_name not in data:
+      if data is not None and key_name not in data:
         data[key_name] = [{}]
-      set_value_by_path(data[key_name][0], keys[i + 1 :], value)
+      if data is not None:
+        set_value_by_path(data[key_name][0], keys[i + 1 :], value)
       return
-
-    data = data.setdefault(key, {})
-
-  existing_data = data.get(keys[-1])
-  # If there is an existing value, merge, not overwrite.
-  if existing_data is not None:
-    # Don't overwrite existing non-empty value with new empty value.
-    # This is triggered when handling tuning datasets.
-    if not value:
-      pass
-    # Don't fail when overwriting value with same value
-    elif value == existing_data:
-      pass
-    # Instead of overwriting dictionary with another dictionary, merge them.
-    # This is important for handling training and validation datasets in tuning.
-    elif isinstance(existing_data, dict) and isinstance(value, dict):
-      # Merging dictionaries. Consider deep merging in the future.
-      existing_data.update(value)
+    if data is not None:
+      data = data.setdefault(key, {})
+
+  if data is not None:
+    existing_data = data.get(keys[-1])
+    # If there is an existing value, merge, not overwrite.
+    if existing_data is not None:
+      # Don't overwrite existing non-empty value with new empty value.
+      # This is triggered when handling tuning datasets.
+      if not value:
+        pass
+      # Don't fail when overwriting value with same value
+      elif value == existing_data:
+        pass
+      # Instead of overwriting dictionary with another dictionary, merge them.
+      # This is important for handling training and validation datasets in tuning.
+      elif isinstance(existing_data, dict) and isinstance(value, dict):
+        # Merging dictionaries. Consider deep merging in the future.
+        existing_data.update(value)
+      else:
+        raise ValueError(
+            f'Cannot set value for an existing key. Key: {keys[-1]};'
+            f' Existing value: {existing_data}; New value: {value}.'
+        )
     else:
-      raise ValueError(
-          f'Cannot set value for an existing key. Key: {keys[-1]};'
-          f' Existing value: {existing_data}; New value: {value}.'
-      )
-  else:
-    data[keys[-1]] = value
+      data[keys[-1]] = value
 
 
-def get_value_by_path(data: Any, keys: list[str]):
+def get_value_by_path(data: Any, keys: list[str]) -> Any:
   """Examples:
 
   get_value_by_path({'a': {'b': v}}, ['a', 'b'])
@@ -227,7 +230,7 @@ class CaseInSensitiveEnum(str, enum.Enum):
   """Case insensitive enum."""
 
   @classmethod
-  def _missing_(cls, value):
+  def _missing_(cls, value: Any) -> Any:
     try:
       return cls[value.upper()]  # Try to access directly with uppercase
     except KeyError:
@@ -295,12 +298,12 @@ def encode_unserializable_types(data: dict[str, object]) -> dict[str, object]:
   return processed_data
 
 
-def experimental_warning(message: str):
+def experimental_warning(message: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
   """Experimental warning, only warns once."""
-  def decorator(func):
+  def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
     warning_done = False
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
       nonlocal warning_done
       if not warning_done:
         warning_done = True

google/genai/_extra_utils.py

@@ -78,16 +78,17 @@ def format_destination(
 
 def get_function_map(
     config: Optional[types.GenerateContentConfigOrDict] = None,
-) -> dict[str, Callable]:
+    is_caller_method_async: bool = False,
+) -> dict[str, Callable[..., Any]]:
   """Returns a function map from the config."""
-  function_map: dict[str, Callable] = {}
+  function_map: dict[str, Callable[..., Any]] = {}
   if not config:
     return function_map
   config_model = _create_generate_content_config_model(config)
   if config_model.tools:
     for tool in config_model.tools:
       if callable(tool):
-        if inspect.iscoroutinefunction(tool):
+        if inspect.iscoroutinefunction(tool) and not is_caller_method_async:
           raise errors.UnsupportedFunctionError(
               f'Function {tool.__name__} is a coroutine function, which is not'
               ' supported for automatic function calling. Please manually'
@@ -199,11 +200,11 @@ def convert_if_exist_pydantic_model(
   return value
 
 
-def invoke_function_from_dict_args(
-    args: Dict[str, Any], function_to_invoke: Callable
-) -> Any:
-  signature = inspect.signature(function_to_invoke)
-  func_name = function_to_invoke.__name__
+def convert_argument_from_function(
+    args: dict[str, Any], function: Callable[..., Any]
+) -> dict[str, Any]:
+  signature = inspect.signature(function)
+  func_name = function.__name__
   converted_args = {}
   for param_name, param in signature.parameters.items():
     if param_name in args:
@@ -213,19 +214,40 @@ def invoke_function_from_dict_args(
           param_name,
           func_name,
       )
+  return converted_args
+
+
+def invoke_function_from_dict_args(
+    args: Dict[str, Any], function_to_invoke: Callable[..., Any]
+) -> Any:
+  converted_args = convert_argument_from_function(args, function_to_invoke)
   try:
     return function_to_invoke(**converted_args)
   except Exception as e:
     raise errors.FunctionInvocationError(
-        f'Failed to invoke function {func_name} with converted arguments'
-        f' {converted_args} from model returned function call argument'
-        f' {args} because of error {e}'
+        f'Failed to invoke function {function_to_invoke.__name__} with'
+        f' converted arguments {converted_args} from model returned function'
+        f' call argument {args} because of error {e}'
+    )
+
+
+async def invoke_function_from_dict_args_async(
+    args: Dict[str, Any], function_to_invoke: Callable[..., Any]
+) -> Any:
+  converted_args = convert_argument_from_function(args, function_to_invoke)
+  try:
+    return await function_to_invoke(**converted_args)
+  except Exception as e:
+    raise errors.FunctionInvocationError(
+        f'Failed to invoke function {function_to_invoke.__name__} with'
+        f' converted arguments {converted_args} from model returned function'
+        f' call argument {args} because of error {e}'
     )
 
 
 def get_function_response_parts(
     response: types.GenerateContentResponse,
-    function_map: dict[str, Callable],
+    function_map: dict[str, Callable[..., Any]],
 ) -> list[types.Part]:
   """Returns the function response parts from the response."""
   func_response_parts = []
@@ -256,6 +278,44 @@ def get_function_response_parts(
         func_response_parts.append(func_response_part)
   return func_response_parts
 
+async def get_function_response_parts_async(
+    response: types.GenerateContentResponse,
+    function_map: dict[str, Callable[..., Any]],
+) -> list[types.Part]:
+  """Returns the function response parts from the response."""
+  func_response_parts = []
+  if (
+      response.candidates is not None
+      and isinstance(response.candidates[0].content, types.Content)
+      and response.candidates[0].content.parts is not None
+  ):
+    for part in response.candidates[0].content.parts:
+      if not part.function_call:
+        continue
+      func_name = part.function_call.name
+      if func_name is not None and part.function_call.args is not None:
+        func = function_map[func_name]
+        args = convert_number_values_for_dict_function_call_args(
+            part.function_call.args
+        )
+        func_response: dict[str, Any]
+        try:
+          if inspect.iscoroutinefunction(func):
+            func_response = {
+                'result': await invoke_function_from_dict_args_async(args, func)
+            }
+          else:
+            func_response = {
+                'result': invoke_function_from_dict_args(args, func)
+            }
+        except Exception as e:  # pylint: disable=broad-except
+          func_response = {'error': str(e)}
+        func_response_part = types.Part.from_function_response(
+            name=func_name, response=func_response
+        )
+        func_response_parts.append(func_response_part)
+  return func_response_parts
+
 
 def should_disable_afc(
     config: Optional[types.GenerateContentConfigOrDict] = None,